Batch Normalization

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

import numpy as np
 
def illustrate_sequence_batchnorm_problem():
    """
    Demonstrate why BatchNorm is problematic for variable-length sequences.
    """
    # Three sequences of different lengths (padded to max length)
    # Values after sequence end are zeros (padding)
    sequences = np.array([
        [1.0, 2.0, 3.0, 0.0, 0.0],  # Length 3
        [4.0, 5.0, 6.0, 7.0, 0.0],  # Length 4
        [8.0, 9.0, 0.0, 0.0, 0.0],  # Length 2
    ])
    lengths = [3, 4, 2]
    
    print("Input sequences (0 = padding):")
    print(sequences)
    print(f"Lengths: {lengths}")
    
    # BatchNorm would compute statistics over ALL positions
    batch_mean = np.mean(sequences, axis=0)
    batch_var = np.var(sequences, axis=0)
    
    print(f"\nBatchNorm statistics per time step:")
    print(f"Means: {batch_mean}")
    print(f"Vars:  {batch_var}")
    
    # Problems:
    print("\nProblems:")
    print("1. Position 2: Mean includes padding from seq 3, biasing statistics")
    print("2. Position 3: Only seq 2 contributes; mean/var dominated by single value")
    print("3. Position 4: All zeros (only padding) - variance is 0!")
    print("4. As training progresses, which sequences end where changes statistics")
    
    # LayerNorm approach: normalize each sequence position independently
    print("\n--- LayerNorm Alternative ---")
    for i, (seq, length) in enumerate(zip(sequences, lengths)):
        valid = seq[:length]
        ln_mean = np.mean(valid)
        ln_var = np.var(valid)
        normalized = (valid - ln_mean) / np.sqrt(ln_var + 1e-5)
        print(f"Sequence {i}: mean={ln_mean:.2f}, var={ln_var:.2f}")
        print(f"  Normalized: {normalized}")
 
illustrate_sequence_batchnorm_problem()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

import numpy as np
 
def layer_norm_forward(x, gamma, beta, eps=1e-5):
    """
    Layer Normalization forward pass.
    
    Args:
        x: Input tensor of shape (batch_size, features) or (batch, seq, features)
        gamma: Scale parameter, shape (features,) or (normalized_shape)
        beta: Shift parameter, shape (features,) or (normalized_shape)
        eps: Small constant for numerical stability
    
    Key difference from BatchNorm:
    - Statistics computed per sample, across features
    - No running statistics needed
    - Same behavior in training and inference
    """
    # For 2D input (batch_size, features)
    # Normalize over the feature dimension (axis=-1)
    
    # Compute mean and variance per sample
    mu = np.mean(x, axis=-1, keepdims=True)    # Shape: (batch_size, 1)
    var = np.var(x, axis=-1, keepdims=True)    # Shape: (batch_size, 1)
    
    # Normalize
    x_norm = (x - mu) / np.sqrt(var + eps)
    
    # Scale and shift
    y = gamma * x_norm + beta
    
    return y, (mu, var, x_norm)
 
# Example: comparing BatchNorm vs LayerNorm normalization axes
np.random.seed(42)
batch_size, features = 4, 8
x = np.random.randn(batch_size, features) * 2 + 1
 
print("Input shape:", x.shape)
print("\n--- Axis Comparison ---")
 
# BatchNorm: statistics over batch dimension (axis 0)
bn_mean = np.mean(x, axis=0)  # Shape: (features,)
print(f"BatchNorm: mean over axis 0 → shape {bn_mean.shape}")
print(f"  (one mean per feature, computed from {batch_size} samples)")
 
# LayerNorm: statistics over feature dimension (axis -1)
ln_mean = np.mean(x, axis=-1)  # Shape: (batch_size,)
print(f"LayerNorm: mean over axis -1 → shape {ln_mean.shape}")
print(f"  (one mean per sample, computed from {features} features)")
 
# Apply LayerNorm
gamma = np.ones(features)
beta = np.zeros(features)
y, _ = layer_norm_forward(x, gamma, beta)
 
print(f"\n--- LayerNorm Output ---")
for i in range(batch_size):
    print(f"Sample {i}: mean = {y[i].mean():.6f}, std = {y[i].std():.6f}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

import torch
import torch.nn as nn
 
def explore_normalized_shape():
    """
    Explore different normalized_shape configurations.
    """
    torch.manual_seed(42)
    
    # Transformer-style input: (batch, sequence, embedding)
    batch, seq, emb = 2, 4, 8
    x = torch.randn(batch, seq, emb)
    
    print(f"Input shape: {x.shape} (batch={batch}, seq={seq}, emb={emb})")
    print()
    
    # Option 1: Normalize over embedding dimension only (MOST COMMON)
    ln1 = nn.LayerNorm(normalized_shape=(emb,))
    y1 = ln1(x)
    
    print("Option 1: normalized_shape = (emb,) - STANDARD FOR TRANSFORMERS")
    print(f"  Normalizes each (batch, seq) position independently")
    print(f"  Number of independent normalizations: {batch * seq}")
    print(f"  Elements per normalization: {emb}")
    
    # Verify: each position should have mean≈0, std≈1 across embedding
    for b in range(batch):
        for s in range(seq):
            mean = y1[b, s].mean().item()
            std = y1[b, s].std(unbiased=False).item()
            if b == 0 and s < 2:  # Print just a few
                print(f"    Position [{b},{s}]: mean={mean:.4f}, std={std:.4f}")
    
    print()
    
    # Option 2: Normalize over sequence AND embedding
    ln2 = nn.LayerNorm(normalized_shape=(seq, emb))
    y2 = ln2(x)
    
    print("Option 2: normalized_shape = (seq, emb)")
    print(f"  Normalizes each batch sample independently")
    print(f"  Number of independent normalizations: {batch}")
    print(f"  Elements per normalization: {seq * emb}")
    
    for b in range(batch):
        mean = y2[b].mean().item()
        std = y2[b].std(unbiased=False).item()
        print(f"    Batch {b}: mean={mean:.4f}, std={std:.4f}")
    
    print()
    
    # Parameters comparison
    print("--- Parameter Comparison ---")
    print(f"Option 1 parameters: gamma {ln1.weight.shape}, beta {ln1.bias.shape}")
    print(f"Option 2 parameters: gamma {ln2.weight.shape}, beta {ln2.bias.shape}")
 
explore_normalized_shape()
 
# In Transformers, LayerNorm(emb_dim) is standard because:
# 1. Each token position should be normalized consistently
# 2. Sequence length can vary; normalizing over it would change statistics
# 3. Matches the dimensionality of feed-forward and attention outputs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

import torch
import torch.nn as nn
 
def compare_bn_ln_behavior():
    """
    Detailed comparison of BatchNorm and LayerNorm behavior.
    """
    torch.manual_seed(42)
    features = 8
    
    bn = nn.BatchNorm1d(features)
    ln = nn.LayerNorm(features)
    
    # Training mode
    bn.train()
    
    print("=== Batch Size Sensitivity ===")
    for batch_size in [1, 2, 4, 32]:
        x = torch.randn(batch_size, features)
        
        try:
            bn_out = bn(x)
            bn_status = f"✓ Works (output mean: {bn_out.mean():.4f})"
        except Exception as e:
            bn_status = f"✗ Error: {str(e)[:50]}"
        
        ln_out = ln(x)
        ln_status = f"✓ Works (output mean: {ln_out.mean():.4f})"
        
        print(f"Batch size {batch_size:2d}: BN: {bn_status}")
        print(f"             LN: {ln_status}")
    
    print("\n=== Output Consistency ===")
    x_single = torch.randn(1, features)
    
    # LayerNorm: always same output
    ln_outputs = [ln(x_single).detach() for _ in range(3)]
    print(f"LayerNorm same input 3 times: {all(torch.allclose(ln_outputs[0], o) for o in ln_outputs)}")
    
    # BatchNorm in eval mode: same output
    bn.eval()
    bn_eval_outputs = [bn(x_single).detach() for _ in range(3)]
    print(f"BatchNorm eval mode, same input: {all(torch.allclose(bn_eval_outputs[0], o) for o in bn_eval_outputs)}")
    
    # BatchNorm in train mode with different batch compositions
    bn.train()
    x_test = torch.randn(1, features)
    bn_train_outputs = []
    for _ in range(3):
        # Different random batch companions
        x_batch = torch.cat([x_test, torch.randn(15, features)], dim=0)
        bn_train_outputs.append(bn(x_batch)[0].detach())
    
    print(f"BatchNorm train mode, varying batch: {all(torch.allclose(bn_train_outputs[0], o, atol=1e-3) for o in bn_train_outputs)}")
    print("  (Different batch compositions → different outputs for same input)")
    
    print("\n=== Gradient Flow ===")
    bn.train()
    x1 = torch.randn(4, features, requires_grad=True)
    x2 = torch.randn(4, features, requires_grad=True)
    
    bn_y = bn(x1)
    ln_y = ln(x2)
    
    # Check if gradient of sample 0 depends on sample 1's values
    bn_y[0, 0].backward(retain_graph=True)
    ln_y[0, 0].backward(retain_graph=True)
    
    print(f"BatchNorm: sample 1's gradient non-zero? {(x1.grad[1] != 0).any()}")
    print(f"LayerNorm: sample 1's gradient non-zero? {(x2.grad[1] != 0).any()}")
    print("  (BatchNorm couples samples; LayerNorm keeps them independent)")
 
compare_bn_ln_behavior()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

import torch
import torch.nn as nn
 
class PostLNTransformerBlock(nn.Module):
    """
    Original Transformer style: LayerNorm AFTER residual addition.
    
    y = LN(x + Sublayer(x))
    """
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
    
    def forward(self, x):
        # Self-attention with residual, then normalize
        attn_out, _ = self.self_attn(x, x, x)
        x = self.ln1(x + attn_out)  # Norm AFTER residual
        
        # FFN with residual, then normalize
        ffn_out = self.ffn(x)
        x = self.ln2(x + ffn_out)  # Norm AFTER residual
        
        return x
 
 
class PreLNTransformerBlock(nn.Module):
    """
    Modern style: LayerNorm BEFORE sublayer.
    
    y = x + Sublayer(LN(x))
    
    Advantages:
    - More stable gradient flow for very deep networks
    - Residual path is "clean" (no normalization)
    - Used in GPT-2/3, modern LLMs
    """
    def __init__(self, d_model, n_heads, d_ff):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model)
        )
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
    
    def forward(self, x):
        # Normalize, then self-attention, then residual
        normed = self.ln1(x)
        attn_out, _ = self.self_attn(normed, normed, normed)
        x = x + attn_out  # Clean residual path
        
        # Normalize, then FFN, then residual
        normed = self.ln2(x)
        ffn_out = self.ffn(normed)
        x = x + ffn_out  # Clean residual path
        
        return x
 
 
# Compare gradient flow
def analyze_gradient_magnitude(block, x):
    """Analyze how gradients flow through the block."""
    x = x.clone().requires_grad_(True)
    y = block(x)
    loss = y.sum()
    loss.backward()
    return x.grad.abs().mean().item()
 
torch.manual_seed(42)
d_model, n_heads, d_ff = 256, 8, 1024
x = torch.randn(2, 10, d_model)
 
post_ln = PostLNTransformerBlock(d_model, n_heads, d_ff)
pre_ln = PreLNTransformerBlock(d_model, n_heads, d_ff)
 
print(f"Post-LN gradient magnitude: {analyze_gradient_magnitude(post_ln, x):.6f}")
print(f"Pre-LN gradient magnitude:  {analyze_gradient_magnitude(pre_ln, x):.6f}")
print("\nPre-LN typically provides more stable gradients in deep networks")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

import torch
import torch.nn as nn
 
class LayerNormLSTMCell(nn.Module):
    """
    LSTM cell with Layer Normalization.
    
    Applies LayerNorm to the pre-activation of each gate.
    This stabilizes training, especially for long sequences.
    """
    def __init__(self, input_size, hidden_size):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        
        # Weights for input-to-hidden and hidden-to-hidden
        self.W_ih = nn.Linear(input_size, 4 * hidden_size, bias=False)
        self.W_hh = nn.Linear(hidden_size, 4 * hidden_size, bias=False)
        
        # LayerNorm for each of the 4 gates (combined)
        self.ln_ih = nn.LayerNorm(4 * hidden_size)
        self.ln_hh = nn.LayerNorm(4 * hidden_size)
        self.ln_cell = nn.LayerNorm(hidden_size)  # For cell state
    
    def forward(self, x, hidden):
        """
        Args:
            x: Input at current timestep, shape (batch, input_size)
            hidden: Tuple of (h, c), each shape (batch, hidden_size)
        """
        h, c = hidden
        
        # Compute pre-activations for all gates
        ih = self.W_ih(x)
        hh = self.W_hh(h)
        
        # Apply LayerNorm separately to input and hidden contributions
        ih = self.ln_ih(ih)
        hh = self.ln_hh(hh)
        
        # Combine and split into gates
        gates = ih + hh
        i, f, g, o = gates.chunk(4, dim=1)
        
        # Apply activations
        i = torch.sigmoid(i)  # Input gate
        f = torch.sigmoid(f)  # Forget gate
        g = torch.tanh(g)     # Cell candidate
        o = torch.sigmoid(o)  # Output gate
        
        # Update cell state
        c_new = f * c + i * g
        
        # Apply LayerNorm to cell state before output
        c_normed = self.ln_cell(c_new)
        
        # Compute hidden state
        h_new = o * torch.tanh(c_normed)
        
        return h_new, (h_new, c_new)
 
 
class LayerNormLSTM(nn.Module):
    """
    Full LSTM with LayerNorm, supporting sequences.
    """
    def __init__(self, input_size, hidden_size, num_layers=1):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        self.cells = nn.ModuleList([
            LayerNormLSTMCell(
                input_size if i == 0 else hidden_size,
                hidden_size
            )
            for i in range(num_layers)
        ])
    
    def forward(self, x, hidden=None):
        """
        Args:
            x: Input sequence, shape (batch, seq_len, input_size)
            hidden: Optional initial hidden state
        """
        batch_size, seq_len, _ = x.shape
        
        if hidden is None:
            hidden = [
                (torch.zeros(batch_size, self.hidden_size, device=x.device),
                 torch.zeros(batch_size, self.hidden_size, device=x.device))
                for _ in range(self.num_layers)
            ]
        
        outputs = []
        for t in range(seq_len):
            input_t = x[:, t, :]
            
            for layer, cell in enumerate(self.cells):
                input_t, hidden[layer] = cell(input_t, hidden[layer])
            
            outputs.append(input_t)
        
        return torch.stack(outputs, dim=1), hidden
 
 
# Demonstrate the difference
torch.manual_seed(42)
batch_size, seq_len, input_size, hidden_size = 8, 50, 32, 64
 
ln_lstm = LayerNormLSTM(input_size, hidden_size)
x = torch.randn(batch_size, seq_len, input_size)
 
# Process the sequence
out, hidden = ln_lstm(x)
print(f"Output shape: {out.shape}")
print(f"Output statistics - mean: {out.mean():.4f}, std: {out.std():.4f}")
print("LayerNorm helps maintain stable activations through long sequences")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

import torch
import torch.nn as nn
import torch.nn.functional as F
 
# PyTorch LayerNorm parameters
ln = nn.LayerNorm(
    normalized_shape=512,   # Can be int or tuple
    eps=1e-5,               # Numerical stability
    elementwise_affine=True # Whether to use γ and β
)
 
print("LayerNorm Parameters:")
print(f"  normalized_shape: {ln.normalized_shape}")
print(f"  eps: {ln.eps}")
print(f"  elementwise_affine: {ln.elementwise_affine}")
print(f"  weight (γ): shape={ln.weight.shape}")
print(f"  bias (β): shape={ln.bias.shape}")
 
# Initialization
print(f"\nDefault initialization:")
print(f"  weight (γ): all ones = {torch.allclose(ln.weight, torch.ones_like(ln.weight))}")
print(f"  bias (β): all zeros = {torch.allclose(ln.bias, torch.zeros_like(ln.bias))}")
 
# Without learnable parameters
ln_no_affine = nn.LayerNorm(512, elementwise_affine=False)
print(f"\nElementwise_affine=False:")
print(f"  weight: {ln_no_affine.weight}")
print(f"  bias: {ln_no_affine.bias}")
 
# Manual implementation for understanding
def layer_norm_manual(x, normalized_shape, weight, bias, eps):
    """
    Manual LayerNorm to understand the internals.
    """
    # Determine which dims to normalize over
    if isinstance(normalized_shape, int):
        normalized_shape = (normalized_shape,)
    
    # normalized_shape specifies the last N dimensions
    # E.g., for x.shape = (B, T, D) and normalized_shape = (D,),
    # we normalize over dimension 2 (the last one)
    
    dims = tuple(range(-len(normalized_shape), 0))  # e.g., (-1,) for (D,)
    
    mean = x.mean(dim=dims, keepdim=True)
    var = x.var(dim=dims, unbiased=False, keepdim=True)
    
    x_norm = (x - mean) / torch.sqrt(var + eps)
    
    # weight and bias have shape = normalized_shape
    # They broadcast over the other dimensions
    return weight * x_norm + bias
 
# Verify manual matches PyTorch
x = torch.randn(2, 10, 512)
y_pytorch = ln(x)
y_manual = layer_norm_manual(x, 512, ln.weight, ln.bias, ln.eps)
print(f"\nManual matches PyTorch: {torch.allclose(y_pytorch, y_manual, atol=1e-6)}")
 
# Using F.layer_norm (functional interface)
y_functional = F.layer_norm(x, [512], ln.weight, ln.bias, ln.eps)
print(f"Functional matches Module: {torch.allclose(y_pytorch, y_functional, atol=1e-6)}")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

import torch
import torch.nn as nn
 
class RMSNorm(nn.Module):
    """
    Root Mean Square Layer Normalization.
    
    Advantages over standard LayerNorm:
    1. Simpler computation (no mean subtraction)
    2. Slightly faster
    3. Works well in practice
    
    Used in: LLaMA, Mamba, and other modern architectures
    """
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    
    def forward(self, x):
        # RMS = sqrt(mean(x^2))
        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        
        # Normalize by RMS and scale
        x_norm = x / rms
        return self.weight * x_norm
 
 
# Compare RMSNorm vs LayerNorm
torch.manual_seed(42)
dim = 512
x = torch.randn(2, 10, dim)
 
ln = nn.LayerNorm(dim)
rms = RMSNorm(dim)
 
import time
 
# Speed comparison (rough)
n_iterations = 1000
 
start = time.time()
for _ in range(n_iterations):
    _ = ln(x)
ln_time = time.time() - start
 
start = time.time()
for _ in range(n_iterations):
    _ = rms(x)
rms_time = time.time() - start
 
print(f"LayerNorm time: {ln_time:.4f}s")
print(f"RMSNorm time:   {rms_time:.4f}s")
print(f"RMSNorm is ~{ln_time/rms_time:.2f}x faster (on CPU)")
 
# Output properties
y_ln = ln(x)
y_rms = rms(x)
 
print(f"\nLayerNorm output - mean: {y_ln.mean():.4f}, std: {y_ln.std():.4f}")
print(f"RMSNorm output   - mean: {y_rms.mean():.4f}, std: {y_rms.std():.4f}")
print("Note: RMSNorm doesn't center the mean, but works well in practice")