Layer Normalization

Start from brute force

The obvious shortcut gets both statistics in one pass from running sums, using the textbook identity var = Σx² / n − mean². It is algebraically correct and costs O(n) time and O(1) extra space, the same as the right answer, so speed is not the problem: in floating point it cancels catastrophically when the values are large compared with their spread. For [1e8+1, 1e8+2, 1e8+3, 1e8+4] it computes a variance of 2.0 instead of 1.25, and with an offset of 1e9 it computes 0.0, so the outer outputs come out near ±474 instead of ±1.34. The other naive version, recomputing the mean and variance inside the per-element loop, is correct but O(n²).

The approach

Compute the statistics in two passes: the mean first, then the variance as the average squared deviation from that mean. Subtracting the mean before squaring keeps the numbers small, so nothing is lost to cancellation, and the result ignores any constant added to every input: [1e8+1, …, 1e8+4] normalizes exactly like [1, 2, 3, 4]. After normalizing, the vector has mean 0 and variance var / (var + eps), which is 1 up to the epsilon; gamma and beta then let the model choose a different scale and center for each feature.

The reference solution does exactly that with plain lists: n = len(x), mean = sum(x) / n, var = sum((xi - mean) ** 2 for xi in x) / n (the population variance, dividing by n), and denom = math.sqrt(var + eps), computed once rather than once per element. The return line builds gamma[i] * (x[i] - mean) / denom + beta[i] for every index i.

Walk through an example

Trace the first example, layer_norm([1, 2, 3, 4], [1, 1, 1, 1], [0, 0, 0, 0]), where gamma is all ones and beta all zeros, so the output is the raw normalized vector.

Complexity

Optimal (two-pass): O(n) time, for one pass to sum, one to accumulate squared deviations and one to build the output, and O(n) space for the returned list (O(1) extra). Brute force: the one-pass sum-of-squares version is also O(n) time and O(1) extra space but numerically unstable; recomputing the statistics for every element is O(n²).

Edge cases to test

• Constant input: [7, 7, 7] → [0.0, 0.0, 0.0]. The variance is 0 and every deviation is exactly 0, so the epsilon only has to keep the denominator nonzero; without it, Python raises ZeroDivisionError and JavaScript produces NaN from 0 / 0.
• Variance below epsilon: [0, 0.002] has variance 1e-6, so eps = 1e-5 dominates and the output is about ±0.3015, not ±1. Epsilon is not negligible for features with a tiny spread.
• The output variance is var / (var + eps), a little under 1: 0.999992 for [1, 2, 3, 4]. A test that compares with exactly 1 needs a tolerance.
• A single element: layer_norm([5], [3], [0.5]) → [0.5]. The deviation is 0, so only beta survives.
• A large common offset: [1e8+1, 1e8+2, 1e8+3, 1e8+4] gives the same output as [1, 2, 3, 4] with the two-pass formula; the one-pass Σx² / n − mean² version does not.
• An empty vector raises ZeroDivisionError at sum(x) / n. Normalizing zero features is undefined, so reject it explicitly.

Common mistakes

• Normalizing over the wrong axis — layer norm is per-example over features, unlike batch norm which normalizes over the batch.
• Forgetting the epsilon, which divides by zero on a constant (zero-variance) vector.
• Dropping the learned γ/β, so the layer cannot rescale or re-center after normalizing.
• Using the sample variance, dividing by n − 1: it does not match the statement or PyTorch’s nn.LayerNorm, and it divides by zero for a single element.

Follow-ups an interviewer asks next

How does this extend to a batch of token vectors?

For activations shaped (batch, sequence, d), normalize each length-d vector on its own: mean and variance over the last axis only. gamma and beta have length d and are shared by every token. Because no statistic crosses tokens or examples, one token’s output never depends on what else is in the batch, so training and inference behave identically.

What is RMSNorm, and why do many LLMs use it?

RMSNorm drops the mean: y = gamma · x / sqrt(mean(x²) + eps), with no beta. That removes one reduction per vector and one parameter vector, the RMSNorm paper reports comparable results with less computation, and LLaMA uses it. The trade-off is that it rescales without re-centering, so a vector with a large constant offset keeps that offset.

How would you compute it in a single pass, for example over a stream?

Use Welford’s algorithm. For the k-th value, set delta = x - mean, then mean += delta / k and m2 += delta * (x - mean); at the end, var = m2 / n. It reads each value once, like the sum-of-squares shortcut, but it accumulates deviations from the running mean, so it avoids the catastrophic cancellation.

Where does layer norm go in a Transformer block?

The original Transformer applied it after each residual addition (post-norm). GPT-2 moved it to the input of each sub-block (pre-norm), leaving the residual path as an unnormalized sum, and added one more layer norm after the final block. Pre-norm is generally easier to train in deep stacks, while post-norm typically needs a learning-rate warmup to stay stable.

Where this shows up

Layer norm (and its cheaper cousin RMSNorm) stabilizes training in every Transformer, wrapped around the attention and feed-forward blocks. Because it is computed per example and not across the batch, it is independent of batch size — which is exactly why it, not batch norm, suits variable-length sequence models.

Reference solution

import math

def layer_norm(x, gamma, beta, eps=1e-5):
    n = len(x)
    mean = sum(x) / n
    var = sum((xi - mean) ** 2 for xi in x) / n
    denom = math.sqrt(var + eps)
    return [gamma[i] * (x[i] - mean) / denom + beta[i] for i in range(n)]


print(layer_norm([1, 2, 3, 4], [1, 1, 1, 1], [0, 0, 0, 0]))  # ~[-1.34, -0.45, 0.45, 1.34]
print(layer_norm([7, 7, 7], [1, 1, 1], [0, 0, 0]))            # -> [0.0, 0.0, 0.0]

function layerNorm(x: number[], gamma: number[], beta: number[], eps = 1e-5): number[] {
  const n = x.length;
  const mean = x.reduce((a, b) => a + b, 0) / n;
  const varr = x.reduce((a, b) => a + (b - mean) ** 2, 0) / n;
  const denom = Math.sqrt(varr + eps);
  return x.map((xi, i) => (gamma[i] * (xi - mean)) / denom + beta[i]);
}

console.log(layerNorm([1, 2, 3, 4], [1, 1, 1, 1], [0, 0, 0, 0])); // ~[-1.34, -0.45, 0.45, 1.34]
console.log(layerNorm([7, 7, 7], [1, 1, 1], [0, 0, 0]));          // -> [0, 0, 0]