Batch Normalization

Summary

Batch Normalization (or mini-batch normalization) is one of the most influential techniques published in the field of deep learning since the advent of convolutional neural networks. Proposed by Ioffe and Szegedy in 2015 in their foundational paper “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift,” this method revolutionized the way we train modern deep networks.

Before Batch Normalization, training very deep networks was a real engineering nightmare. Practitioners had to meticulously tune learning rates, carefully initialize weights, and often deal with gradients that exploded or vanished across layers. Batch Normalization changed the game by enabling the training of networks with hundreds of layers with remarkable stability.

Today, Batch Normalization is integrated into virtually all modern architectures: ResNet, DenseNet, EfficientNet, and many others. This complete guide explores its mathematical functioning, its deep intuition, its practical implementation in Python with PyTorch, and its concrete use cases.

Mathematical Principle

The principle of Batch Normalization is elegant in its simplicity. The central idea is to normalize the activations of each layer for each mini-batch (lot) of training. Here are the fundamental equations that govern this technique.

For a mini-batch B containing m samples, we first compute the mean and variance of the activations:

• Batch mean: μ_B = (1/m) · Σ x_i — this is the mean of all activations in the current batch
• Batch variance: σ²_B = (1/m) · Σ (x_i − μ_B)² — this is the variance measuring dispersion around the mean

Then, we normalize each activation:

• Normalized activation: x̂_i = (x_i − μ_B) / √(σ²_B + ε)

The term ε (epsilon) is a small numerical constant, typically 10⁻⁵, added for numerical stability to avoid division by zero when the variance is zero or extremely small.

However, systematically normalizing to zero mean and unit variance would be too restrictive: it would prevent the network from representing distributions that naturally require other parameters. The solution is brilliant:

• Learnable affine transformation: y_i = γ · x̂_i + β

The parameters γ (gamma) and β (beta) are learnable parameters that the network adjusts during training via backpropagation. They allow the model to restore any distribution it deems useful — including the identity (γ = √σ²_B, β = μ_B), which amounts to completely canceling the normalization if the network so chooses.

Inference phase

During training, Batch Normalization uses the statistics of the current batch (μ_B and σ²_B). But during the inference (prediction) phase, we often have only a single sample — making it impossible to compute a batch mean and variance. The solution is to maintain cumulative statistics during training:

• running_mean: exponential moving average of the μ_B
• running_var: exponential moving average of the σ²_B

These statistics are updated at each training step according to the formula:

• running_new = (1 − momentum) · running_old + momentum · μ_B

The momentum (default 0.1 in PyTorch) controls the adaptation speed: a low momentum gives more weight to historical statistics, while a high momentum makes the running statistics more responsive to recent changes. During inference, these running_mean and running_var are used instead of the current batch statistics.

Intuition

To truly understand why Batch Normalization works so well, you first need to understand the problem it solves: internal covariate shift.

Imagine you’re trying to learn to walk, but the ground keeps shifting under your feet. Every time you take a step, the surface changes — sometimes sloping, sometimes bumpy, sometimes slippery. This is exactly what happens in a deep network without Batch Normalization. Each layer modifies the distribution of data it passes to the next layer. The second layer receives a distribution that evolves as the first layer adjusts its weights. The third layer receives a distribution that has been transformed by two layers in constant evolution, and so on.

Batch Normalization stabilizes this shifting ground. By normalizing the activations of each layer, it ensures that each subsequent layer receives data with a relatively stable distribution — centered and scaled. It’s the equivalent of cruise control in a car: it maintains a constant speed even on hills and descents, freeing the driver (the network) to focus on what matters — steering.

The benefits of this stabilization are multiple:

1. Higher learning rates: since gradients are better conditioned, more aggressive learning rates can be used without risking divergence.
2. Faster convergence: the network reaches good performance in far fewer epochs.
3. Implicit regularization: the noise introduced by estimating statistics on a mini-batch (rather than the entire dataset) acts as a regularizer, sometimes reducing the need for Dropout.
4. Reduced sensitivity to initialization: the network becomes much more robust to the choice of initial weights.
5. Eliminates the need for Dropout in many architectures: Batch Normalization layers already provide a form of regularization.

Python Implementation with PyTorch

1. From-scratch implementation

Here is a complete from-scratch implementation of BatchNorm2d in PyTorch, including management of running statistics for the inference phase:

import torch
import torch.nn as nn
import torch.nn.functional as F


class BatchNorm2dFromScratch(nn.Module):
    """
    From-scratch implementation of BatchNorm2d.
    Reproduces the behavior of PyTorch's nn.BatchNorm2d.
    """

    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self.affine = affine

        # Learnable gamma and beta parameters
        if self.affine:
            self.weight = nn.Parameter(torch.ones(num_features))   # gamma
            self.bias = nn.Parameter(torch.zeros(num_features))    # beta
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

        # Running statistics for inference
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))
        self.register_buffer('num_batches_tracked', torch.tensor(0))

    def forward(self, x):
        # x has shape (N, C, H, W)
        if self.training:
            # Compute batch statistics
            # Average over dimensions N, H, W for each channel C
            batch_mean = x.mean(dim=(0, 2, 3))       # shape: (C,)
            batch_var = x.var(dim=(0, 2, 3), unbiased=False)  # shape: (C,)

            # Update running statistics
            self.running_mean = (
                (1 - self.momentum) * self.running_mean
                + self.momentum * batch_mean.detach()
            )
            self.running_var = (
                (1 - self.momentum) * self.running_var
                + self.momentum * batch_var.detach()
            )
            self.num_batches_tracked += 1

            # Normalization
            x_normalized = (x - batch_mean.view(1, -1, 1, 1)) / torch.sqrt(
                batch_var.view(1, -1, 1, 1) + self.eps
            )
        else:
            # Inference mode: use running statistics
            x_normalized = (x - self.running_mean.view(1, -1, 1, 1)) / torch.sqrt(
                self.running_var.view(1, -1, 1, 1) + self.eps
            )

        # Affine transformation
        if self.affine:
            x_normalized = self.weight.view(1, -1, 1, 1) * x_normalized
            x_normalized = x_normalized + self.bias.view(1, -1, 1, 1)

        return x_normalized

2. Comparison with PyTorch’s nn.BatchNorm2d

Let’s verify that our implementation produces results identical to PyTorch’s:

import torch
import torch.nn as nn

# Create test input
torch.manual_seed(42)
x = torch.randn(8, 16, 32, 32)  # (N, C, H, W)

# Our implementation
bn_custom = BatchNorm2dFromScratch(num_features=16)
bn_custom.train()
output_custom = bn_custom(x)

# PyTorch implementation
bn_pytorch = nn.BatchNorm2d(num_features=16)
bn_pytorch.train()
output_pytorch = bn_pytorch(x)

print(f"Custom output shape: {output_custom.shape}")
print(f"PyTorch output shape: {output_pytorch.shape}")
print(f"Maximum difference: {torch.max(torch.abs(output_custom - output_pytorch)):.6f}")
print(f"Custom mean: {output_custom.mean():.6f}")
print(f"PyTorch mean: {output_pytorch.mean():.6f}")

Both implementations produce numerically very close results, validating our understanding of the internal mechanism of Batch Normalization.

3. Training a ResNet with and without Batch Normalization

Now let’s compare the impact of Batch Normalization on training a simplified ResNet-18:

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt


class ResidualBlock(nn.Module):
    """Residual block optionally with Batch Normalization."""

    def __init__(self, in_channels, out_channels, use_bn=True, stride=1):
        super().__init__()
        self.use_bn = use_bn

        self.conv1 = nn.Conv2d(
            in_channels, out_channels, kernel_size=3,
            stride=stride, padding=1, bias=not use_bn
        )
        self.bn1 = nn.BatchNorm2d(out_channels) if use_bn else None
        self.conv2 = nn.Conv2d(
            out_channels, out_channels, kernel_size=3,
            stride=1, padding=1, bias=not use_bn
        )
        self.bn2 = nn.BatchNorm2d(out_channels) if use_bn else None

        # Skip connection with convolution if needed
        if in_channels != out_channels:
            self.skip = nn.Conv2d(in_channels, out_channels, kernel_size=1,
                                  stride=stride, bias=not use_bn)
            self.skip_bn = nn.BatchNorm2d(out_channels) if use_bn else None
        else:
            self.skip = None
            self.skip_bn = None

        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        identity = x

        out = self.conv1(x)
        if self.use_bn:
            out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        if self.use_bn:
            out = self.bn2(out)

        if self.skip is not None:
            identity = self.skip(x)
            if self.skip_bn is not None:
                identity = self.skip_bn(identity)

        out += identity
        out = self.relu(out)
        return out


class SimpleResNet(nn.Module):
    """Simplified ResNet for CIFAR-10."""

    def __init__(self, use_bn=True, num_classes=10):
        super().__init__()
        self.conv_in = nn.Conv2d(3, 64, kernel_size=3, padding=1, bias=not use_bn)
        self.bn_in = nn.BatchNorm2d(64) if use_bn else None
        self.relu = nn.ReLU(inplace=True)

        self.layer1 = ResidualBlock(64, 64, use_bn=use_bn)
        self.layer2 = ResidualBlock(64, 128, use_bn=use_bn, stride=2)
        self.layer3 = ResidualBlock(128, 256, use_bn=use_bn, stride=2)
        self.layer4 = ResidualBlock(256, 512, use_bn=use_bn, stride=2)

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, num_classes)

    def forward(self, x):
        x = self.conv_in(x)
        if self.bn_in:
            x = self.bn_in(x)
        x = self.relu(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x


def train_and_compare():
    """Trains two models and compares convergence curves."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # CIFAR-10 transforms
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomCrop(32, padding=4),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                           (0.2470, 0.2435, 0.2616))
    ])

    trainset = datasets.CIFAR10(root='./data', train=True,
                                download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,
                                              shuffle=True, num_workers=2)

    epochs = 30
    criterion = nn.CrossEntropyLoss()

    results = {"with_bn": [], "without_bn": []}

    for config_name, use_bn in [("with_bn", True), ("without_bn", False)]:
        print(f"\n{'='*50}")
        print(f"Training {config_name}")
        print(f"{'='*50}")

        model = SimpleResNet(use_bn=use_bn).to(device)
        optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9,
                              weight_decay=5e-4)
        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1)

        for epoch in range(epochs):
            model.train()
            running_loss = 0.0
            correct = 0
            total = 0

            for inputs, targets in trainloader:
                inputs, targets = inputs.to(device), targets.to(device)
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                loss.backward()
                optimizer.step()

                running_loss += loss.item()
                _, predicted = outputs.max(1)
                total += targets.size(0)
                correct += predicted.eq(targets).sum().item()

            scheduler.step()
            epoch_loss = running_loss / len(trainloader)
            epoch_acc = 100.0 * correct / total
            results[config_name].append(epoch_loss)

            print(f"  Epoch {epoch+1:2d}/{epochs} — "
                  f"Loss: {epoch_loss:.4f} — Accuracy: {epoch_acc:.2f}%")

    # Convergence curve visualization
    plt.figure(figsize=(12, 5))

    plt.subplot(1, 2, 1)
    plt.plot(results["with_bn"], 'b-o', label='With Batch Norm', markersize=4)
    plt.plot(results["without_bn"], 'r-s', label='Without Batch Norm', markersize=4)
    plt.xlabel('Epoch')
    plt.ylabel("Training Loss")
    plt.title('Convergence curves: with vs without Batch Normalization')
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.subplot(1, 2, 2)
    plt.plot([0.1 if e < 10 else 0.01 if e < 20 else 0.001
              for e in range(epochs)], 'g--', label='Learning Rate')
    plt.xlabel('Epoch')
    plt.ylabel("Learning Rate")
    plt.title("Learning Rate Schedule")
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig('batch_norm_convergence.png', dpi=150)
    print("\nChart saved: batch_norm_convergence.png")


if __name__ == "__main__":
    train_and_compare()

4. What the convergence curves reveal

In practice, it is consistently observed that the model with Batch Normalization:

• Converges faster: the loss decreases more rapidly from the first epochs.
• Reaches better final accuracy: implicit regularization improves generalization.
• Supports higher learning rates: without BN, a learning rate of 0.1 often causes divergence; with BN, it’s a standard and stable choice.
• Exhibits smoother gradient evolution: updates are more regular, without the violent oscillations characteristic of unnormalized deep networks.

Batch Normalization Hyperparameters

Batch Normalization introduces a few hyperparameters that are important to understand:

Epsilon (ε) — Numerical stability

• Default value: 1e-5 (PyTorch), 1e-3 (TensorFlow/Keras)
• Role: Avoids division by zero when the batch variance is zero or near-zero.
• Impact: Generally, this parameter is not touched. Values that are too large disrupt normalization; values that are too small risk numerical instability.

Momentum — Running statistics adaptation speed

• Default value: 0.1 in PyTorch (note this is the inverse of TensorFlow which uses 0.99)
• Role: Controls the relative weight of new batch statistics versus historical statistics: running_new = (1 - momentum) × running_old + momentum × batch_stat
• Impact: A low momentum (0.1) gives more weight to history — running statistics are very smooth. A high momentum makes running stats more responsive but potentially noisier.

Affine — Learnable parameters activation

• Default value: True
• Role: Determines whether the γ (weight) and β (bias) parameters are learnable.
• Impact: In the vast majority of cases, affine=True is kept. Setting affine=False forces strict zero-mean, unit-variance normalization, which can limit the model’s capacity.

Recommended batch sizes

Batch Normalization depends heavily on batch size. A batch that is too small (less than 16-32) produces noisy mean and variance estimates. For very small batches, alternatives like Group Normalization or Layer Normalization are recommended.

Advantages and Limitations

Advantages

1. Significant training speedup: convergence reached in 5 to 10 times fewer epochs.
2. Gradient stabilization: considerably reduces exploding or vanishing gradient problems.
3. Enables higher learning rates: better-conditioned gradients allow larger optimization steps.
4. Implicit regularization: mini-batch statistics noise helps generalization.
5. Reduces sensitivity to initialization: the network is more robust to initial weight choices.
6. Simplifies architecture: reduces the need for Dropout in many convolutional architectures.

Limitations

1. Batch size dependence: small batches (less than 16) give unreliable statistics — problematic for object detection or video processing where batches are naturally small.
2. Incompatibility with sequential RNN processing: BN is not trivially applicable to variable-length sequences.
3. Noise during training vs silence at inference: the difference in behavior between training mode (noisy due to batch statistics) and inference mode (deterministic) can create a slight performance gap.
4. Memory and computational cost: computing statistics and storing running parameters add overhead, especially for very large models.
5. Interaction with Dropout: BN and Dropout together can sometimes have antagonistic effects. Best practices recommend either using them separately or applying Batch Norm before Dropout.

4 Concrete Use Cases

Use case #1: Image classification with ResNet

The classic use case for Batch Normalization. ResNet architectures integrate it after each convolution, before the ReLU activation function. This Conv → BN → ReLU combination has become the industry standard for image classification, enabling the training of networks with over 1000 layers (ResNet-1000+) with remarkable stability.

Use case #2: Generative networks (GANs)

GANs (Generative Adversarial Networks) are notoriously unstable to train. Batch Normalization plays a crucial role in stabilizing the training of generators and discriminators. In DCGAN (Deep Convolutional GAN), BN is applied everywhere except on the generator’s output layer and the discriminator’s input layer, which prevents mode collapse and improves the quality of generated images.

Use case #3: Transfer Learning and Fine-Tuning

When doing transfer learning by taking a pre-trained model (such as ResNet-50 on ImageNet), Batch Normalization requires special attention. The running statistics of the pre-trained model reflect the ImageNet distribution, not that of our new dataset. Two common strategies exist: completely freeze the BN layers (use the original stats) or put them back in training mode (recalculate them on our data). The choice depends on the similarity between the distributions.

Use case #4: Networks with Skip Connections and dense architectures

Architectures like DenseNet, where each layer is connected to all subsequent layers, benefit enormously from Batch Normalization. Without BN, activations would accumulate and explode through the multiple skip connection paths. Batch Normalization keeps each activation at a controlled scale, making these deeply stacked dense architectures efficient and effective.

See also

• Maîtriser les Duodigits en Python : Guide Complet pour Optimiser vos Calculs Numériques
• Maîtriser les Permutations de Projet avec Python : Guide Complet pour Développeurs