ML Made Simple

Understanding Batch Normalization

Batch Normalization (BN) is a technique designed to address internal covariance shift in neural networks, which is the change in distribution of layer inputs during training.

Batch Normalization: Before and After

• Without BN: Training can be slow due to changing input distributions, leading to vanishing or exploding gradients.
• With BN: Normalizing inputs within each mini-batch stabilizes the distribution, reducing internal covariance shift and accelerating convergence.

Why Batch Normalization Matters

BN is essential when different channels in a network have varying ranges of values, leading to redundancy in loss convergence. Unlike dataset normalization, BN is applied on feature maps within a batch.

How Batch Normalization Works

BN counteracts internal covariance shift by scaling and shifting the input feature maps of each mini-batch. The formula BN(x) = γx^ + β demonstrates how normalized inputs are adjusted using learnable scale (γ) and shift (β) parameters, allowing for non-linear transformations and more complex relationship learning.

Ideal Scenarios for Using Batch Normalization

BN is particularly useful in datasets with high variance among classes and in deeper networks with complex architectures. It’s beneficial in many modern professional tasks.

Experiment: Demonstrating the Benefits of Batch Normalization

A Python-based experiment using torch and matplotlib compares models with and without batch normalization.

import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim

# Function to create a neural network without batch normalization
class ModelWithoutBN(nn.Module):
    def __init__(self):
        super(ModelWithoutBN, self).__init__()
        self.fc1 = nn.Linear(1, 200)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(200, 200)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(200, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        x = self.relu2(x)
        x = self.fc3(x)
        return x

# Function to create a neural network with batch normalization
class ModelWithBN(nn.Module):
    def __init__(self):
        super(ModelWithBN, self).__init__()
        self.fc1 = nn.Linear(1, 200)
        self.bn1 = nn.BatchNorm1d(200)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(200, 200)
        self.bn2 = nn.BatchNorm1d(200)
        self.relu2 = nn.ReLU()
        self.fc3 = nn.Linear(200, 1)

    def forward(self, x):
        x = self.fc1(x)
        x = self.bn1(x)
        x = self.relu1(x)
        x = self.fc2(x)
        x = self.bn2(x)
        x = self.relu2(x)
        x = self.fc3(x)
        return x
    
# Function to train a model
def train_model(model, x, y, epochs=100, batch_size=32, use_batch_norm=False):
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters())

    if use_batch_norm:
        print("Training with Batch Normalization")
    else:
        print("Training without Batch Normalization")

    losses = []
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(x)
        loss = criterion(outputs, y)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())

    return losses


# Function to generate random data with more noise
def generate_data():
    x = np.linspace(-5, 5, 100)
    y = np.sin(x) + 0.5 * np.random.normal(0, 1, size=x.shape)
    x = torch.tensor(x, dtype=torch.float32).view(-1, 1)
    y = torch.tensor(y, dtype=torch.float32).view(-1, 1)
    return x, y

# Function to plot training loss
def plot_loss(title, losses_without_bn, losses_with_bn):
    plt.plot(losses_without_bn, label='Without Batch Norm')
    plt.plot(losses_with_bn, label='With Batch Norm')
    plt.title(title)
    plt.xlabel('Epochs')
    plt.ylabel('Loss')
    plt.legend()
    plt.show()

# Generate data
x, y = generate_data()

# Train models with and without batch normalization
model_without_bn = ModelWithoutBN()
losses_without_bn = train_model(model_without_bn, x, y, use_batch_norm=False)

model_with_bn = ModelWithBN()
losses_with_bn = train_model(model_with_bn, x, y, use_batch_norm=True)

# Plot training history for comparison
plot_loss("Training Loss Convergence", losses_without_bn, losses_with_bn)

The results clearly show a faster convergence rate in the model with BN, indicating its effectiveness.