Layer Normalization - Transformer Architecture

Layer Normalization

What you'll build

LayerNorm over the last (feature) axis, with a learned scale gamma and shift beta. Every Transformer block uses it.

Background

BatchNorm normalizes each feature over the batch, which couples examples together and breaks down for variable-length sequences and batch size 1 at inference. LayerNorm normalizes each token's feature vector on its own, so it behaves identically in training and inference and needs no running statistics. In pre-norm Transformers (GPT-2 onwards) it sits at the start of every residual branch: x = x + Attn(LN(x)). LLaMA-style models use the cheaper RMSNorm (see SwiGLU Feed-Forward Block in LLM Internals).

Formula

For each position (each vector x along the last axis, length D):

mu      = mean(x)                          # scalar per position
sigma^2 = mean((x - mu)^2)                 # population (biased) variance
y       = gamma * (x - mu) / sqrt(sigma^2 + eps) + beta

Inputs

• X: float array (..., D), e.g. (seq_len, D) or (batch, seq_len, D)
• gamma, beta: float arrays (D,)
• eps: small float, default 1e-5

Expected Return Format

Array with the same shape as X.

Edge cases and conventions

• Statistics are over axis=-1 only, never over the batch or sequence axes.
• Use the biased variance (np.var default, ddof=0), as PyTorch's nn.LayerNorm does.
• eps goes inside the square root. A constant row gives x - mu = 0, so the output is just beta.

Constraints

NumPy only; use keepdims=True so the statistics broadcast for any number of leading axes.

Follow-up

• Write the backward pass. Why do the gradients with respect to x sum to zero along the feature axis?
• Pre-norm vs post-norm: which one trains stably without learning-rate warmup, and why?

Input: layer_norm(np.array([[1.0, 2.0, 3.0, 4.0]]), np.ones(4), np.zeros(4))
Output: [[-1.3416, -0.4472, 0.4472, 1.3416]]

Input: layer_norm(np.array([[2.0, 2.0, 2.0, 2.0]]), np.ones(4), np.array([0.0, 0.0, 0.0, 1.0]))
Output: [[0.0, 0.0, 0.0, 1.0]]

import numpy as np

def layer_norm(X, gamma, beta, eps=1e-5):
    """
    Apply layer normalization over the last axis.

    Args:
        X: Input (..., D), e.g. (batch, seq_len, D) or (seq_len, D)
        gamma: Scale (D,)
        beta: Shift (D,)
        eps: Small constant added to the variance

    Returns:
        Array with the same shape as X
    """
    # Your code here
    pass