What problem does batch normalization solve?

Batch normalization addresses the problem of internal covariate shift during neural network training. This phenomenon occurs when the distribution of inputs to each layer changes during training, which slows down and complicates learning.

Batch normalization stabilizes and normalizes the inputs of each layer, bringing them to zero mean and unit variance within a mini-batch. This allows for higher learning rates, reduces dependence on weight initialization, and helps prevent vanishing or exploding gradients.

Example usage in PyTorch:

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(100, 50)
        self.bn1 = nn.BatchNorm1d(50)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.fc1(x)
        x = self.bn1(x)  # batch normalization
        x = self.relu(x)
        return x

Thus, batch normalization speeds up and stabilizes the training of deep neural networks.