Multilabel text classification with highly imbalanced training data

Question

I'm trying to train a multilabel text classification model using BERT. Each piece of text can belong to 0 or more of a total of 485 classes. My model consists of a dropout layer and a linear layer added on top of the pooled output from the bert-base-uncased model from Hugging Face. The loss function I'm using is the BCEWithLogitsLoss in PyTorch.

I have millions of labeled observations to train on. But the training data are highly unbalanced, with some labels appearing in less than 10 observations and others appearing in more than 100K observations! I'd like to get a "good" recall.

My first attempt at training without adjusting for data imbalance produced a micro recall rate of 70% (good enough) but a macro recall rate of 45% (not good enough). These numbers indicate that the model isn't performing well on underrepresented classes.

How can I effectively adjust for the data imbalance during training to improve the macro recall rate? I see we can provide label weights to BCEWithLogitsLoss loss function. But given the very high imbalance in my data leading to weights in the range of 1 to 1M, can I actually get the model to converge? My initial experiments show that a weighted loss function is going up and down during training.

Alternatively, is there a better approach than using BERT + dropout + linear layer for this type of task?

Answer 1

The architecture selection is reasonable. BERT itself has plenty of parameters. There is no need to use anything more complex.

• If the labels are mutually exclusive, you should use softmax + categorical cross-entropy instead of binary cross-entropy. In this case, you can also upsample the less frequent classes.
• If the labels are not mutually exclusive, you can still try some upsampling. You can also assign a higher weight to some classes: set reduction='none' in the loss function, and it will get you the loss per instance and per class, and you can weight it as you want (but clever data upsampling will be better).
• BERT can quickly overfit during fine-tuning, be careful with learning rates: you can use a higher learning rate for the classifier than for the BERT layers. You can also fine-tune only the last few layers of BERT.
• There are better models than BERT, e.g., RoBERTa.