Exploding NLLLoss

Hari_Krishnan: <p>I am using the Transformer module provided by the PyTorch for training a model for text generation. I am using <code>NLLLoss()</code> for measuring the quality of reconstruction. After a certain number of iterations, the loss explodes and changes all weights to <code>nan</code>. This is a log generated by the training program.</p>
<pre><code class="lang-auto">root - WARNING - Loss: 203.81146240234375
root - WARNING - Loss: 124.32596588134766
root - WARNING - Loss: 62.59440612792969
root - WARNING - Loss: 59.84109115600586
root - WARNING - Loss: 59.247005462646484
root - WARNING - Loss: 48.832725524902344
root - WARNING - Loss: 57.592288970947266
root - WARNING - Loss: 50.18443298339844
root - WARNING - Loss: 46.474849700927734
root - WARNING - Loss: 52.12908172607422
root - WARNING - Loss: 50.090736389160156
root - WARNING - Loss: 66.04253387451172
root - WARNING - Loss: 49.094024658203125
root - WARNING - Loss: 36.69044494628906
root - WARNING - Loss: 48.54591369628906
root - WARNING - Loss: 60.71137237548828
root - WARNING - Loss: 40.35478591918945
root - WARNING - Loss: 49.070556640625
root - WARNING - Loss: 54.33742141723633
root - WARNING - Loss: 47.14014434814453
root - WARNING - Loss: 55.043060302734375
root - WARNING - Loss: 47.63726043701172
root - WARNING - Loss: 46.314571380615234
root - WARNING - Loss: 41.330291748046875
root - WARNING - Loss: 48.85242462158203
root - WARNING - Loss: 50.59345245361328
root - WARNING - Loss: 48.508975982666016
root - WARNING - Loss: 43.35681915283203
root - WARNING - Loss: 45.875431060791016
root - WARNING - Loss: 51.701438903808594
root - WARNING - Loss: 39.1783561706543
root - WARNING - Loss: 30.14274024963379
root - WARNING - Loss: 44.33928680419922
root - WARNING - Loss: 40.88005447387695
root - WARNING - Loss: 62.682804107666016
root - WARNING - Loss: 45.18329620361328
root - WARNING - Loss: 39.7137451171875
root - WARNING - Loss: 47.31813049316406
root - WARNING - Loss: 50.755348205566406
root - WARNING - Loss: 40.52918243408203
root - WARNING - Loss: 49.48160934448242
root - WARNING - Loss: 58.29778289794922
root - WARNING - Loss: 45.660675048828125
root - WARNING - Loss: 55.13115692138672
root - WARNING - Loss: 50.72150421142578
root - WARNING - Loss: 33.377098083496094
root - WARNING - Loss: 48.404151916503906
root - WARNING - Loss: 60.24494934082031
root - WARNING - Loss: 46.290470123291016
root - WARNING - Loss: 9.493173539216099e+24
</code></pre>
<p>As you can see, the loss goes down for some time as it should and spikes up. I have tried using gradient clipping to mitigate the issue but it did not solve the problem.</p>
<pre><code class="lang-auto">criterion_1 = nn.NLLLoss()
y_hat = model(X_train)
y_hat = y_hat.transpose(0,1)
mask = (tgt!=pad_idx).bool()
y_hat = nn.functional.log_softmax(y_hat, dim = -1)
cel = criterion_1(y_hat.reshape(-1,vocab_size), tgt.reshape(-1))
loss = cel.masked_select(mask.reshape(-1)).sum()
loss.backward()
torch.nn.utils.clip_grad_value_(model.parameters(), 100)
optimizer.step()
</code></pre>
<p>The above given is the code I am using for calculating the loss.</p>

googlebot: <p>perhaps perfect predictors exist and training reaches (1,0,0,…) state. y_hat = y_hat.clamp(-b,b) should solve that (with b like 10…20, before softmax)</p>

Hari_Krishnan: <p>For some reason, clamping the predictions is causing the loss to increase after a certain point. This continues until some of model weights becomes <code>nan</code>.</p>

googlebot: <p>Actually, I suggested early clamping, and that’s tricky with log_softmax. Post log_softmax clamping (-20.,-1e-6) or an additional loss mask may work instead. Or it is something else, I’d place a breakpoint and inspect problematic network output.</p>

Abhilash_Srivastava: <p>Few things before trying gradient clipping:</p>
<ol>
<li>What does your input data look like? Make sure it’s in the correct form, you would expect. Sometimes unnormalized input can cause huge loss values.</li>
<li>What optimizer and lr are you using?</li>
<li>Not sure if it’s a good idea to <code>sum</code> the loss values before <code>loss.backward()</code></li>
</ol>
<aside class="quote no-group" data-username="Hari_Krishnan" data-post="1" data-topic="100781">
<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/hari_krishnan/48/16377_2.png" class="avatar"> Hari_Krishnan:</div>
<blockquote>
<p><code>loss = cel.masked_select(mask.reshape(-1)).sum()</code></p>
</blockquote>
</aside>

Hari_Krishnan: <ol>
<li>Input data is in the shape <code>(batch_size, max_len)</code> and output is in the shape <code>(batch_size, max_len, vocab_size)</code>
</li>
<li>I am using Adam optimizer with lr of 0.001</li>
<li>I tried training the model taking the mean loss instead of the sum, I am still getting a spike in the loss.</li>
</ol>

Hari_Krishnan: <p>I tried clamping the output post log softmax, this is the log generated</p>
<pre><code class="lang-auto">root - WARNING - Loss: 7753.49169921875
root - WARNING - Loss: 6186.9287109375
root - WARNING - Loss: 5434.07861328125
root - WARNING - Loss: 6422.82568359375
root - WARNING - Loss: 6344.4873046875
root - WARNING - Loss: 5779.78515625
root - WARNING - Loss: 5681.9140625
root - WARNING - Loss: 5288.10498046875
root - WARNING - Loss: 5314.443359375
root - WARNING - Loss: 4506.3115234375
root - WARNING - Loss: 5896.3134765625
root - WARNING - Loss: 6842.0830078125
root - WARNING - Loss: 9111.4599609375
root - WARNING - Loss: 7685.61328125
root - WARNING - Loss: 8802.61328125
root - WARNING - Loss: 11280.5126953125
root - WARNING - Loss: 14238.529296875
root - WARNING - Loss: 13673.314453125
root - WARNING - Loss: 13150.68359375
root - WARNING - Loss: 13360.0
root - WARNING - Loss: 13180.0
root - WARNING - Loss: nan
</code></pre>
<p>The NLLLoss becomes nan after few batches</p>

Abhilash_Srivastava: <p>A few things you can check:</p>
<ol>
<li>Ensure that this output is like what you would expect (ie the scale is same as tgt).</li>
</ol>
<aside class="quote no-group" data-username="Hari_Krishnan" data-post="1" data-topic="100781">
<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/hari_krishnan/48/16377_2.png" class="avatar"> Hari_Krishnan:</div>
<blockquote>
<p><code>y_hat = nn.functional.log_softmax(y_hat, dim = -1)</code></p>
</blockquote>
</aside>
<ol start="2">
<li>I’m not sure what this part below is doing.</li>
</ol>
<aside class="quote no-group" data-username="Hari_Krishnan" data-post="1" data-topic="100781">
<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/hari_krishnan/48/16377_2.png" class="avatar"> Hari_Krishnan:</div>
<blockquote>
<p><code>loss = cel.masked_select(mask.reshape(-1)).sum()</code></p>
</blockquote>
</aside>
<p>Instead can this variable <code>loss</code> be removed and simply <code>cel.backward()</code> be used.</p>
<p>I usually go for <code>criterion = nn.CrossEntropyLoss()</code> to avoid confusion.</p>

Hari_Krishnan: <p><code>masked_select</code> is for removing the loss corresponding to the <code>&lt;pad&gt;</code> token. I’m building an architecture similar to a variational autoencoder which uses log-likelihood for the loss which is why I used <code>NLLLoss</code> over <code>CrossEntropyLoss</code></p>

googlebot: <p>Then it is something else, probably. If you use sampling from trainable distributions, the issue can be there.</p>
<p>Generally autograd.set_detect_anomaly(True) should show the problematic part (NaN inducing). As it slows down training, it is better to enable it late.</p>