Validation showing huge fluctuations. What could be the cause?

Question

I'm training a CNN for a 3-class image classification problem. My training loss decreased smoothly, which is the expected behaviour. However, my validation loss shows a lot of fluctuation.

Is this something that I should be worried about, or should I just pick the model which scores the best on my performance measure (accuracy)?

Additional info: I'm fine-tuning the last layer of a Resnet-18 that was pre-trained on ImageNet data in PyTorch. I must note that I'm using a weighted loss function for the training phase since my data is highly unbalanced. To plot the losses, however, I use the unweighted loss as to be able to compare the validation and training loss. I would use the unweighted loss, were it not that the distributions of the training dataset and validation dataset are somewhat different (they are both highly unbalanced however).

Answer 1

From simple inspection of your plot, I could make a few conclusions and list things to try. (This is without knowing any more about your setup: training parameters and model hyperparameters).

It looks like the loss is decreasing (put a line of best fit through the validation loss). It also looks like you might be able to train for longer to improve results, as the curve is still headed downwards.

First I will try answer your title question:

what is the cause of the fluctation in the validation loss?

I can think of three possibilities:

1. Regularisation - to help smoothing the learning process and make the model weights more robust. Adding/increasing your regularisation will prevent large updates to weights being introduced.
2. Batch size - is it relatively small (e.g. < 20?). This would mean that the measured mean error at the end of the network is computed using only a few samples. With a batch size of, say 8, then getting 4/8 correct and compared to getting 6/8 correct has a large relative difference when looking at the loss. Taking the mean of the errors with such small batches will lead to a not-so-smooth loss curve. If you have enough GPU memory/RAM, try increasing batch size.
3. Learning Rate - might be too large. This is similar to the first point regarding regularisation. To make smoother improvements, you might need to slow down the pace of learning as you approach a loss-minimum. You can make this perhaps run on a schedule, whereby is is reduce by some factor (e.g. multiply it by 0.5) every time the validation loss has not improved after, say 6 epochs. This will prevent you from taking big steps and then maybe overshooting a minumum and just bouncing around it.

Specific to your task, I would also suggest trying to perhaps unfreeze another layer, to increase the scope of your fine-tuning. This will give the Resnet-18 a little more freedom to learn, based on your data.

Regarding your last question:

Is this something that I should be worried about, or should I just pick the model which scores the best on my performance measure (accuracy)?

Should you be worried? In short, no. A validation loss curve like yours can be perfectly fine and deliver reasonable results; however, I would try some of the steps I mentioned above before settling for it.

Should you just pick the best performing model? If you mean taking the model at its point with best validation loss (validation accuracy), then I would say to be more careful. On your plot above, this might equate to around epoch 30, but I would personally take a point that has trained a little more, where the curve gets a little less volatile. Again, after having tried some of the steps outlined above.

Answer 2

I would also recommend that you use methods for data augmentation and oversampling to compensate the unbalance of the classes. This standford paper explains some of the ideas that you can implement The Effectiveness of Data Augmentation in Image Classification using Deep Learning.

I hope this helps!