Applying log_softmax manually in a stable way

berkay_berabi: <p>Hi all,</p>
<p>I have a multiclass classification problem and my network structure is a bit complex than usual.  In a nutshell, I have 2 types of sets for labels. The ground-truth is always one label from one of the sets. (think like, labels from 0 to C are from one set and labels from C+1 to N are from another set)</p>
<p>My network calculates 2 diferent logits for each set with different architectures. What I do is:</p>
<p>I convert these logits to probability distributions via softmax and now I have 2 probability distributions one for each target set: p1 and p2.</p>
<p>I have a learnable scalar s(in range [0,1], which weights the learnt probability distributions.<br>
I calculate p1 = s * p1 and p2 = (1-s) * p2</p>
<p>Basically, depending on the case, I generate the label from one of the sets. At the end I concetanete both distributions and it is still a valid probability distribution since s is in range [0,1]. Lastly, I want to compute the loss on that.</p>
<p>I cannot use cross entropy loss because it requires raw logits. Therefore I have to use NLLLoss, However, it requires log probabilities. Since I already applied softmax, I have to apply log manually but this leads to instabilities and numerical issues as mentioned in the documentation.</p>
<p>But I need to apply the softmax in the first place because if i apply log_softmax then I do not have valid probability distributions and I can not weight them in a reasonable way with parameter s.</p>
<p>Any suggestions ?</p>

googlebot: <p>won’t something like x.clamp_min(1e-6).log() be enough?</p>

KFrank: <p>Hi Berkay!</p>
<aside class="quote no-group quote-modified" data-username="berkay_berabi" data-post="1" data-topic="98679" 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/berkay_berabi/48/16985_2.png" class="avatar"> berkay_berabi:</div>
<blockquote>
<p>I have 2 types of sets for labels. The ground-truth is always one label from one of the sets. (think like, labels from 0 to C are from one set and labels from C+1 to N are from another set)<br>
…<br>
I convert these logits to probability distributions via softmax and now I have 2 probability distributions one for each target set: p1 and p2.</p>
<p>I have a learnable scalar s(in range [0,1], which weights the learnt probability distributions.<br>
I calculate p1 = s * p1 and p2 = (1-s) * p2<br>
…<br>
Since I already applied softmax, I have to apply log manually but this leads to instabilities and numerical issues as mentioned in the documentation.</p>
<p>But I need to apply the softmax in the first place because if i apply log_softmax then I do not have valid probability distributions and I can not weight them in a reasonable way with parameter s.</p>
</blockquote>
</aside>
<p>Perform your <code>s</code>, <code>(1 - s)</code> weighting in “log space” so that you work<br>
directly with the log-probabilities and only have to call <code>log_softmax()</code>,<br>
with its better numerical stability.</p>
<p>That is, because:</p>
<pre data-code-wrap="python"><code class="lang-python">log (s * prob) = log (s) + log_prob,
</code></pre>
<p>just add <code>log (s)</code> (and <code>log (1 - s)</code>) to your results of <code>log_softmax()</code>,<br>
rather that multiplying the results of <code>softmax()</code> with <code>s</code> (and <code>(1 - s)</code>).</p>
<aside class="quote no-group">
<blockquote>
<p>I cannot use cross entropy loss because it requires raw logits. Therefore I have to use NLLLoss, However, it requires log probabilities.</p>
</blockquote>
</aside>
<p>As an aside, after computing the log-probabilities as outlined above, you<br>
could use <code>CrossEntropyLoss</code> if you wanted to.  The log-probabilities<br>
are legitimate raw-score logits – they’ve just already been “normalized”<br>
(whatever I mean by that).  In my mind it’s a little more conceptually<br>
straightforward to feed your log-probabilities into <code>NLLLoss</code>, but you’ll<br>
get the same answer if you use <code>CrossEntropyLoss</code>.</p>
<p>Here is a script that illustrates the above:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__

import math

torch.manual_seed (2020)

# 7 logits
# first group contains 4, second group 3
N = 7
C = 4
nBatch = 2   # batch size of 2, for example

s = 1/3

logits = torch.randn ((nBatch, N))
logits

lp1 = torch.nn.functional.log_softmax (logits[:, :C], dim = 1)  # log-probs for first group
lp2 = torch.nn.functional.log_softmax (logits[:, C:], dim = 1)  # log-probs for second group

log_probs = torch.cat ((math.log (s) + lp1, math.log (1.0 - s) + lp2), dim = 1)
log_probs   # all 7 log-probs

torch.exp (log_probs).sum (dim = 1)   # check that it's a valid probability distribution

# the log_probs are legitimate logits -- verify that log_softmax() doesn't change them
torch.nn.functional.log_softmax (log_probs, dim = 1)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; torch.__version__
'1.6.0'
&gt;&gt;&gt;
&gt;&gt;&gt; import math
&gt;&gt;&gt;
&gt;&gt;&gt; torch.manual_seed (2020)
&lt;torch._C.Generator object at 0x7eff574ed930&gt;
&gt;&gt;&gt;
&gt;&gt;&gt; # 7 logits
&gt;&gt;&gt; # first group contains 4, second group 3
&gt;&gt;&gt; N = 7
&gt;&gt;&gt; C = 4
&gt;&gt;&gt; nBatch = 2   # batch size of 2, for example
&gt;&gt;&gt;
&gt;&gt;&gt; s = 1/3
&gt;&gt;&gt;
&gt;&gt;&gt; logits = torch.randn ((nBatch, N))
&gt;&gt;&gt; logits
tensor([[ 1.2372, -0.9604,  1.5415, -0.4079,  0.8806,  0.0529,  0.0751],
        [ 0.4777, -0.6759, -2.1489, -1.1463, -0.2720,  1.0066, -0.0416]])
&gt;&gt;&gt;
&gt;&gt;&gt; lp1 = torch.nn.functional.log_softmax (logits[:, :C], dim = 1)  # log-probs for first group
&gt;&gt;&gt; lp2 = torch.nn.functional.log_softmax (logits[:, C:], dim = 1)  # log-probs for second group
&gt;&gt;&gt;
&gt;&gt;&gt; log_probs = torch.cat ((math.log (s) + lp1, math.log (1.0 - s) + lp2), dim = 1)
&gt;&gt;&gt; log_probs   # all 7 log-probs
tensor([[-2.0768, -4.2745, -1.7725, -3.7219, -1.0388, -1.8665, -1.8443],
        [-1.5592, -2.7127, -4.1858, -3.1831, -2.1721, -0.8934, -1.9416]])
&gt;&gt;&gt;
&gt;&gt;&gt; torch.exp (log_probs).sum (dim = 1)   # check that it's a valid probability distribution
tensor([1.0000, 1.0000])
&gt;&gt;&gt;
&gt;&gt;&gt; # the log_probs are legitimate logits -- verify that log_softmax() doesn't change them
&gt;&gt;&gt; torch.nn.functional.log_softmax (log_probs, dim = 1)
tensor([[-2.0768, -4.2745, -1.7725, -3.7219, -1.0388, -1.8665, -1.8443],
        [-1.5592, -2.7127, -4.1858, -3.1831, -2.1721, -0.8934, -1.9416]])
&gt;&gt;&gt;
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

berkay_berabi: <p>Dear KFrank,</p>
<p>Thank you very much for your reply. It makes perfect sense what you said. Now, I implemented my model as you said and it is fine. However, now the loss becomes nan after awhile. I ran torch anomaly detection and the problem is caught there. It tells me</p>
<p>RuntimeError: Function ‘LogBackward’ returned nan values in its 0th output.</p>
<p>and the last called line is the following: (Basically where I do log(1-s))</p>
<p>local_distribution = torch.log(1-st) + F.log_softmax(attention_weights.squeeze(dim=2), dim=1)</p>
<p>I thought about adding a small perturbation to both logs. Somethinglike:</p>
<p>perturbation = 1e-6<br>
w1 = log(s + perturbation)<br>
w2 = log(1 - s + perturbation)</p>
<p>Do you think that it makes sense or any do you have any other suggestions ?</p>

KFrank: <p>Hi Berkay!</p>
<aside class="quote no-group" data-username="berkay_berabi" data-post="4" data-topic="98679" 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/berkay_berabi/48/16985_2.png" class="avatar"> berkay_berabi:</div>
<blockquote>
<p>RuntimeError: Function ‘LogBackward’ returned nan values in its 0th output.</p>
<p>and the last called line is the following: (Basically where I do log(1-s))</p>
<p>local_distribution = torch.log(1-st) + F.log_softmax(attention_weights.squeeze(dim=2), dim=1)</p>
</blockquote>
</aside>
<p>Have you verified that the divergence occurs in <code>log (1 - s)</code>, and<br>
not, perhaps in a <code>log()</code> embedded in <code>log_softmax()</code>?  Can you<br>
check that <code>s</code> has become greater than (or equal to or close to) <code>1.0</code>?</p>
<aside class="quote no-group">
<blockquote>
<p>perturbation = 1e-6<br>
w1 = log(s + perturbation)<br>
w2 = log(1 - s + perturbation)</p>
</blockquote>
</aside>
<p>This would likely cure the immediate problem, but I would probably<br>
use <code>torch.clamp()</code>, rather than adding <code>perturbation</code>.</p>
<aside class="quote no-group">
<blockquote>
<p>Do you think that it makes sense or any do you have any other suggestions ?</p>
<p>I have a learnable scalar s (in range [0,1])</p>
</blockquote>
</aside>
<p>Is there anything in your training or the structure of your problem<br>
that “knows” that <code>s</code> is supposed to be a probability, and keeps it<br>
between <code>0</code> and <code>1</code>?  Your training might happily push <code>s</code> outside its<br>
valid range, leading to <code>nan</code>s in <code>log()</code>.</p>
<p>I would would suggest training the <em>logit</em> of <code>s</code>.  (It will run from <code>-inf</code> to<br>
<code>inf</code>.)  Let’s call it <code>t</code>.  Then replace <code>log (s)</code> and <code>log (1 - s)</code> with:</p>
<pre data-code-wrap="python"><code class="lang-python">torch.nn.logsigmoid (t)
torch.nn.logsigmoid (1.0 - t)
</code></pre>
<p>(where we’re using pytorch’s <code>logsigmoid()</code> to avoid the potential<br>
numerical instability of calling <code>log (sigmoid (t))</code> in two steps).</p>
<p>Best.</p>
<p>K. Frank</p>

berkay_berabi: <p>Hi KFrank,</p>
<p>Thanks again for your reply. Yes, I actually verified that problem is coming from the log and yes s was becoming 1.</p>
<p>I was already using a sigmoid function for s. So yes, in my model there exist a layer that ensures that it is in the range [0,1] and s is a learnable parameter. However, I was applying the log separately, after the sigmoid. I did not know that there exist a LogSigmoid. Now i am using it and it works just fine.</p>
<p>Thank you very much for your help !<br>
Best,<br>
Berkay</p>