Core ML

Layer Normalization

Standardise every feature of one example against its own mean and spread, never against the batch, so the same example normalises identically whether it arrives alone or with ninety-nine others.

Why Does This Exist?

Batch normalization works well for image models with large, fixed-size batches, but a sequence model translating sentences hits two problems batch norm was never built for. First, sentences arrive at different lengths — an 8-word sentence and a 40-word sentence don't share a shape, so there's no clean way to line them up along a batch axis and compute one shared statistic per position. Second, generation often runs one sequence at a time, batch size one, and batch normalization's statistics computed from a single example are just that example's own values — no averaging effect at all, and no sensible way to distinguish training statistics from a running average at inference either.

The transformer architecture needed a normalization scheme that never looks across the batch dimension at all — one that computes its mean and variance from a single example's own features, so the same example normalises to the same result whether it's sitting in a batch of one or a batch of a hundred. That's layer normalization, and the "layer" in the name refers to normalizing across a layer's full feature vector for one example, not across the layer's outputs for many examples.

Think of It Like This

Grading on a curve within one exam versus across many exams

Batch normalization is grading on a curve computed across an entire room of students taking the same exam — your score depends on how everyone else in the room did, and if the room has only one student in it, there's no curve to compute. Layer normalization is grading on a curve computed from your own answers alone: take every question you answered, standardise your scores against your own average and your own spread across those questions, and hand back a result that depends on nothing but what you personally wrote.

One student's curved grade under the second scheme comes out identical whether they took the exam alone in an empty room or alongside ninety-nine other students, because the curve was never computed from the room in the first place — only from that one student's own answers.

How It Actually Works

Which axis the statistics come from

For one example's feature vector $x \in \mathbb{R}^d$ — the full set of $d$ features at one position, one token, one row:

\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad y = \gamma \odot \hat{x} + \beta

$\mu$ and $\sigma^2$ are the mean and variance of that example's own $d$ features — one pair of numbers per example, computed by looking across the feature axis rather than the batch axis. $\epsilon$ keeps the division stable, $\gamma$ and $\beta$ are learned per-feature scale and shift, applied after normalization exactly as in batch normalization. The formula looks almost identical to batch norm's; the entire difference is which axis $\mu$ and $\sigma^2$ get computed across.

Why that one change fixes both problems

Because $\mu$ and $\sigma^2$ never touch the batch dimension, a sequence of length 8 and a sequence of length 40 each normalise every one of their own positions independently — there's no shared statistic across positions or across examples that a length mismatch could break. And because nothing here comes from "the other examples in this batch," a batch of size one computes exactly the same normalization a batch of size one hundred would compute for that same example — verified directly in the code below. Training and inference run the identical computation, with no running-average bookkeeping to keep in sync between the two, which is the piece of batch normalization's design that variable-length sequence models found hardest to live with.

Where it sits in a transformer block

Every transformer block wraps both its attention sub-layer and its feed-forward sub-layer in a layer normalization step — historically after each sub-layer's output is added back to the residual stream, though pre-norm vs post-norm covers a since-preferred alternative ordering. Either way, the reason layer normalization specifically was the choice, rather than batch normalization, traces straight back to variable sequence lengths and the small per-device batches typical of large-model training.

Show Me the Code

The same feature vector, normalised alone and again inside a batch of one hundred wildly different examples — checking directly that the result never changes.

1import numpy as np2
3
4def layer_norm(x: np.ndarray, eps: float = 1e-5) -> np.ndarray:5    mean = x.mean(axis=-1, keepdims=True)6    var = x.var(axis=-1, keepdims=True)7    return (x - mean) / np.sqrt(var + eps)8
9
10row = np.array([10.0, 12.0, 8.0, 14.0])11rng_others = [np.random.default_rng(i).normal(size=4) * 50 for i in range(99)]12alone = layer_norm(np.array([row]))[0]13in_a_batch = layer_norm(np.array([row] + rng_others))[0]14
15print(np.round(alone, 4))       # -> [-0.4472  0.4472 -1.3416  1.3416]16print(np.round(in_a_batch, 4))  # -> [-0.4472  0.4472 -1.3416  1.3416]17print(np.allclose(alone, in_a_batch))  # -> True — identical, batch size never entered the formula

Ninety-nine other rows with entirely different scales sit in the same batch, and the first row's normalised output doesn't move by a single digit — nothing in the computation ever reads those other ninety-nine rows.

Watch Out For

Normalizing across the wrong axis by accident

The formula for layer normalization and batch normalization differ only in which axis the mean and variance are computed across, and a transposed tensor or a mistaken axis argument in a hand-written implementation can silently compute one when the other was intended. The code still runs and produces a same-shaped output, so nothing crashes — but a sequence model given batch-axis statistics loses the exact batch-independence property it was chosen for, and length-mismatched batches start producing subtly wrong results that only show up as poor convergence, not an error.

Expecting layer normalization to fix the same internal-covariate-shift story batch norm was originally credited with

Layer normalization is often reached for as a drop-in swap for batch normalization on the assumption it solves the identical problem batch norm was originally described as solving. The two do share the broad effect of stabilising the scale of activations flowing through a network, but they normalise across entirely different axes and were adopted for different, specific reasons — batch independence and variable sequence lengths for layer norm, large fixed-size batches of images for batch norm. Treating them as interchangeable without checking which axis the model actually needs invariance across is a common source of an unexplained accuracy gap after swapping one in for the other.

The Quick Version

Layer normalization computes its mean and variance across one example's own features, never across the batch.
That single change makes it batch-independent: the same example normalises identically whether it's alone or with ninety-nine others.
It handles variable-length sequences and small or single-example batches, which is exactly where batch normalization struggles.
Training and inference run the identical computation, with no running-average bookkeeping to keep synchronized.
Transformer blocks wrap both sub-layers in it, precisely because sequence models need batch independence batch norm can't provide.