Batch Normalization

1. Introduction

The training of deep neural networks is often hampered by the phenomenon of Internal Covariate Shift. This refers to the change in the distribution of layer inputs during the training process, as the parameters of the preceding layers are updated. The originators of Batch Normalization posited that this shift in distribution can significantly slow down the convergence of the network, necessitating lower learning rates and more careful parameter initialization.
Furthermore, as network depth increases, the model becomes more susceptible to overfitting, where it learns patterns specific to the training data that do not generalize to new data. Batch Normalization introduces a regularizing effect that can mitigate this issue.

2. Mathematical Formulation

Batch Normalization standardizes the activations of a layer for each mini-batch. Given an input mini-batch X={x1,...,xm}X = \{x_1, ..., x_m\}X={x1​,...,xm​}, the process is defined by the following steps:

1. Calculate the mean and variance of the mini-batch:
   $${\mu }_B=\frac{1}{m}\sum _{i=1}^m{x}_i$$
   $${\sigma }_B^2=\frac{1}{m}\sum _{i=1}^m(x_i-{\mu }_B{)}^2$$
2. Normalize the input:
   $${\hat{x}}_i=\frac{x_i-{\mu }_B}{\sqrt{{\sigma }_B^2+ϵ}}$$
   where $ϵ$ is a small constant added for numerical stability.
3. Scale and shift the normalized value:
   $$y_i=\gamma {\hat{x}}_i+\beta$$
   Here, $\gamma$ (scale) and $\beta$ (shift) are learnable parameters that allow the model to recover the original data distribution if it is deemed optimal for the task.

3. Application in Neural Network Architectures

Batch Normalization is strategically placed within a network’s layers. In fully connected layers, it is typically applied between the affine transformation (linear layer) and the non-linear activation function.

For convolutional layers, Batch Normalization is applied after the convolution operation and before the non-linear activation. When a convolution has multiple output channels, normalization is performed independently for each channel. Each channel maintains its own set of learnable scale ($\gamma$) and shift ($\beta$) parameters, which are scalars. Assuming a mini-batch contains $m$ samples and the convolutional output for a single channel has a height and width of $p$ and $q$ respectively, the normalization is applied across all $m\times p\times q$ elements for that channel simultaneously.

4. Implementation

4.1. Core Batch Normalization Function

The following Python function implements the Batch Normalization logic, distinguishing between training and prediction modes. During training, it computes the mean and variance from the current mini-batch. During prediction, it uses the running averages of these statistics.

import torch
from torch import nn
from d2l import torch as d2l
def batch_norm(X, gamma, beta, moving_mean, moving_var, eps, momentum):
    # 通过is_grad_enabled来判断当前模式是训练模式还是预测模式
    if not torch.is_grad_enabled():
        # 如果是在预测模式下，直接使用传入的移动平均所得的均值和方差
        X_hat = (X - moving_mean) / torch.sqrt(moving_var + eps)
    else:
        assert len(X.shape) in (2, 4)
        if len(X.shape) == 2:
            # 使用全连接层的情况，计算特征维上的均值和方差
            mean = X.mean(dim=0)
            var = ((X - mean) ** 2).mean(dim=0)
        else:
            # 使用二维卷积层的情况，计算通道维上（axis=1）的均值和方差。
            # 这里我们需要保持X的形状以便后面可以做广播运算
            mean = X.mean(dim=(0, 2, 3), keepdim=True)
            var = ((X - mean) ** 2).mean(dim=(0, 2, 3), keepdim=True)
        # 训练模式下，用当前的均值和方差做标准化
        X_hat = (X - mean) / torch.sqrt(var + eps)
        # 更新移动平均的均值和方差
        moving_mean = momentum * moving_mean + (1.0 - momentum) * mean
        moving_var = momentum * moving_var + (1.0 - momentum) * var
    Y = gamma * X_hat + beta  # 缩放和移位
    return Y, moving_mean.data, moving_var.data

The learnable parameters gamma and beta are updated during training via backpropagation. The running mean and variance are stored in moving_mean and moving_var and are utilized during the prediction mode.

4.2. Custom BatchNorm Layer

To integrate Batch Normalization into a neural network, we encapsulate the logic within a custom nn.Module. This layer manages the learnable parameters (gamma, beta) and the non-trainable running statistics (moving_mean, moving_var).

class BatchNorm(nn.Module):
    # num_features：完全连接层的输出数量或卷积层的输出通道数。
    # num_dims：2表示完全连接层，4表示卷积层
    def __init__(self, num_features, num_dims):
        super().__init__()
        if num_dims == 2:
            shape = (1, num_features)
        else:
            shape = (1, num_features, 1, 1)
        # 参与求梯度和迭代的拉伸和偏移参数，分别初始化成1和0
        self.gamma = nn.Parameter(torch.ones(shape))
        self.beta = nn.Parameter(torch.zeros(shape))
        # 非模型参数的变量初始化为0和1
        self.moving_mean = torch.zeros(shape)
        self.moving_var = torch.ones(shape)
    def forward(self, X):
        # 如果X不在内存上，将moving_mean和moving_var
        # 复制到X所在显存上
        if self.moving_mean.device != X.device:
            self.moving_mean = self.moving_mean.to(X.device)
            self.moving_var = self.moving_var.to(X.device)
        # 保存更新过的moving_mean和moving_var
        Y, self.moving_mean, self.moving_var = batch_norm(
            X, self.gamma, self.beta, self.moving_mean,
            self.moving_var, eps=1e-5, momentum=0.9)
        return Y
net = nn.Sequential(
    nn.Conv2d(1, 6, kernel_size=5), BatchNorm(6, num_dims=4), nn.Sigmoid(),
    nn.AvgPool2d(kernel_size=2, stride=2),
    nn.Conv2d(6, 16, kernel_size=5), BatchNorm(16, num_dims=4), nn.Sigmoid(),
    nn.AvgPool2d(kernel_size=2, stride=2), nn.Flatten(),
    nn.Linear(16*4*4, 120), BatchNorm(120, num_dims=2), nn.Sigmoid(),
    nn.Linear(120, 84), BatchNorm(84, num_dims=2), nn.Sigmoid(),
    nn.Linear(84, 10))

5. Controversy and Further Discussion

Despite its widespread adoption and empirical success, the exact mechanism behind Batch Normalization’s effectiveness remains a topic of debate. The original explanation of reducing Internal Covariate Shift has been challenged by subsequent research. Alternative hypotheses suggest that its benefits stem from smoothing the optimization landscape, thereby allowing for larger learning rates and faster convergence. Furthermore, its regularizing effect, which reduces the need for other techniques like Dropout, is also a significant factor contributing to its popularity.