I am building a multi-class sentiment analysis BERT model that's optimized to give the best f1 score.

More specifically, I train each epoch by optimizing binary cross entropy per class, taking the mean, and then run back propagation to optimize the parameters. At the end of each epoch, I then compute the soft-f1 score on the validation set, and then save that model only if the score has improved. The predicted class is chosen by taking an argmax, so there is no threshold.

There are a few observations here, some of which are troubling:

  1. If I simply use log-loss as the criteria for choosing the best model, then the model will typically only train for 3-4 epochs before the losses start to degrade. This is the standard approach, and gives descent f1 scores across the training/validation and test sets.

  2. If I instead use the soft-f1 criteria as specified above, the model will typically train for 15 epochs (where I use a stopping condition of 5 epochs if there is no improvement). The f1 scores are significantly better on the training set, and marginally better on the validation and test sets. I'm guessing that the improvement on the training set is simply due to overfitting, as the model has run substantially longer. But then there are also improvements in the validation and test sets too. The log-loss for this second model is (naturally) worse than model #1.

My question, which is the best model? Is it correct to use the f1 score to choose the best model while training, or am I just overfitting without realizing? Thanks!



Your Answer

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