Model loss difference between BertForSequenceClassification and Bert + nn.Linear

username: qattp
cooked: <p>I was trying to fine tune BERT for a continuous outcome (ranges between 0-400).  I was doing this with transformers and PyTorch Lightning in Google Colab. I noticed a big difference in validation loss during training between loading the pre-trained BERT with BertForSequenceClassification and loading with BertModel + writing nn.Linear, dropout, loss myself.</p>
<p>Specifically, using BertForSequenceClassification seems to work fine, with validation loss decreasing in each epoch. But if l load the same pre-trained BERT using BertModel.from_pretrained and writing the linear, dropout, and loss myself, the validation loss quickly stagnates:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/7/27978bf657767e2f3d708440725abdf0741cbf22.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/27978bf657767e2f3d708440725abdf0741cbf22" title="Screen Shot 2020-12-12 at 1.17.32 AM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/27978bf657767e2f3d708440725abdf0741cbf22_2_289x250.jpeg" alt="Screen Shot 2020-12-12 at 1.17.32 AM" data-base62-sha1="5Efj8HbUCgIV5ILMAm2873NQJrk" width="289" height="250" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/27978bf657767e2f3d708440725abdf0741cbf22_2_289x250.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/27978bf657767e2f3d708440725abdf0741cbf22_2_433x375.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/7/27978bf657767e2f3d708440725abdf0741cbf22_2_578x500.jpeg 2x" data-dominant-color="C5DFE6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screen Shot 2020-12-12 at 1.17.32 AM</span><span class="informations">1114×962 170 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The data, seed, hardware and much of the code are the same, with the only difference being the snippets below:</p>
<pre><code class="lang-auto">        self.config.num_labels = self.hparams.num_labels
        self.model = transformers.BertForSequenceClassification.from_pretrained(self.hparams.bert_path, config=self.config)
        self.tokenizer = transformers.BertTokenizer.from_pretrained(self.hparams.bert_path)

    def forward(self, **inputs):
        return self.model(**inputs)

    def training_step(self, batch, batch_idx):
        outputs = self(**batch)
        loss = outputs[0]
        return loss
</code></pre>
<p>and</p>
<pre><code class="lang-auto">        self.config.num_labels = self.hparams.num_labels
        self.model = transformers.BertModel.from_pretrained(self.hparams.bert_path, config=self.config)
        self.tokenizer = transformers.BertTokenizer.from_pretrained(self.hparams.bert_path)
        self.drop = torch.nn.Dropout(p=self.hparams.dropout)
        self.out = torch.nn.Linear(self.model.config.hidden_size, self.hparams.num_labels)
        self.loss = torch.nn.MSELoss()

    def forward(self, input_ids, att_mask):
        res = self.model(input_ids = input_ids, attention_mask = att_mask)
        dropout_output = self.drop(res.pooler_output)
        out = self.out(dropout_output)
        return out

    def training_step(self, batch, batch_idx):
        outputs = self(input_ids = batch["input_ids"], att_mask = batch["attention_mask"])
        loss = self.loss(outputs.view(-1), batch['labels'].view(-1))
        return loss
</code></pre>
<p>I’m lost as to why this is the case. I especially want the second approach to work so I can build on this further. Any advice is greatly appreciated!</p>

username: Abhilash_Srivastava
cooked: <p>I am wondering if this is a classification task should you be using CrossEntropy loss instead of MSELoss?</p>