Batch Normalization

Topics: Neural Networks

Motivation

In traditional neural networks, as the data propagates through the network, the distribution of each layer's inputs changes. This is known as internal covariance shift, and it can make training difficult and slow.

Batch Normalization is a technique that addresses this issue by normalizing the data across each mini-batch. It operates by calculating the mean and variance of the activations for each feature in the mini-batch and then normalizing the activations using these statistics. The normalized activations are then scaled and shifted using learnable parameters, allowing the model to adapt to the optimal activation distribution.

Batch Normalization is typically applied after the linear transformation of a layer (e.g., after the matrix multiplication in a fully connected layer or after the convolution operation in a convolutional layer) and before the non-linear activation function (e.g., ReLU).

The benefits of batch normalization:

• Solves the problem of internal covariate shift.
• Makes training faster and more stable.
• Allows use of higher learning rates.
• Helps avoid vanishing or exploding gradients.
• Can act like a regularizer sometimes reduce the need for dropout.

Details of Batch Normalization

The main idea to keep in mind is that batch normalization, as the name indicates, is applied on the mini-batch level. So, the input to the "algorithm" is the whole mini-batch.

Step 1: Compute Mean and Variance of Mini-Batch

Given a mini-batch with $m$ features, $x_1, \, x_2, \, ... \, x_m$, the mean $\mu_i$ and variance $\sigma^2_i$ for each feature across the batch is computed.
$\mu_B = \frac{1}{m} \sum_{i=1}^m x_i$
$\sigma^2_B = \frac{1}{m} \sum_{i=1}^m(x_i - \mu_i)^2$
For example, if the first layer has $256$ outputs, and a batch size of $32$ is used, then, $256$ means and variances are computed across the $32$ elements of the mini-batch.

Step 2: Normalization

Each feature $x_i$ is normalized using its computed mean and variance. The process is basically z-score normalization; so, each normalized feature has a mean of $0$ and a variance of $1$.
$\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma^2_B + \epsilon}}$
where $\epsilon$ is a small constant added for numerical stability.

Step 3: Scale and Shift the Normalized Activations

The normalized activations are then scaled using learnable parameters $\gamma$ and $\beta$:
$y_i = \gamma \hat{x}_i + \beta$
These parameters allow the model to learn the optimal scaling and shifting of the normalized activations giving the network additional flexibility.

Batch Normalization During Inference

During inference, the batch statistics $(\mu_i, \, \sigma^2_i)$ are replaced by running statistics computed during training. They are usually updated using a moving average with a momentum factor, that is, a convex combination that keeps part of the previous values and replaces the other with the current batch value.
$\mu_{running} = \alpha \, \mu_{running} + (1 - \alpha) \mu_i$
$\sigma^2_{running} = \alpha \, \sigma^2_{running} + (1 - \alpha) \sigma^2_i$
where $\alpha \in [0,1)$ is the momentum factor.

PyTorch Example

Problems

Batch normalization suffers from issues that sometimes make it difficult to use:

• Training/Testing Discrepancy: during training, BN dynamically shifts and scales the data based on batch statistics. During inference, however, it freezes and relies on estimated population statistics (running averages). If training batches do not accurately reflect the global dataset, this mismatch can severely hurt generalization.
• Batch Size Dependency: BN computes mean and variance per mini-batch. If the batch size is too small, these statistics are inaccurate, causing performance degradation. It restricts architectures from utilizing extremely small batch sizes required for large inputs or heavy memory constraints.
• Unsuitable for RNNs: in recurrent neural networks (RNNs) and sequence modeling, sequence lengths vary, often resulting in varying and smaller effective batch sizes. Applying BN across different time steps can be problematic and unstable.

References

1. https://www.geeksforgeeks.org/deep-learning/what-is-batch-normalization-in-deep-learning/
2. https://www.datacamp.com/tutorial/batch-normalization-tensorflow