Layer Normalization Integration

Learning Objectives

By the end of this section, you will:

1. Distinguish batch normalization from layer normalization and when to use each
2. Master the layer normalization formula and its application
3. Understand why layer norm is preferred for RNNs
4. Integrate layer normalization with the BiLSTM output
5. Configure layer norm parameters in PyTorch

Why This Matters: Batch normalization transformed CNN training, but it has limitations for recurrent networks. Layer normalization provides similar benefits without the batch dependency, making it the preferred choice for LSTM architectures.

Batch Norm vs Layer Norm

Both normalization techniques stabilize training, but they normalize across different dimensions.

Batch Normalization (Recap)

For a batch of activations, batch norm normalizes across the batch dimension:

1Input: (N, T, H)  ← Batch, Time, Hidden
2
3For each feature h:
4  μ_h = mean across all N samples and T timesteps
5  σ_h = std across all N samples and T timesteps
6  normalize: x̂_{n,t,h} = (x_{n,t,h} - μ_h) / σ_h

Layer Normalization

Layer norm normalizes across the feature dimension within each sample:

1Input: (N, T, H)  ← Batch, Time, Hidden
2
3For each sample n and timestep t:
4  μ_{n,t} = mean across all H features
5  σ_{n,t} = std across all H features
6  normalize: x̂_{n,t,h} = (x_{n,t,h} - μ_{n,t}) / σ_{n,t}

Visual Comparison

1Tensor shape: (N, T, H) = (Batch, Time, Hidden)
2
3BatchNorm:                    LayerNorm:
4    H ──────────►                 H ──────────►
5  ┌─────────────┐               ┌─────────────┐
6T │ ■ ■ ■ ■ ■ ■ │             T │ ▒ ▒ ▒ ▒ ▒ ▒ │  ← normalize this row
7  │ ■ ■ ■ ■ ■ ■ │               │ ░ ░ ░ ░ ░ ░ │  ← normalize this row
8  │ ■ ■ ■ ■ ■ ■ │               │ ▓ ▓ ▓ ▓ ▓ ▓ │  ← normalize this row
9  └─────────────┘               └─────────────┘
10        ↑                             ↑
11    normalize                   each row normalized
12    this column                  independently
13  across samples

Key Differences

Layer Normalization Formula

Layer normalization normalizes each sample independently across features.

Mathematical Formulation

For a single sample at timestep t with hidden dimension H:

Where:

• $\mu_t$: Mean across features at timestep t
• $\sigma_t^2$: Variance across features at timestep t
• $\epsilon$: Small constant (1e-5) for numerical stability
• $\gamma_h, \beta_h$: Learnable scale and shift per feature

Why Layer Norm for RNNs?

Several properties make layer normalization superior to batch normalization for recurrent networks.

Problem 1: Variable Sequence Lengths

RNNs often process sequences of different lengths. With batch norm:

1Batch with varying lengths:
2  Sequence 1: [x₁, x₂, x₃, x₄, x₅]  (length 5)
3  Sequence 2: [x₁, x₂, x₃, PAD, PAD]  (length 3 + padding)
4
5BatchNorm at position 4:
6  Only uses Sequence 1's value (Sequence 2 is padding)
7  Statistics become unreliable!

Layer norm avoids this—each position is normalized independently.

Problem 2: Sequential Dependency

RNN hidden states at different timesteps are correlated. Batch normalizing across time can damage this structure:

• State at t depends on state at t-1
• Normalizing all t positions together ignores temporal order
• Layer norm respects the sequential structure

Problem 3: Small Batch Sizes

RNN training often uses smaller batches (memory constraints). Batch norm statistics become noisy with small N:

Integration with BiLSTM

We apply layer normalization to the BiLSTM output before passing to attention.

Placement

1BiLSTM Output: (B, 30, 256)
2       ↓
3   LayerNorm(256)  ← Normalize across 256 features
4       ↓
5Normalized: (B, 30, 256)
6       ↓
7   Attention Layer

Why After BiLSTM?

• Stabilizes attention inputs: Prevents attention scores from exploding/vanishing
• Consistent scale: All 30 timesteps have comparable magnitudes
• Faster convergence: Better-conditioned optimization landscape

Parameter Count

LayerNorm for 256 features:

• $\gamma \in \mathbb{R}^{256}$: Scale parameters
• $\beta \in \mathbb{R}^{256}$: Shift parameters
• Total: 512 learnable parameters

PyTorch Configuration

1import torch.nn as nn
2
3# Layer normalization after BiLSTM
4self.layer_norm = nn.LayerNorm(
5    normalized_shape=256,  # Hidden dimension (2 × 128)
6    eps=1e-5,              # Numerical stability
7    elementwise_affine=True  # Learnable γ and β
8)
9
10# In forward pass
11lstm_output, _ = self.bilstm(cnn_features)  # (B, 30, 256)
12normalized = self.layer_norm(lstm_output)    # (B, 30, 256)

Summary

In this section, we integrated layer normalization with the BiLSTM:

1. Layer vs Batch norm: Layer norm is per-sample, batch-independent
2. Formula: Normalize across features, then scale and shift
3. RNN advantages: Handles variable lengths, small batches, sequential structure
4. Placement: After BiLSTM output, before attention
5. Parameters: 512 (γ and β for 256 dimensions)

Looking Ahead: We have designed all components of the BiLSTM encoder: two bidirectional layers with dropout and layer normalization. The next section brings everything together with the complete PyTorch implementation.

With all BiLSTM components designed, we now implement the complete encoder module.