Are there any advantages to using log softmax over softmax? What are the reasons to choose one over the other?
2 Answers
There are a number of advantages of using log softmax over softmax including practical reasons like improved numerical performance and gradient optimization. These advantages can be extremely important for implementation especially when training a model can be computationally challenging and expensive. At the heart of using log-softmax over softmax is the use of log probabilities over probabilities, which has nice information theoretic interpretations.
When used for classifiers the log-softmax has the effect of heavily penalizing the model when it fails to predict a correct class. Whether or not that penalization works well for solving your problem is open to your testing, so both log-softmax and softmax are worth using.
Log softmax is $$\log(\exp(x)/\sum(\exp(x))) =x - \log(\sum(\exp(x))).$$ Now $\log(\sum(\exp(x))) \approx \max(x)$, since the sum is dominated by the largest entry.
We see that log softmax is nearly just $x-\max(x)$ which is naturally much faster to compute than anything involving logarithms and exponentials. We are also guaranteed that the output won't be of a vastly different scale than the input.
Another reason is that when softmax is used with log likelihood loss, we're gonna take the logarithm of the entries anyway.