Why do we use the softmax activation function on the last layer?

Suppose $i$ is the index that has the highest value (in the case when we don't use softmax at all). If we use softmax and take $i$th value, it would be the highest value because $e$ is an increasing function, so that's why I am asking this question. Taking argmax(vec) and argmax(softmax(vec)) would give us the same value.


Short answer: Generally, you don't need to do softmax if you don't need probabilities. And using raw logits leads to more numerically stable code.

Long answer: First of all, the inputs of the softmax layer are called logits.

During evaluation, if you are only interested in the highest-probability class, then you can do argmax(vec) on the logits. If you want probability distribution over classes, then you'll need to exponentiate and normalize to 1 - that's what softmax does.

During training, you'd need to have a loss function to optimize. Your training data contains true classes, so you have your target probability distribution $p_i$, which is 1 at your true class and 0 at all other classes. You train the network to produce a probability distribution $q_i$ as an output. It should be as close to the target distribution $p_i$ as possible. The "distance" measure between two probability distribution is called cross-entropy:

$$ H = - \sum p_i \log q_i $$ As you can see, you only need logs of the output probabilities - so the logits will suffice to compute the loss. For example, the keras standard CategoricalCrossentropy loss can be configured to compute it from_logits and it mentions that:

Using from_logits=True is more numerically stable.

  • $\begingroup$ thank you, now i get it $\endgroup$ May 7 at 18:42
  • $\begingroup$ Note that the cross-entropy is really defined for 2 probability distributions, and I would say that's the reason why we use the softmax. Note that, if the logits are zero or negative, then the cross-entropy is not defined because of the logarithm. I think that TensorFlow accepts logits because it probably performs a (stable) conversion of the logits to a probability distribution under the hood, but I have not looked at the source code. So, it's a matter of diving into the TF implementation of the CE loss. $\endgroup$
    – nbro
    May 7 at 19:09
  • 1
    $\begingroup$ If you directly use the logits as "the logarithms of the probability", then what you say should make sense. However, why should you interpret the output of the neural network as the logarithm of the probability of belonging to a class? Still, it would be nice to show what TF exactly does when you set from_logits=True. $\endgroup$
    – nbro
    May 7 at 19:21
  • 1
    $\begingroup$ @nbro I kinda started explaining why use of logits is more numerically stable, but it turned out to be too mush of a tangent. So that deserves a separate question and answer in my opinion. $\endgroup$
    – Kostya
    May 7 at 19:28

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

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