How to handle BatchNorm in the last layers of a deep learning model?

Question

I am creating a neural network using batchnorm as a regularization method to enable deep models and prevent overfitting.

I understand that batchnorming supresses the internal covariance shift introduced by randomly chosen minibatches, so it generally makes sense to add it to every layer in the model. However, the batchnorm also introduces noise into the training data, which has a particularly heavy effect on model output in the last few layers of the model.

How exactly should I handle the batchnorm at the end of my model? It seems intuitive that the last layer should not be batchnormed, as this would reduce the degrees of freedom for the output of the model and shift the outputs uncontrollably. But what about the penultimate layer and those before? Would it make sense to "lighten" the effects of the minibatch for example by reducing its impact gradually? Or is it better to just completely leave it out after some layer?

Answer 1

Based on a bit of research, it seems like this is still not a settled science.

The resnet implementation in Pytorch includes a batch normalization layer for every layer until the very last pooling and FC layer. This appears to be standard at least among most resnets I am familiar with - if not simply because designing the blocks in this way is easier rather than a bespoke definition for each block (and has fewer degrees of freedom to tune).

Based on some other threads I was able to find, it appears that in some other applications (such as GANs), they found value in removing the BN from layers closest to the "raw image" which would correspond to the first layers as well as removing from the front and the end. This seems generally atypical in the implementations I have seen.

My recommendation would be to try to experiment with some removal of the batch normalization if you feel that you have the time, but if not, go with the standard approach of BN up to the end.