Justification for LogSoftmax being better than Log(Softmax)

username: Zhihan_Yang
cooked: <p>There’ve been other questions on this forum asking about LogSoftmax vs Softmax. This question is more focused on why LogSoftmax is claimed to be better (both numerically and in terms of speed) than applying Log to the output of Softmax. The claim is mentioned in this doc page:</p>
<p><a href="https://pytorch.org/docs/stable/generated/torch.nn.functional.log_softmax.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/generated/torch.nn.functional.log_softmax.html</a></p>
<p>But, softmax by itself is actually numerically stable, and also uses the max trick for numerical stability (see link below). Therefore, it’s unclear to me why adding log to it makes it unstable.</p>
<aside class="onebox stackexchange" data-onebox-src="https://stackoverflow.com/questions/42599498/numercially-stable-softmax">
  <header class="source">

      <a href="https://stackoverflow.com/questions/42599498/numercially-stable-softmax" target="_blank" rel="noopener nofollow ugc">stackoverflow.com</a>
  </header>

  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/3646408/abhishek-bhatia" target="_blank" rel="noopener nofollow ugc">
    <img alt="Abhishek Bhatia" src="https://www.gravatar.com/avatar/307b419030c0265f697f1d5d29be439a?s=256&amp;d=identicon&amp;r=PG&amp;f=1" class="thumbnail onebox-avatar" width="256" height="256">
  </a>

<h4>
  <a href="https://stackoverflow.com/questions/42599498/numercially-stable-softmax" target="_blank" rel="noopener nofollow ugc">Numercially stable softmax</a>
</h4>

<div class="tags">
  <strong>python, numpy, nan, scientific-computing</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stackoverflow.com/users/3646408/abhishek-bhatia" target="_blank" rel="noopener nofollow ugc">
    Abhishek Bhatia
  </a>
  on <a href="https://stackoverflow.com/questions/42599498/numercially-stable-softmax" target="_blank" rel="noopener nofollow ugc">06:11PM - 04 Mar 17 UTC</a>
</div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>I’d really appreciate if someone can help me understand this via some math / code.</p>

username: KFrank
cooked: <p>Hi Zhihan!</p>
<aside class="quote no-group" data-username="Zhihan_Yang" data-post="1" data-topic="140130" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zhihan_yang/48/12773_2.png" class="avatar"> Zhihan_Yang:</div>
<blockquote>
<p>Therefore, it’s unclear to me why adding log to it makes it unstable.</p>
</blockquote>
</aside>
<p>First try some numerical experiments.</p>
<p>Use the double-precision version of the naive expression as your<br>
assumed “correct” result.  (The double-precision calculation will also<br>
have the numerical overflow issue, but it won’t set in as soon.)</p>
<pre data-code-wrap="python"><code class="lang-python">true_result = torch.log (torch.softmax (alpha * torch.tensor ([-1.0, 0.0, 1.0]).double(), dim = 0))
</code></pre>
<p>Then compare the single-precision results from</p>
<pre><code class="lang-auto">torch.log (torch.softmax (alpha * torch.tensor ([-1.0, 0.0, 1.0]), dim = 0))
# and
torch.log_softmax (alpha * torch.tensor ([-1.0, 0.0, 1.0]), dim = 0)
</code></pre>
<p>with one another and with your “true” result for increasing values of <code>alpha</code>,<br>
say, <code>alpha = 2, 5, 10, 20, 50, 100, ...</code>, and see how the results behave.</p>
<p>Second, write two of your own versions of <code>log_softmax()</code>, one where you<br>
just use the naive <code>log (softmax())</code> approach, and a second where you<br>
apply the <a href="https://en.wikipedia.org/wiki/LogSumExp#log-sum-exp_trick_for_log-domain_calculations" rel="noopener nofollow ugc">“log-sum-exp trick”</a> to the sum of exponentials in the denominator<br>
of the formula for <code>softmax()</code>.</p>
<p>Does the log-sum-exp trick significantly reduce the overflow issue?  This<br>
is pretty much how <code>log_softmax()</code> is implemented in pytorch.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Zhihan_Yang
cooked: <p>Hi KFrank!</p>
<p>Thanks a lot for the code example you gave, I gained a much better understanding of this issue. I’m sharing my results and interpretations below for you and others.</p>
<p>Try <code>alpha=100</code> gives:</p>
<pre><code class="lang-python">tensor([-200., -100.,    0.], dtype=torch.float64)  # log + softmax
tensor([-200., -100.,    0.])  # logsoftmax
</code></pre>
<p>Try <code>alpha=1000</code> gives:</p>
<pre><code class="lang-python">tensor([-inf, -inf, 0.])  # log + softmax
tensor([-2000., -1000.,     0.])  # logsoftmax
</code></pre>
<p>This immediately suggests to me that, if we apply log and softmax separately, when the output of softmax becomes very close to zero, then log would yield negative infinity.</p>
<p>For an even more succinct example, where the input of log is very close to zero (exp is just one way to achieve this):</p>
<pre><code class="lang-python">torch.log(torch.exp(torch.tensor([-2000])))  # -inf
</code></pre>
<p>but if we adopt the log-softmax idea then the answer is clearly just -2000.</p>
<p>*Numerical overflow might not be relevant in this context though, since it’s ruled out by the max-trick in softmax implementation.</p>
<p>Merry Christmas,<br>
Zhihan</p>