Fine-tuning with BertForSequenceClassification on custom dataset yields a model that outputs only the label with highest support in training set
Hello!
I have a custom dataset that I wish to fine tune BERT on for classification. The examples consist of 3 sequences each and the label set is {0, 1, ..., 9}. The training data has 570 examples, and validation has 150 examples. 
The encoding of the input examples is as follows:
`tmp_enc = sequence_a + ' [SEP] ' + sequence_b + ' [SEP] ' + sequence_c`
 `enc = tokenizer.encode(tmp_enc, add_special_tokens=True)`
where
`tokenizer =  BertTokenizer.from_pretrained('bert-base-uncased')`

my training loop is:
```
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, 
num_training_steps=5*len(X_train))

train_losses = []
valid_losses = []
for epoch in range(5):
    #train
    model.train();
    epoch_train_loss = 0
    for i in range(len(X_train)):
        model.zero_grad()
        loss, logits = model(input_ids=X_train[i], labels=y_train[i].unsqueeze(0))
        loss.backward()
        
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        epoch_train_loss += loss.item()
        if i % 100 == 0:
            print(f'example {i + 1}, loss: {loss.item()}')
    
    avg_loss = epoch_train_loss / len(X_train)
    train_losses.append(avg_loss)
    print(f'epoch {epoch + 1}, average train loss: {avg_loss}')
    
    #validation
    model.eval();
    epoch_valid_loss = 0
    for i in range(len(X_test)):
        with torch.no_grad():
            loss, logits = model(input_ids=X_test[i], labels=y_test[i].unsqueeze(0))
        epoch_valid_loss += loss.item()
    avg_loss = epoch_valid_loss / len(X_test)
    valid_losses.append(avg_loss)
    print(f'epoch {epoch + 1} done. average valid loss: {avg_loss}')
```
where
`model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=10)`
and all layers but the classifier are frozen:
```
for param in model.bert.parameters():
    param.requires_grad = False
```
My issue is that loss is not decreasing and the model eventually always outputs a single label which is the one with highest support in the data. (validated by removing the examples with highest occurring label and getting similar results). Also, the logits for all examples, passed to the model after training are quite similar. e.g.:
```
logits: tensor([[ 1.5107, -0.0595,  0.3490, -0.8669, -0.8848, -0.8097,  0.2685,  0.7246,
         -0.3133,  0.4215]]), true label: 4
logits: tensor([[ 1.4187, -0.3009,  0.3776, -0.5615, -0.7881, -0.5849,  0.3391,  0.5756,
         -0.3861,  0.3639]]), true label: 6
logits: tensor([[ 1.3919, -0.4227,  0.3455, -0.4626, -0.7795, -0.5608,  0.2996,  0.5791,
         -0.4275,  0.3700]]), true label: 8
```
I have verified that a linear classifier can achieve better results (65% accuracy as opposed to 30%) on the embeddings from pytorch_transformers' BertModel.  
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
<|||||>@royeis did you find solution for this? I am facing the same issue. <|||||>@SarikGhazarian 
Ended up solving this with a workaround. Constructed a Pytorch model that had  huggingface's BertModel as a module and another linear layer that received Bert's outputs and acted as a classifier. Bert module parameters were frozen and training worked properly from there.