Loss function(s) not reflecting optimum values and increasing with time

username: jimmy_f
cooked: <p>I’ve encountered some strange behavior in my model’s training outputs. It seems that the steps the optimizer takes do not help with convergence, and in many cases do the opposite of what the loss function would entail. For the sake of testing, I make all labels of the dataset 0.86 (arbitrary value) to see if the model can at least predict a constant value. Using the MSELoss here are the outputs at each training step :</p>
<pre><code class="lang-auto">              Loss                       |                 Model Output.               |       Label    

tensor(0.3838, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.2405]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.2788, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.3320]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.1464, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.4774]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0539, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.6278]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0121, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.7500]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0006, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.8354]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0009, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.8908]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0043, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.9257]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0077, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.9475]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0104, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.9619]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0124, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.9713]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0138, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.9776]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0149, grad_fn=&lt;MseLossBackward0&gt;) tensor([[0.9820]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
</code></pre>
<p>The same thing also happens when changing the loss to L1Loss:</p>
<pre><code class="lang-auto">              Loss                    |                 Model Output.             |       Label    
tensor(0.8063, grad_fn=&lt;SumBackward0&gt;) tensor([[0.0537]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.7222, grad_fn=&lt;SumBackward0&gt;) tensor([[0.1378]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.5698, grad_fn=&lt;SumBackward0&gt;) tensor([[0.2902]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.3777, grad_fn=&lt;SumBackward0&gt;) tensor([[0.4823]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.1891, grad_fn=&lt;SumBackward0&gt;) tensor([[0.6709]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0419, grad_fn=&lt;SumBackward0&gt;) tensor([[0.8181]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0509, grad_fn=&lt;SumBackward0&gt;) tensor([[0.9109]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.0926, grad_fn=&lt;SumBackward0&gt;) tensor([[0.9526]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.1121, grad_fn=&lt;SumBackward0&gt;) tensor([[0.9721]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
tensor(0.1224, grad_fn=&lt;SumBackward0&gt;) tensor([[0.9824]], grad_fn=&lt;MeanBackward1&gt;) tensor([[0.8600]])
</code></pre>
<p>Here the steps the optimizer is taking is not conducive to lowering the loss, but instead moves in the direction of converging to 1.</p>
<p>my model is defined as follows :</p>
<pre><code class="lang-python">
class DualBertForClassification(nn.Module):
    def __init__(self, bert_model_a, bert_model_b):
        super(DualBertForClassification, self).__init__()
        
        self.bert_model_wt = bert_model_a
        self.bert_model_mutant = bert_model_b
        self.layer_1 = nn.Linear(1024, 512)
        self.layer_2 = nn.Linear(512, 128)
        self.layer_3 = nn.Linear(128, 16)
        self.layer_4 = nn.Linear(16, 1)

    def forward(self, x):
        x_a = x[0]  
        x_b = x[1]
        x = torch.cat(
            (
                self.bert_model_wt(**x_a).last_hidden_state, # [batch_sz, sequence_sz, 1024]
                self.bert_model_mutant(**x_b).last_hidden_state # [batch_sz, sequence_sz, 1024]
            ), 
            1
        )  # [batch_sz, 2*sequence_sz, 1024, ]
        x = torch.tanh(self.layer_1(x)) # [batch_sz, 2*sequence_sz, 512, ]
        x = torch.tanh(self.layer_2(x)) # [batch_sz, 2*sequence_sz, 128, ]
        x = torch.tanh(self.layer_3(x)) # [batch_sz, 2*sequence_sz, 16, ]
        x = torch.tanh(self.layer_4(x)) # [batch_sz, 2*sequence_sz, 1, ]
        x = torch.mean(x, dim = 1)  # [batch_sz, 1]
        return x
</code></pre>
<p>bert_model_a , bert_model_b are both pretrained identical BERT models. Whose output sizes are shown in the comments. Here is also the training loop which produced the printed output shown above.</p>
<pre><code class="lang-python">    model1 = BertModel.from_pretrained(model_name)
    model2 = BertModel.from_pretrained(model_name)
    model = DualBertForClassification(model1, model2)

    
    optimizer = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9)

    
    loss_fct = nn.L1Loss(reduction="sum") # also with nn.MSELoss(reduction = "sum")
    model.train()
    for index, (input_a, input_b) in enumerate(zip(wt_inputs, alt_inputs)):
        label_val = torch.tensor([[0.86]], dtype = torch.float32)
        model_output = model((input_a, input_b))
        loss = loss_fct(model_output, label_val)
        print(loss, model_output, label_val)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
</code></pre>
<p>My thoughts are that this could be an issue with the way autograd is handling the two input BERT models, but I still cannot think of an explanation that could explain the growing loss over time. Any help is greatly appreciated,  thank you</p>

username: KFrank
cooked: <p>Hi Jimmy!</p>
<aside class="quote no-group" data-username="jimmy_f" data-post="1" data-topic="188492" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/j/8e8cbc/48.png" class="avatar"> jimmy_f:</div>
<blockquote>
<p>I still cannot think of an explanation that could explain the growing loss over time.</p>
</blockquote>
</aside>
<p>Your code looks okay (but who knows what goblins might be lurking in your<br>
<code>BertModel</code>).</p>
<p>Try turning off <code>momentum</code> in your optimizer and start by running with a very<br>
small learning rate.  Plain-vanilla <code>SGD</code> takes steps in the direction of smaller<br>
loss (That’s what gradient descent is.) and if the steps are small – small learning<br>
rate – they should indeed reduce the loss, rather than “jumping” to a point with<br>
a higher loss.</p>
<p>Note, however, that if your batches are “noisy,” then even if the optimizer step you<br>
take using the gradients for, say, batch <span class="hashtag-raw">#12</span> would reduce the loss for <em>batch <span class="hashtag-raw">#12</span>,</em><br>
it is possible that that step could increase the loss for batch <span class="hashtag-raw">#13</span>.  I don’t thing that<br>
this is your issue, but to rule it out you could try passing in the same batch over<br>
and over again.</p>
<p>To rule out something weird in <code>BertModel</code> try training (with <code>SGD</code>, no <code>momentum</code>,<br>
and a small learning rate) with the <code>BertModel</code> parameters frozen.  The “head”<br>
of your model (<code>layer_1</code>, …, <code>layer_4</code>) is very standard and should easily train<br>
to spit out your fixed target value of <code>0.86</code> and train down to a loss of zero.  (With<br>
<code>L1Loss</code>, your training could end up jumping back and forth across the minimum,<br>
necessitating an additional reduction your learning rate, but with <code>MSELoss</code> this<br>
shouldn’t happen.)</p>
<p>Best.</p>
<p>K. Frank</p>

username: jimmy_f
cooked: <p>Thanks a lot for the response and suggestions! It seems as though there isn’t anything super suspicious going on under the hood of loss computation, backprop etc. I tried all of the things you suggested</p>
<ul>
<li>lower learning rate</li>
<li>no gradient computation on the two <code>BertModel</code> instances</li>
<li>using the same batch each step as input</li>
</ul>
<p>And the model did in fact converge to 0.86 as expected. It looks like your comment about noise across batches was correct since using the same batch each time seemed to solve it. Also I ran the exact same configuration as described above but without freezing the weights of the <code>BertModel</code>, and again it seems to converge at 0.86 predictions, so it doesn’t seem like there is anything unusual about those components either.</p>
<p>Thanks again for the help!</p>