Layer Normalization & Residual Connections in Transformers

🎯 Learning Objectives

• Explain why Transformers need normalization and residual connections
• Compute layer normalization step by step
• Distinguish layer norm from batch norm
• Understand how residual connections enable deep networks
• Analyze gradient flow through residual paths

📋 Prerequisites

• Basic neural network training concepts
• Gradient descent and backpropagation
• Transformer architecture

1. 📖 Core Content

1.1 Why Normalization?

Batch Normalization vs Layer Normalization

1. Transformers process variable-length sequences — batch statistics are unreliable
2. LayerNorm is independent of batch size
3. LayerNorm works the same at training and inference
4. LayerNorm preserves the per-token representation structure

1.2 Layer Normalization — Formal Definition

• $\mu = \frac{1}{d}\sum_{i=1}^d x_i$ (mean across features)
• $\sigma^2 = \frac{1}{d}\sum_{i=1}^d (x_i - \mu)^2$ (variance across features)
• $\epsilon$: Small constant for numerical stability (e.g., $10^{-5}$)
• $\gamma \in \mathbb{R}^d$: Learnable scale parameter
• $\beta \in \mathbb{R}^d$: Learnable shift parameter

Worked Example

python
Run Inline
# runnable
import numpy as np
def layer_norm(x, gamma=None, beta=None, eps=1e-5):
    """
    Layer Normalization
    Args:
        x: Input array (..., d_model)
        gamma: Learnable scale (d_model,)
        beta: Learnable shift (d_model,)
        eps: Numerical stability constant
    Returns:
        normalized: Layer-normalized output
    """
    mean = np.mean(x, axis=-1, keepdims=True)
    variance = np.var(x, axis=-1, keepdims=True)
    x_norm = (x - mean) / np.sqrt(variance + eps)
    if gamma is not None:
        x_norm = x_norm * gamma
    if beta is not None:
        x_norm = x_norm + beta
    return x_norm
# Test with our example
x = np.array([2.0, 4.0, 1.0, 3.0])
gamma = np.ones(4)
beta = np.zeros(4)
result = layer_norm(x, gamma, beta)
print(f"Input: {x}")
print(f"Normalized: {np.round(result, 4)}")
print(f"Mean after norm: {np.mean(result):.6f}")
print(f"Std after norm: {np.std(result):.6f}")

1.3 Pre-LN vs Post-LN

1.4 Residual Connections

Gradient Flow

1.5 Why Both Are Necessary

1.6 Implementation Example

python
Run Inline
# runnable
import numpy as np
class TransformerBlock:
    """A complete Transformer block with Pre-LN"""
    def __init__(self, d_model, d_ff, num_heads):
        self.d_model = d_model
        self.d_ff = d_ff
        self.num_heads = num_heads
        # Learnable LayerNorm parameters (simplified)
        self.gamma1 = np.ones(d_model)
        self.beta1 = np.zeros(d_model)
        self.gamma2 = np.ones(d_model)
        self.beta2 = np.zeros(d_model)
        # Attention projections (simplified - random init)
        np.random.seed(42)
        self.W_Q = np.random.randn(d_model, d_model) * 0.1
        self.W_K = np.random.randn(d_model, d_model) * 0.1
        self.W_V = np.random.randn(d_model, d_model) * 0.1
        self.W_O = np.random.randn(d_model, d_model) * 0.1
        # FFN weights
        self.W1 = np.random.randn(d_model, d_ff) * 0.1
        self.W2 = np.random.randn(d_ff, d_model) * 0.1
    def layer_norm(self, x, gamma, beta):
        mean = np.mean(x, axis=-1, keepdims=True)
        var = np.var(x, axis=-1, keepdims=True)
        x_norm = (x - mean) / np.sqrt(var + 1e-5)
        return x_norm * gamma + beta
    def forward(self, x):
        # Pre-LN: Norm -> Sublayer -> Residual
        # Sublayer 1: Multi-head Attention
        x_norm1 = self.layer_norm(x, self.gamma1, self.beta1)
        # Simplified attention
        Q = x_norm1 @ self.W_Q
        K = x_norm1 @ self.W_K
        V = x_norm1 @ self.W_V
        scores = Q @ K.T / np.sqrt(self.d_model)
        attn = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
        attn_out = attn @ V @ self.W_O
        x = x + attn_out  # Residual connection
        # Sublayer 2: FFN
        x_norm2 = self.layer_norm(x, self.gamma2, self.beta2)
        ffn_out = np.maximum(x_norm2 @ self.W1, 0) @ self.W2
        x = x + ffn_out  # Residual connection
        return x
# Test
block = TransformerBlock(d_model=64, d_ff=256, num_heads=8)
x = np.random.randn(10, 64)  # 10 tokens, d_model=64
output = block.forward(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {output.shape}")
print(f"Input norm: {np.linalg.norm(x, axis=-1).mean():.4f}")
print(f"Output norm: {np.linalg.norm(output, axis=-1).mean():.4f}")

4. 📐 Key Formulas / Concepts

5. ⚠️ Common Pitfalls

Pitfall 1: Confusing LayerNorm with BatchNorm

• Sequence lengths vary across samples
• Batch size may be 1 during inference
• LayerNorm preserves per-token statistics

Pitfall 2: Placing LayerNorm after the residual addition (Post-LN)

Pitfall 3: Forgetting that LayerNorm has learnable parameters

6. 📝 Practice Questions

Q1: Compute LayerNorm for x=[3, 5, 2, 6] with γ=1, β=0

μ = (3+5+2+6)/4 = 4 σ² = ((3-4)² + (5-4)² + (2-4)² + (6-4)²)/4 = (1+1+4+4)/4 = 2.5 σ = 1.581

x_norm = [(3-4)/1.581, (5-4)/1.581, (2-4)/1.581, (6-4)/1.581] = [-0.632, 0.632, -1.265, 1.265] Q2: If γ=[0.5, 1.0, 1.5, 2.0] and β=[0, 0, 0, 0], what is the output of LayerNorm for x=[3, 5, 2, 6]?

From Q1, x_norm = [-0.632, 0.632, -1.265, 1.265]

y = x_norm * γ + β y = [-0.632×0.5, 0.632×1.0, -1.265×1.5, 1.265×2.0] y = [-0.316, 0.632, -1.898, 2.530] Q3: Why does Pre-LN make training more stable than Post-LN?

In Post-LN: output = LN(x + Sublayer(x)). Gradients must pass through LN's normalization computation, which can amplify/attenuate gradients.

In Pre-LN: output = x + Sublayer(LN(x)). The residual path provides a direct gradient highway (the "1" term in $\frac{\partial L}{\partial x} = \frac{\partial L}{\partial y}(1 + \frac{\partial F}{\partial x})$). Gradients bypass the sublayer entirely.

This direct path ensures that even 100+ layer networks receive gradient signals. Q4: What happens if ε (epsilon) is too large in LayerNorm?

If ε is too large (e.g., 1.0 instead of 1e-5), the normalized values become: x_norm = (x - μ) / √(σ² + 1.0)

For small σ² (uniform activations), the denominator is dominated by ε, not σ. This means:

• Normalized values are much smaller than they should be
• The output becomes less sensitive to input variations
• The model loses representational capacity

Typical ε values: 1e-5 to 1e-8. Q5: In a Transformer with d_model=512, how many learnable parameters does each LayerNorm have?

LayerNorm has γ and β, each of shape (d_model,):

• γ: 512 parameters
• β: 512 parameters
• Total: 1,024 parameters per LayerNorm

The Transformer has 2 LayerNorms per encoder layer + 3 per decoder layer, so for 6 encoder + 6 decoder: Total LN params = (6 × 2 + 6 × 3) × 1024 = 30 × 1024 = 30,720 learnable normalization parameters. Q6: What is RMS Norm and how does it differ from LayerNorm?

RMS Norm (Root Mean Square Layer Normalization) simplifies LayerNorm by removing the mean subtraction:

$\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^d x_i^2 + \epsilon}} \cdot \gamma$

RMS Norm only scales by the RMS, no shift or mean subtraction. It's used in some modern LLMs (like LLaMA) because it's computationally simpler and empirically works as well as full LayerNorm. Q7: If a residual connection is removed from a 12-layer Transformer, what happens during training?

Without residual connections in a 12-layer network:

1. Gradients vanish: The gradient at layer 1 is the product of 12 Jacobians. Each Jacobian has eigenvalues typically < 1, so the product becomes exponentially small.
2. Layer 1 almost doesn't learn: The first layer receives negligible gradient updates.
3. Effective depth is much less than 12: Only the last few layers learn meaningful representations.

This is exactly why residual connections were introduced — they make very deep networks trainable. Q8: How would you implement Pre-LN in a Transformer encoder layer using PyTorch-like pseudocode?

python
Run Inline
class TransformerEncoderLayer:
    def forward(self, x):
        # Pre-LN pattern
        x_norm = layer_norm(x)
        attn_out = self_attention(x_norm)
        x = x + attn_out  # First residual

        x_norm = layer_norm(x)
        ffn_out = ffn(x_norm)  # FFN: ReLU(xW₁+b₁)W₂+b₂
        x = x + ffn_out  # Second residual

        return x

Note the order: LN → Sublayer → Add (residual).

7. 🔗 Cross-References

• Next: Pre-training Objectives (Week 3)
• Related: Normalization techniques in deep learning
• Video: BSDA5004 Week 2 transcripts Join Discord PreviousDecoder & Cross-AttentionNextPre-training Objectives