35. LayerNorm From Scratch

Problem

Every Transformer sub-layer, and most SSM block variants, wraps its core computation in a normalization layer to keep activation scales stable across dozens of stacked layers and thousands of training steps. LayerNorm is the classic choice: it re-centers and re-scales each token's feature vector to zero mean / unit variance — per token, not per batch — then applies a learned affine transform so the network can undo the normalization if it needs to.

Implement LayerNorm's forward pass from the raw math, normalizing over the last dimension only, independently per position.

Source: src/35_layer_norm_from_scratch.py

def layer_norm(x: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
    ...

Examples:
>>> x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
>>> gamma = torch.ones(4)
>>> beta = torch.zeros(4)
>>> layer_norm(x, gamma, beta).shape
torch.Size([1, 4])

Step-by-Step Approach

1. Identify the normalization axis: LayerNorm normalizes over the feature dimension (the last dim), independently for every batch/sequence position — unlike BatchNorm, which normalizes across the batch dimension.
2. Compute the per-position mean: x.mean(dim=-1, keepdim=True). Keeping the dim lets it broadcast back against the original tensor.
3. Compute the per-position variance with unbiased=False — LayerNorm uses the biased (population) estimator, dividing by d not d - 1.
4. Normalize: subtract the mean, divide by sqrt(var + eps). eps lives inside the square root, added to the variance before it's taken — that's what keeps the division finite when a row's variance is exactly zero.
5. Apply the learned affine transform: multiply by gamma and add beta, both shape (d,), broadcasting over every leading dimension.
6. Sanity check the constant-row edge case: if every value in a row is identical, its variance is exactly 0, and the output should still be finite (no NaN/Inf) purely because of eps.

The key insight is that mean and variance are reductions over the last axis only, computed independently per token — getting the axis wrong (e.g. normalizing over the batch dim instead) is the single most common bug, and it's exactly the kind of mistake that silently corrupts training without crashing anything.

Reference solution

def layer_norm(x: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
    # mean/var are computed over the last dim only, keepdim so they broadcast
    # back against x -- everything else (batch, sequence, ...) is untouched.
    mean = x.mean(dim=-1, keepdim=True)                       # shape (..., 1)
    # unbiased=False: LayerNorm uses the *biased* (population) variance,
    # dividing by d not d-1 -- using the unbiased estimator here is a classic
    # subtle bug that only shows up as a small numerical mismatch.
    var = x.var(dim=-1, unbiased=False, keepdim=True)         # shape (..., 1)
    # eps sits inside the sqrt, added to the variance -- this is what keeps
    # the division finite even when a row's variance is exactly 0.
    x_norm = (x - mean) / torch.sqrt(var + eps)                # shape (..., d)
    return x_norm * gamma + beta                               # shape (..., d)

Key Functions & Tricks

• x.mean(dim=-1, keepdim=True) — per-token mean, keepdim preserves broadcastability.
• x.var(dim=-1, unbiased=False, keepdim=True) — biased variance, matching nn.LayerNorm's exact convention (divide by d, not d-1).
• torch.sqrt(var + eps) — eps inside the square root, not added after, is what prevents a divide-by-zero on a constant-valued row.
• Broadcasting * gamma + beta — a (d,) affine pair applied across every leading dim without an explicit loop or reshape.
• F.layer_norm(x, (d,), weight=gamma, bias=beta, eps=eps) — the library reference used to build ground-truth expected values in the test suite.

How to Recognize This Pattern