I'm training a BERT sequence classifier on a custom dataset. When the training starts, the loss is at around ~0.4 in a few steps. I print the absolute sum of gradients for each layer/item in the model and the values are high. The model converges initially but when left to be trained for a few hours and sometimes even early as well it gets stuck. I am calculating gradients with the below code. Also the logs are at - https://pastecode.io/s/v2s3mr3e (initial convergence) . I'm printing gradients, loss, metrics and logits.

for name, param in model.named_parameters():
    print(name, param.grad.abs().sum())

While the model is stuck, the loss value is around ~0.69 with worse performance metrics (precision/recall) on the training set but the gradients are very small compared to the initial training phase. Also it seems that predictions are swinging with most of the values predicted as either 0 or 1. Following are the logs for the stuck phase - https://pastecode.io/s/cjuxog44

Training code Link - Training Code

It seems that the model is stuck at a local minima where the gradient values are relatively smaller even though the loss is high. How can I mitigate it ? One option I see is using a higher learning rate or a cyclic learning rate but not sure if that's the right approach since the the learning rate is 5e-5 with LR scheduler disabled. Below is the plot for Loss, Bert pooler and classifier gradients sum over steps.

Also the data is 50-50 balanced. Batch size is 32. I'm using AdamW. I have also tried SGD but the convergence is very slow.

Or there might be some error/reason which I am not able to identify. Please help

Plot for Loss, Bert pooler and classifier gradients sum over steps



Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Browse other questions tagged or ask your own question.