Nan gradients when using F.nll_loss

jcSun: <p>Hello everyone!</p>
<p>I am trying to  train a RBM in a discriminative way. The forward of the net compute the log-conditional probabilities. The normalization I need to perform in order to get the probabilities, however, does not involve a softmax (hence, I cannot use F.log_softmax) (see  <a href="http://www.dmi.usherb.ca/~larocheh/publications/icml-2008-discriminative-rbm.pdf" rel="noopener nofollow ugc">DRBM paper</a>, p(y|x), at page 2).</p>
<p>In the problem I’m trying to solve, it is possible to have 0 probabilities. Let <code>x</code> be the binary input, <code>x[i] = 0</code> implies <code>p(y=i|x) = 0</code>.<br>
This is usually not a problem when using <code>F.nll_loss(logits, target).backward()</code>, since the 0 probability class will never be a target. Nevertheless, this seems to be a problem when using <code>torch.log(logits/logits.sum(-1))</code>.</p>
<p>I report here an example to reproduce the error:</p>
<pre><code class="lang-auto">w = torch.tensor([1.,2.,3.], requires_grad=True)
x = torch.tensor([[1,0,1]])
c = torch.tensor([0])

logits = w*x
</code></pre>
<p>If I had to normalize using softmax, Then I could use</p>
<pre><code class="lang-auto">l_mod = logits.clone()
l_mod[l_mod==0] = float("-inf")
print(F.log_softmax(l_mod, dim=-1)) 
# out: tensor([[-2.1269,    -inf, -0.1269]], grad_fn=&lt;LogSoftmaxBackward&gt;)
F.nll_loss(F.log_softmax(l_mod, dim=-1), c).backward() #or F.cross_entropy(...)
print(w.grad)
# out: tensor([-0.8808,  0.0000,  0.8808])
</code></pre>
<p>If, instead, I normalize without using the softmax, I get:</p>
<pre><code class="lang-auto">print(torch.log(logits/logits.sum(-1)))
#out: tensor([[-1.3863,    -inf, -0.2877]], grad_fn=&lt;LogBackward&gt;)
F.nll_loss(-torch.log(logits/logits.sum(-1)), c).backward()
print(w.grad)
#out: tensor([nan, nan, nan])
</code></pre>
<p>For the moment I solved the problem by either</p>
<ol>
<li>using <code>(logits\logits.sum(-1)).clamp(min=1e-16)</code> before taking the log</li>
<li>Taking the log only of positive probabilities.:<br>
<code>probs = logits\logits.sum(-1)        return torch.log(probs[probs&gt;0])</code>
</li>
</ol>
<p>However my question remains: Why do I get nan without using such tricks? Are there better ways to avoid this problem from happening?</p>
<p>Thank you in advance!</p>

tom: <p>Edit: The explanation I initially had was bogus.<br>
It’s 0<em>inf</em>0 in a way (see below) and using log_softmax has the “correct limit” to avoid (inf*0) in the backward.</p>

jcSun: <p>Hi Tom! first of all, thank you for you answer. Could you please point to me where, in particular, I end up doing <code>-inf/-inf</code>? When derivating the loss of my example (<code>-1.3863</code>) I should not encounter any -inf in the bacward steps, no? last but not least: I have -inf log-probabilities also in the log_softmax case. Why are the gradients well defined in that case?</p>

tom: <p>Sorry, my bad…<br>
So you get 0 logit and thus infty score and feed it in the nll. This gives you 0 grad in the backward for this entry, which then gets multiplied by 1/0=inf (the derivative of log at 0), and so 0*inf=NaN. This is the curse of backpropagation because next it’ll be multiplied by 0 again.<br>
With log_softmax you would not run into the problem, because it can operate entirely in log-space (i.e. you have logsoftmax(x) = x - logsumexp(x) and so you avoid the derivative of log) and so it can tell that you only have 0.<br>
But is your loss actually a vector of log probabilities? I would seem the logsumexp is inf due to the sign flip. It would need to be log probabilities for mathematical meaning.</p>
<p>Best regards</p>
<p>Thomas</p>

jcSun: <p>Thank you a lot, Thomas! This helps to understand why the code works for <code>F.log_softmax()</code> but not with <code>torch.log(...)</code>. Still, I was not understanding why the backward() pass had to go through the problematic derivative in the first place. In particular, I can avoid the error by doing the following:</p>
<pre><code class="lang-auto">w = torch.tensor([1.,2.,3.], requires_grad=True)
x = torch.tensor([[1,0,1]])
c = torch.tensor([0])

p = w*x / (w*x).sum(-1)
print(p)

nll = -torch.log(p[0][0]) #target class = c = 0
print(nll)

nll.backward()

print(w.grad)
</code></pre>
<p>which produces the following output:</p>
<pre><code class="lang-auto">tensor([[0.2500, 0.0000, 0.7500]], grad_fn=&lt;DivBackward0&gt;). #p
tensor(1.3863, grad_fn=&lt;NegBackward&gt;) #nll_loss
tensor([-0.7500,  0.0000,  0.2500]) #  gradients
</code></pre>
<p>As expected, the derivative of the nll_loss (equivalent to the example above) works fine, since no <code>-inf</code> score is involved in its computation. Instead, if I use the following to produce the log-probabilities (and the nll_loss), I get the nan gradients.</p>
<pre><code class="lang-auto">logits = torch.log(p)
nll = -logits[0][0]
print(nll)

nll.backward()

print(w.grad)
</code></pre>
<p>output:</p>
<pre><code class="lang-auto">tensor([[0.2500, 0.0000, 0.7500]], grad_fn=&lt;DivBackward0&gt;) #p (same as before!)
tensor(1.3863, grad_fn=&lt;NegBackward&gt;) #nll_loss (same as before!)
tensor([nan, nan, nan]) # gradients
</code></pre>
<p>The only difference with the previous examples is that in this case I build the computational graph for the useless (and problematic) case, which comes from <code>torch.log(logits[0][1])</code>. Even if I don’t care about this value in the nll_loss computation, it seems its backward graph is still processed when I do  <code>F.nll_loss(torch.log(p), [0]).backward()</code>.</p>
<p>Thank you again.</p>