Difference between mathematical and Tensorflow implementation of Softmax Crossentropy with logit

Question

Softmax cross entropy with logits can be defined as:

$a_i = \frac{e^{z_i}}{\sum_{\forall j} e^{z_j}}$

$l={\sum_{\forall i}}y_ilog(a_i)$

Where $l$ is the actual loss.

But when you look deep inside into C++ Tensorflow implementation of SoftmaxCrossEntropyWithLogits operation, the exact formula which they use is descibed as:

$l={\sum_{\forall j}}y_j ((z_j-max(z))-log({\sum_{\forall i}}e^{z_i-max(z)}))$

The part: $z-max(z)$ - is perfectly understood - it is just normalization which helps to avoid under/overflow.

BUT:

• Where is the actual Softmax in their implementation?

• Why from each $z_j$ they subtract $log({\sum_{\forall i}}e^{z_i-max(z)})$ before multiply it by $y_j$?

Note: One may argue that the code I provide is just Tensorflow's implementation of CrossEntropyWithLogits operation, but the actual SoftmaxCrossEntropyWithLogits operation - additionaly checks only dimentions and do not perform any more computation.

Answer 1

As I understand it, the softmax function for $z_i$ is given by $a_i$. Then just taking the loss you've defined you get back exactly the formula that is implemented. The way it is written down however is, as you mentioned, to avoid underflow/overflow.

For instance, suppose you want to compute the following:

$A=\log(\sum_{i=1}^{4}\exp(z_i))$, with $z_i=(-1000.5,-2000.5,-3000.5,-4000.5)$

Clearly, if you just type in the formula directly, you will get an underflow error. Instead if you isolate the main contribution in the exponential by taking the $\max(z_i)$, the same formula can be written as:

$A=\max_i(z_i)+\log(\sum_{i=1}^{4}\exp(z_i-\max_i(z_i)))$

The difference now is that the expression is "numerically stable" and we see that $A\approx -1000.5$.

Thus, let's make the softmax numerically stable: \begin{align} \log(a_i)&=z_i-\log(\sum_j e^{z_j})\\ &=z_i-\max_j(z_j)-\log(\sum_je^{z_j-\max_j(z_j)}) \end{align} which is the expression that is implemented for the loss (just multiply by $y_i$ and sum over $i$).