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.


  • 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.


1 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:


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$).

  • $\begingroup$ Thanks for answer! I've probably got it when I read your answer, but just for clarification: $log(\frac{e^{z_i}}{\sum_{\forall j} e^{z_j}}) = z_i - log(\sum_{\forall j} e^{z_j})$ because in Tensorflow $ln(x) = log(x)$ and of course $ln(e^{a}) = a$ ? $\endgroup$
    – Ziemo
    May 28, 2018 at 7:25
  • $\begingroup$ Indeed, in Tensorflow $\log$ is used for the natural logarithm by default. In my answer I assume that we are also dealing with base $e$. $\endgroup$ May 29, 2018 at 18:33

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.