24. Layer Normalization From Scratch

Problem

Layer normalization normalizes each individual example's activations across the feature dimension, independent of the rest of the batch — unlike batch norm, which normalizes across the batch for each feature. That batch-independence is exactly why it became the default normalization layer for transformers and sequence models: it works identically whether you're training with batch size 512 or doing greedy decoding one token at a time.

DeepMind's research spans a wide range of normalization variants built on this same idea (RMSNorm, used in Gopher and Chinchilla, drops the mean-centering step entirely), so knowing the exact mechanics of the original LayerNorm formula is the right starting point: normalize over the last dimension using that example's own mean and (biased) variance, then apply a learned per-feature affine transform.

Source: src/24_layer_norm_from_scratch.py

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

>>> x = torch.randn(2, 3, 8)
>>> gamma, beta = torch.ones(8), torch.zeros(8)
>>> layer_norm(x, gamma, beta).shape
torch.Size([2, 3, 8])
>>> layer_norm(x, gamma, beta).mean(dim=-1).abs().max().item() < 1e-5
True

Step-by-Step Approach

1. Compute the per-example mean over the last dimension with keepdim=True, so it broadcasts cleanly against the original shape.
2. Compute the per-example variance over the last dimension with unbiased=False — LayerNorm uses the population variance (divide by d), not the sample variance (divide by d-1).
3. Normalize: (x - mean) / sqrt(var + eps), with eps inside the square root for numerical stability when variance is near zero.
4. Apply the learned affine transform: multiply by gamma and add beta, both shape (d,), broadcasting against every leading dimension of x.
5. Sanity-check with an identity affine (gamma=1, beta=0): the output's per-example mean should be ~0 and variance ~1 along the last dimension.

The entire operation only ever reduces over the last dimension, regardless of how many leading batch/sequence dimensions x has — using dim=-1, keepdim=True throughout is what makes the same three lines of code correct for a 2D (batch, d) input and a 3D (batch, seq, d) input alike.

Reference solution

import torch
import torch.nn.functional as F


def layer_norm(
    x: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, eps: float = 1e-5,
) -> torch.Tensor:
    # per-example statistics over the feature dim only -- keepdim so the
    # subtraction/division below broadcast against x's leading dims cleanly
    mean = x.mean(dim=-1, keepdim=True)
    # unbiased=False: LayerNorm uses the population (1/d) variance, not the
    # sample (1/(d-1)) variance -- using the wrong one is a common bug
    var = x.var(dim=-1, unbiased=False, keepdim=True)

    x_hat = (x - mean) / torch.sqrt(var + eps)
    return x_hat * gamma + beta

Key Functions & Tricks

• Tensor.mean(dim=-1, keepdim=True) — per-example mean over the feature dimension, kept as a broadcastable shape
• Tensor.var(dim=-1, unbiased=False, keepdim=True) — population variance over the feature dimension; the unbiased=False flag is the detail interviewers probe for
• torch.sqrt(var + eps) — eps inside the square root, not just added after, keeps the denominator strictly positive even at zero variance
• F.layer_norm — PyTorch's own built-in, used here only as an independent ground-truth check, not as the required implementation

How to Recognize This Pattern