Why is the implementation of cross entropy different in Pytorch and Tensorflow?

user_1575648 0
Viewed 649

I am going through the documentation of Cross Entropy in Pytorch and Tensorflow. I understand that they are modifying the naive implementation of Cross Entropy to solve for the potential numeric over/underflows. However, I am unable to understand as to how these modifications are helping at all.

The implementation of Cross Entropy in Pytorch follows the following logic -

enter image description here

where enter image description here is the softmax score and enter image description here is the raw score.

This doesn't seem to solve the problem because enter image description here also leads to numeric overflow.

Now, we contrast it with Tensorflow's implementation (I got it from a discussion in Github. This might be completely wrong) -

Let enter image description here is the vector of all k raw logit scores.
enter image description here

While this solves the problem of overflow but it runs into problems of underflow because it is possible that enter image description here which would lead to an even smaller enter image description here

Can someone please help me in understanding what's going on here?

pytorch tensorflow neural-network precision
2 Answers

Answering here by combining answers from the comment section for the benefit of the community.

Since you have addressed the issue of numeric overflow in PyTorch, that is handled by subtracting the max value like below(from here).

scalar_t z = std::exp(input_data[d * dim_stride] - max_input);

Coming to TensorFlow's implementation of Cross entropy the issue of underflow is not that major since it is numerically ignored with the dominating large value.

user_142906970

In PyTorch

loss = F.cross_entropy(x, target)

is equivalent to:

lp = F.log_softmax(x, dim=-1)
loss = F.nll_loss(lp, target)

For the log_softmax we do not calculate log of softmax directly instead use log-sum-exp trick:

def log_softmax(x): 
    return x - x.exp().sum(-1).log().unsqueeze(-1)

Which is numerically sable. Both PyTorch and Tensorflow use this trick.

user_58849550
Related