Batch Normalization with CUDNN

Question

I want to introduce Batch Normalization in my C++/CUDNN implementation of CNN. The implementation is currently performing well (without BN) on the MNIST dataset.

I am using the CUDNN implementation of Batch Norm, but after having read the Batch Norm paper and the CUDNN documentation carefully, still there are some points that are not clear to me.

1. From what I understood, batch norm layers are placed between the output of convolutional/dense layers and the non-linearity, like:

u --->[conv/dense (weights W)]--->[BN]--->[ReLU]---> y = ReLU(BN(W*u))

In some literature I read that placing the batch norm layer after the non-linearity is a better option. Is this true?

2. The CUDNN documentation says to use the BATCHNORM_MODE_SPATIAL for convolutional layers, and BATCHNORM_MODE_PER_ACTIVATION for dense layers. However, in another implementation (YOLO / Darknet), I only see BATCHNORM_MODE_SPATIAL being used. So I ask myself whether SPATIAL is also OK for dense layers?

3. As mentioned, from what I understood, BATCHNORM_MODE_SPATIAL should be used for conv layers: but on their input (i.e. before the convolution), or on their output (i.e. after the convolution) ?

From what I understood, it is to be used on the output, like:

u ---->[ conv (weights W) ] ---->[BN with mode SPATIAL]------>[ReLU]---> y

is this correct?

4. I see that in Darknet / YOLO they do not update the parameters explicitly with the deltas obtained from the cudnnBatchNormalizationBackward(). Does this method also perform the parameters update directly? I don't find any info on this, and I am actually applying a Gradient Descent update after having obtained the Gamma and Beta gradients from the cudnnBatchNormalizationBackward method.

5. I see that in the CUDNN implementation, a Forward method for Inference is available. Is it OK to completely skip the BN units during the inference phase? In this scenario, I would train with the BN units, and remove them from the network during inference.

I really thank you in advance for your patience and your kind help. I tried to get rid of these questions alone for quite long, but I am now blocked... your help is really appreciated.

Answer 1

In general, you perform batch normalization before the activation. The entire point of the scaling/bias parameters ($\beta$ and $\gamma$) in the original paper is to scale the normalized value ($\hat{x}$ in the paper) so that it fully captures the dynamic range of the activation operator.

For example (and this is the example used in the paper), suppose the normalized output of the layer in question (convolution or dense) had unit variance (pointwise within the matrix). Then if you applied the sigmoid function pointwise as activation, you would be constrained to be within the linear region of sigmoid function, hence negating any non-linear properties that are supposedly crucial to making the network work.

However, I'm sure you can find an architecture where the order is switched. In the end, people will do what gets the best results empirically for their use-case.

As far as the updating the parameters go (your fourth question), he does update the parameters. I took a look at the source. Parameter updates occur in the "update" functions associated with the convolutional layer. I.e. he only packages batch norm with other layers, which is a reasonable optimization to make.

For questions 2&3: You are correct except for the following caveat: if you are only using dense layers in the head of a CNN, you don't need to suddenly switch strategies. You would use the MODE_PER_ACTIVATION when you expect each dimension of the feature map to have a different distribution. Think non-vision use cases where the input map might be different types of features. If you are using dense layers in a CNN, you would still still assume the normalization is the same across the "spatial" scale. After you flatten in the final layer, you essentially flatten into $C$ channels with 1x1 spatial dimension. So your "spatial" Batch Norm will have $C$ parameters anyways.

For question 5: No. Think of the BatchNorm as a learned re-scaling and a bias operation together. You wouldn't expect the output of your neural network to be consistent if you trained it with bias then removed the biases at inference time.