Why have a model output raw logits instead of softmax?

username: vishak_bharadwaj
cooked: <p>A General question in mind. If I was training a classification model that I would then want to use for inference, wouldn’t it always be preferable to have a softmax layer at the end, instead of a regular linear that is input into a cross entropy loss function?</p>
<p>Wouldn’t the former make it easier to understand NN outputs, and also make it easier during inference, since all you need is a forward call and nothing else, unlike in the latter case where you need to process the NN output with some code?</p>

username: Eta_C
cooked: <p>I think it’s because that <code>logit.argmax()</code> equals to <code>softmax(logit).argmax()</code> during inference.</p>

username: vishak_bharadwaj
cooked: <p>but a logit is just a tensor. Why would that happen like that? Is there some change that happens when you mark the model to be eval() ?</p>

username: KFrank
cooked: <p>Hi Vishak!</p>
<aside class="quote no-group" data-username="vishak_bharadwaj" data-post="1" data-topic="114289" 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/vishak_bharadwaj/48/18796_2.png" class="avatar"> vishak_bharadwaj:</div>
<blockquote>
<p>If I was training a classification model that I would then want to use for inference, wouldn’t it always be preferable to have a softmax layer at the end,</p>
</blockquote>
</aside>
<p>No.  For reasons of numerical stability, it is better to have your model<br>
output the logits from its final <code>Linear</code> layer so that pytorch can use the<br>
<a href="https://en.wikipedia.org/wiki/LogSumExp" rel="noopener nofollow ugc">log-sum-exp trick</a>, either in <code>CrossEntropyLoss</code> or in <code>LogSoftmax</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>Wouldn’t the former make it easier to understand NN outputs,</p>
</blockquote>
</aside>
<p>Yes, looking at probabilities rather than logits could be easier to<br>
understand.  But it’s not worth the potential damage to the numerical<br>
stability of the computation.  It’s also easy enough enough to run the<br>
logits through <code>softmax()</code> outside of the model (and loss function) if<br>
you want to look at probabilities.</p>
<aside class="quote no-group">
<blockquote>
<p>and also make it easier during inference, since all you need is a forward call and nothing else,</p>
</blockquote>
</aside>
<p>If you were <em>only</em> going to use this version of your model for inference<br>
(that is, not train it), you could include <code>softmax()</code> in your model.  You<br>
would still have the potential for reduced numerical stability, but it’s<br>
less likely to cause trouble than when you’re training.</p>
<aside class="quote no-group">
<blockquote>
<p>unlike in the latter case where you need to process the NN output with some code?</p>
</blockquote>
</aside>
<p>But, again, if you really need probabilities for something (and usually<br>
you don’t), it’s really not that hard or confusing to add a call to<br>
<code>softmax()</code> outside of your model.</p>
<p>Best.</p>
<p>K. Frank</p>

username: vishak_bharadwaj
cooked: <p>That was great. Thanks!</p>