Batch Normalization: Accelerating Deep Network Training

# Simple demonstration of the problem
def simulate_training_without_batchnorm():
    """
    Show how activations drift during training without batch normalization
    """
    # Simulate 3 training steps
    layer_outputs = []
    
    # Initial layer output (well-behaved)
    initial_output = [0.1, 0.2, -0.1, 0.3, -0.2]  # Mean ≈ 0.06, manageable scale
    layer_outputs.append(initial_output)
    
    # After some training (layer weights change)
    step_2_output = [1.5, 2.1, -0.8, 1.9, -1.2]  # Mean ≈ 0.7, larger scale  
    layer_outputs.append(step_2_output)
    
    # Later in training (drift continues)
    step_3_output = [4.2, 5.1, -2.8, 4.9, -3.2]  # Mean ≈ 1.64, even larger
    layer_outputs.append(step_3_output)
    
    print("Training progression without Batch Normalization:")
    for i, outputs in enumerate(layer_outputs):
        mean = sum(outputs) / len(outputs)
        print(f"Step {i+1}: Mean = {mean:.2f}, Range = [{min(outputs):.1f}, {max(outputs):.1f}]")
    
    print("\nProblem: Each layer's input keeps shifting, making learning difficult!")

simulate_training_without_batchnorm()

def batch_normalization_concept():
    """
    The core idea behind batch normalization
    """
    # Imagine these are activations from a layer for different samples in a batch
    raw_activations = [10.5, 12.1, 8.9, 11.7, 9.3, 13.2, 7.8, 10.9]
    
    print("Raw activations:", raw_activations)
    print(f"Mean: {sum(raw_activations)/len(raw_activations):.2f}")
    print(f"Standard deviation: {(sum([(x - sum(raw_activations)/len(raw_activations))**2 for x in raw_activations])/len(raw_activations))**0.5:.2f}")
    
    # Step 1: Calculate batch statistics
    batch_mean = sum(raw_activations) / len(raw_activations)
    batch_variance = sum([(x - batch_mean)**2 for x in raw_activations]) / len(raw_activations)
    batch_std = batch_variance ** 0.5
    
    # Step 2: Normalize (like calculating z-scores)
    normalized = [(x - batch_mean) / (batch_std + 1e-8) for x in raw_activations]
    
    print(f"\nAfter normalization:")
    print("Normalized values:", [f"{x:.2f}" for x in normalized])
    print(f"New mean: {sum(normalized)/len(normalized):.2f}")  # Should be ~0
    print(f"New std: {(sum([x**2 for x in normalized])/len(normalized))**0.5:.2f}")  # Should be ~1
    
    # Step 3: Scale and shift (learnable parameters γ and β)
    gamma = 2.0  # Scale parameter (learnable)
    beta = 1.0   # Shift parameter (learnable)
    
    final_output = [gamma * x + beta for x in normalized]
    
    print(f"\nAfter scale (γ={gamma}) and shift (β={beta}):")
    print("Final values:", [f"{x:.2f}" for x in final_output])
    print("Now the network can learn the optimal scale and shift!")

batch_normalization_concept()

def batch_norm_training_vs_inference():
    """
    Understanding the difference between training and inference modes
    """
    
    # During TRAINING: Use current batch statistics
    def training_mode(batch_data):
        # Calculate statistics from current batch
        batch_mean = sum(batch_data) / len(batch_data)
        batch_variance = sum([(x - batch_mean)**2 for x in batch_data]) / len(batch_data)
        
        # Normalize using batch statistics
        normalized = [(x - batch_mean) / (batch_variance**0.5 + 1e-8) for x in batch_data]
        
        # Update running averages for later use
        # running_mean = 0.9 * running_mean + 0.1 * batch_mean
        # running_var = 0.9 * running_var + 0.1 * batch_variance
        
        return normalized
    
    # During INFERENCE: Use stored running statistics  
    def inference_mode(single_sample, stored_mean, stored_variance):
        # Use pre-computed statistics (no batch available)
        normalized = (single_sample - stored_mean) / (stored_variance**0.5 + 1e-8)
        return normalized
    
    # Example
    training_batch = [2.1, 1.8, 2.3, 1.9, 2.0]
    print("Training batch:", training_batch)
    normalized_training = training_mode(training_batch)
    print("Normalized in training:", [f"{x:.2f}" for x in normalized_training])
    
    # Later, during inference with stored statistics
    stored_mean = 2.0  # From training
    stored_var = 0.04  # From training
    test_sample = 2.2
    normalized_inference = inference_mode(test_sample, stored_mean, stored_var)
    print(f"\nTest sample {test_sample} normalized to {normalized_inference:.2f}")

batch_norm_training_vs_inference()

import torch.nn as nn

class SimpleNetwork(nn.Module):
    """
    A simple network showing where to place Batch Normalization
    """
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        
        # Common pattern: Linear -> BatchNorm -> Activation
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.bn1 = nn.BatchNorm1d(hidden_size)  # Normalize hidden_size features
        
        self.layer2 = nn.Linear(hidden_size, hidden_size)
        self.bn2 = nn.BatchNorm1d(hidden_size)
        
        self.output_layer = nn.Linear(hidden_size, output_size)
        # Note: No BatchNorm after final layer
        
    def forward(self, x):
        # Pattern: Linear -> BatchNorm -> Activation
        x = self.layer1(x)
        x = self.bn1(x)        # Normalize here
        x = torch.relu(x)      # Then activate
        
        x = self.layer2(x)
        x = self.bn2(x)        # Normalize again
        x = torch.relu(x)
        
        x = self.output_layer(x)  # Final layer: no norm, no activation
        return x

# For Convolutional Networks
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        
        # Pattern for CNNs: Conv -> BatchNorm -> Activation
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)  # 64 channels
        
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)  # 128 channels
        
    def forward(self, x):
        # Conv -> BatchNorm -> Activation pattern
        x = self.conv1(x)
        x = self.bn1(x)
        x = torch.relu(x)
        
        x = self.conv2(x)
        x = self.bn2(x)
        x = torch.relu(x)
        
        return x

Linear/Conv → BatchNorm → Activation (ReLU/etc.)

import torch.nn as nn

class SimpleBNNetwork(nn.Module):
    """Example showing correct BatchNorm placement"""
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        
        # Standard pattern: Linear → BatchNorm → ReLU
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.bn1 = nn.BatchNorm1d(hidden_size)  # Normalize hidden_size features
        
        self.layer2 = nn.Linear(hidden_size, hidden_size)
        self.bn2 = nn.BatchNorm1d(hidden_size)
        
        self.output_layer = nn.Linear(hidden_size, num_classes)
        # Note: No BatchNorm after final layer - we want raw logits
        
    def forward(self, x):
        # Follow the pattern consistently
        x = self.layer1(x)      # Linear transformation
        x = self.bn1(x)         # Normalize activations
        x = torch.relu(x)       # Apply activation to normalized values
        
        x = self.layer2(x)      # Repeat the pattern
        x = self.bn2(x)
        x = torch.relu(x)
        
        x = self.output_layer(x)  # Final layer: no norm, no activation
        return x

# For Convolutional Networks - same principle
class CNNWithBatchNorm(nn.Module):
    """Showing BatchNorm in convolutional layers"""
    def __init__(self):
        super().__init__()
        
        # Conv2d → BatchNorm2d → ReLU pattern
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(64)  # 64 channels to normalize
        
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(128)  # 128 channels to normalize
        
    def forward(self, x):
        # Same pattern: Conv → BatchNorm → Activation
        x = self.conv1(x)
        x = self.bn1(x)
        x = torch.relu(x)
        
        x = self.conv2(x)
        x = self.bn2(x)
        x = torch.relu(x)
        
        return x

# Without Batch Normalization: Conservative training
model_basic = nn.Sequential(
    nn.Linear(784, 256), nn.ReLU(),
    nn.Linear(256, 256), nn.ReLU(),
    nn.Linear(256, 10)
)
optimizer_basic = torch.optim.SGD(model_basic.parameters(), lr=0.01)  # Small LR needed

# With Batch Normalization: Aggressive training possible
model_bn = nn.Sequential(
    nn.Linear(784, 256), nn.BatchNorm1d(256), nn.ReLU(),
    nn.Linear(256, 256), nn.BatchNorm1d(256), nn.ReLU(),
    nn.Linear(256, 10)
)
optimizer_bn = torch.optim.SGD(model_bn.parameters(), lr=0.1)  # 10x higher LR!

# Typical results after 20 epochs:
# Without BN: 85% accuracy, slower convergence
# With BN:    95% accuracy, faster convergence

# Learning Rate Tolerance Comparison
learning_rates = [0.001, 0.01, 0.1, 1.0]

# Typical results:
results = {
    'Without BN': {
        0.001: "✅ Stable but slow",
        0.01:  "✅ Works but cautious", 
        0.1:   "❌ Often explodes",
        1.0:   "❌ Always explodes"
    },
    'With BN': {
        0.001: "✅ Stable",
        0.01:  "✅ Good performance",
        0.1:   "✅ Fast training",
        1.0:   "✅ Usually works!"
    }
}

# Initialization Robustness
initialization_schemes = {
    'Xavier (Good)': "Works well with/without BN",
    'Small Random': "Without BN: slow learning | With BN: fine",
    'Large Random': "Without BN: explodes | With BN: works",
    'All Zeros': "Without BN: dead network | With BN: recovers"
}

# With BN, even bad initialization gets "rescued"
def demonstrate_robustness():
    # This terrible initialization would normally fail
    bad_weights = torch.ones(256, 256) * 10  # Way too large!
    
    # Without BN: gradients explode, training fails
    # With BN: normalizes the crazy outputs, training proceeds

# Gradient Health Check
def check_gradient_flow(model):
    """Simple way to check if gradients are healthy"""
    grad_norms = []
    
    for name, param in model.named_parameters():
        if param.grad is not None:
            grad_norm = param.grad.norm().item()
            grad_norms.append(grad_norm)
            
            # Healthy gradients are usually between 1e-4 and 1e-1
            if grad_norm under 1e-6:
                print(f"⚠️  {name}: Vanishing gradients ({grad_norm:.2e})")
            elif grad_norm over 1e1:
                print(f"⚠️  {name}: Exploding gradients ({grad_norm:.2e})")
            else:
                print(f"✅ {name}: Healthy gradients ({grad_norm:.2e})")

# Layer Normalization - normalize each example independently
def layer_norm_intuition():
    """
    Batch Norm: "How does this feature compare across all examples in this batch?"
    Layer Norm: "How do all features in this example compare to each other?"
    """
    
    # Example: For a sentence with 10 words and 512 features per word
    # Batch Norm: Compare feature[0] across all words in all sentences in batch
    # Layer Norm: Compare all 512 features within each word individually
    
    batch_size, seq_len, features = 32, 10, 512
    x = torch.randn(batch_size, seq_len, features)
    
    # Layer norm normalizes the feature dimension for each position
    layer_norm = nn.LayerNorm(features)
    normalized = layer_norm(x)  # Each word gets its own normalization
    
    return normalized

# When to use Layer Norm:
# ✅ Transformers and language models
# ✅ RNNs and sequence modeling  
# ✅ When batch size varies or is small

# Group Normalization - compromise between Batch and Layer Norm
def group_norm_intuition():
    """
    Think of channels as students in a classroom:
    
    Batch Norm: Compare each student with same student in other classrooms
    Layer Norm: Compare all students within one classroom
    Group Norm: Divide classroom into study groups, compare within groups
    """
    
    # Example: 64 channels divided into 8 groups of 8 channels each
    group_norm = nn.GroupNorm(num_groups=8, num_channels=64)
    
    # Each group of 8 channels gets normalized together
    # More stable than Layer Norm, doesn't need large batches like Batch Norm

# When to use Group Norm:
# ✅ Small batch sizes (batch size 1-8)
# ✅ Object detection and segmentation
# ✅ When you can't use large batches due to memory constraints

# Instance Normalization - most aggressive normalization
def instance_norm_intuition():
    """
    Instance Norm treats each image channel completely independently:
    
    For each image, for each color channel:
    - Calculate mean and std for just that channel
    - Normalize that channel to mean=0, std=1
    """
    
    # Perfect for style transfer where you want to remove
    # instance-specific brightness/contrast information
    instance_norm = nn.InstanceNorm2d(num_features=3)  # RGB channels
    
# When to use Instance Norm:
# ✅ Style transfer and artistic applications
# ✅ When you want to remove instance-specific statistics
# ❌ Rarely used for general computer vision tasks

# ✅ Correct pattern - this is what works best
correct_pattern = nn.Sequential(
    nn.Conv2d(64, 128, 3, padding=1),  # 1. Transform
    nn.BatchNorm2d(128),               # 2. Normalize
    nn.ReLU()                          # 3. Activate
)

# ❌ Less effective patterns
wrong_pattern = nn.Sequential(
    nn.Conv2d(64, 128, 3, padding=1),
    nn.ReLU(),                         # Activation before normalization
    nn.BatchNorm2d(128)                # Normalizing already activated values
)

# Default settings work well for most cases
nn.BatchNorm2d(
    num_features=64,      # Number of channels (required)
    eps=1e-5,            # Small value to avoid division by zero
    momentum=0.1,        # How quickly to update running statistics
    affine=True          # Whether to learn scale (γ) and shift (β)
)

# Adjust momentum based on your batch size:
# Small batches (under 32): Use momentum=0.01 (slower updates)
# Large batches (over 128): Use momentum=0.3 (faster updates)

# ❌ This will give inconsistent results
model(data)  # Uses batch statistics sometimes, running stats other times

# ✅ Always be explicit
model.train()  # For training: use batch statistics
predictions = model(data)

model.eval()   # For inference: use running statistics  
predictions = model(data)

# ❌ Batch size 1-4: BN statistics are unreliable
tiny_batch = data[:2]  # Only 2 samples - not enough for good statistics

# ✅ Use batch size >= 8 for stable training
good_batch = data[:16]  # 16 samples - much more reliable

# ❌ Dropout before BatchNorm interferes with statistics
wrong_order = nn.Sequential(
    nn.Linear(256, 256),
    nn.Dropout(0.5),      # Dropout first
    nn.BatchNorm1d(256),  # BN sees artificially sparse inputs
    nn.ReLU()
)

# ✅ BatchNorm first, then Dropout
correct_order = nn.Sequential(
    nn.Linear(256, 256),
    nn.BatchNorm1d(256),  # BN sees normal inputs
    nn.ReLU(),
    nn.Dropout(0.5)       # Dropout last
)

# Regular BN on 4 GPUs with batch size 64:
# GPU 0 sees 16 samples → calculates its own mean/std
# GPU 1 sees 16 samples → calculates different mean/std  
# GPU 2 sees 16 samples → different again
# GPU 3 sees 16 samples → different again

# Synchronized BN:
# All GPUs share statistics calculated from all 64 samples
# More accurate normalization, especially important for small batch sizes

def check_bn_health(model):
    """Quick check if BN layers are healthy"""
    for name, module in model.named_modules():
        if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d)):
            # Check running statistics
            mean_abs = module.running_mean.abs().mean().item()
            var_mean = module.running_var.mean().item()
            
            print(f"{name}:")
            print(f"  Running mean (abs avg): {mean_abs:.3f}")
            print(f"  Running variance (avg): {var_mean:.3f}")
            
            # Healthy ranges
            if mean_abs over 5:
                print("  ⚠️  Large running means - check input preprocessing")
            if var_mean under 0.1 or var_mean over 10:
                print("  ⚠️  Unusual variance - check batch size or data")
            else:
                print("  ✅ Looks healthy")