skip to content

Which axis does LayerNorm normalize over, and why does batch size not change its output?

level: juniorimportance: must knowfreq 72%

answer

  1. Ask which axis the mean runs along
  2. Not the batch axis
  3. Across that one sample's own features
  4. Same arithmetic at train and serve

basics

~20 s

LayerNorm computes a mean and variance across the feature dimensions of each sample on its own, then rescales with a learned per-feature gain and bias. No other sample enters the statistic, so batch size and batch composition cannot change the output.

solid answer

~40 s

LayerNorm pools over the feature axis within a single sample: for one activation vector `x` of length `H` it computes `mu = mean(x)` and `var = mean((x - mu)^2)` over those `H` features, forms `(x - mu)/sqrt(var + eps)`, and applies a learned gain and bias held per feature. Because the sum runs over one sample's own features, the output for a sample is identical whether it is processed alone or in a large batch. That gives three practical properties: it works unchanged at batch size 1, it stores no statistics so training and serving compute the same thing, and padded positions in a ragged batch cannot contaminate a real position's mean. The cost is that the vector's overall level and spread are deleted as information, and you get no batch-noise regularization.

go deeper

for a junior

Be ready to state the axis in one sentence and not confuse it with the batch-axis layer. Knowing that each sample is normalized using only its own feature values is the screening bar here.

for a middle

Explain the arithmetic end to end: the two statistics, the epsilon in the denominator, and that the gain and bias are per-feature parameters shared across samples, not statistics.

for a senior

Show why you would reach for it: single-example serving, tiny or ragged batches, and the absence of a train-versus-serve gap. Mention what you lose -- the overall level of the vector, and any regularization that batch noise was giving you.

for a principal

Own the design consequence: making the function independent of batch composition turns batch size into a pure throughput knob, which changes how you scale training and how reproducible a served prediction is. Be able to argue when that guarantee is worth its costs.

## What a normalization layer has to decide Every normalization layer does the same three things: it computes a location statistic (a mean), computes a spread statistic (a variance or a root mean square), divides the activations by that spread after removing the location, and then hands the network back a learned gain and bias so it can undo the constraint where it needs to. The only real design decision is **which values are pooled into that statistic**. LayerNorm answers: the features of one sample, and nothing else. ## The computation Take a single sample's activation vector `x` with `H` components, where `H` is the size of whatever is being normalized (the feature or channel dimension for a per-position vector, or all of the channel and spatial positions together for a convolutional map). For that sample alone: ``` mu = (1/H) * sum_i x_i var = (1/H) * sum_i (x_i - mu)^2 xhat_i = (x_i - mu) / sqrt(var + eps) y_i = g_i * xhat_i + b_i ``` `eps` is a small constant that keeps the divide finite when a vector is nearly constant. `g` and `b` are learned, one pair per normalized position, and they are shared across every sample -- they are ordinary parameters, not statistics. The sum index `i` runs over the features of one sample. No value from any other sample appears anywhere in the expression. ## What that buys **Batch independence.** The output for a given sample is bit-for-bit the same whether it was processed alone or alongside a thousand others. Batch size becomes purely a throughput and gradient-noise decision, not something that changes the function the network computes. **No stored statistics.** There is nothing to estimate across training and nothing to freeze for inference. The layer performs the identical arithmetic in training and in serving, which removes a whole class of "it was fine in training and wrong in production" bugs from this layer. **Batch size 1 works.** A single-example inference path on a device, or a training run whose sample is a whole 3D volume that barely fits in memory, still has a well-defined statistic: a mean over `H` features exists even when there is only one sample. A statistic pooled over the batch axis has nothing to pool at batch size 1. **Padding cannot leak.** When variable-length inputs are packed into a rectangular batch, the padded positions hold values that mean nothing. A statistic pooled across the batch at a given position mixes those junk values into the mean and variance that real positions are divided by. LayerNorm's statistic for a real position is computed from that position's own features, so a padded neighbour cannot move it. ## What it costs and what it destroys The layer is invariant to a per-sample shift and rescale of its input: feeding `a * x + c` (for any positive `a` and any constant `c` added to every feature) produces exactly the same `xhat`. That is the point -- but it means the overall magnitude and the overall offset of that vector are *deleted* as information before the next layer sees them. If the absolute level of a signal is meaningful, it must reach later layers by some other route, because this layer throws it away every time. Because the mean and variance are functions of all `H` inputs, the backward pass is not elementwise: each input's gradient depends on every output gradient in the same vector. Practically the layer is bandwidth-bound -- it reads and writes the activations a couple of extra times -- rather than compute-bound. It also gives you no free regularization. A statistic pooled over a batch injects noise that depends on which samples happened to share the batch, and that noise acts a little like a regularizer; a per-sample statistic is deterministic, so if you swap toward it you may need to add regularization elsewhere. ## Choosing the axis in practice For per-position feature vectors -- a sequence of vectors, one per element -- normalizing the feature dimension separately for every position is the standard choice, and it is why sequence models with ragged lengths and short batches adopted it. For convolutional feature maps, pooling every channel and every spatial location of the sample into one statistic ties all the channels' scales together, which is often too coarse; that is exactly the gap GroupNorm was designed to fill by pooling over a subset of channels instead. ## Mistakes interviewers listen for Saying it normalizes "over the batch, per feature" -- that is the batch-axis layer, and getting the two backwards is the single most common error on this topic. Claiming the layer keeps a running estimate for inference -- it keeps none, because it never needed one. And claiming batch size still matters "a little" -- for this layer it does not matter at all, which is precisely the property people pick it for.

  • You pack variable-length sequences into a rectangular batch with padding. Why does that not corrupt this layer's statistics?
    Because the mean and variance for a real position are computed from that position's own feature values inside its own sample. Padded positions never enter the sum, so whatever junk they hold cannot move the divisor that real activations are scaled by. A statistic pooled across the batch at a fixed position would mix padding into the estimate.
  • Does the layer behave differently at inference time?
    No. It holds no statistics estimated during training -- only the learned gain and bias, which are ordinary parameters. The same mean, variance and rescale are computed from the incoming sample in both modes, so there is no train-versus-serve gap originating in this layer.
  • What information does this layer throw away?
    The vector's overall offset and overall magnitude. Feeding `a * x + c` for any positive `a` and any constant `c` added to every feature yields the identical normalized vector. If the absolute level of that representation is meaningful to the task, it has to reach later layers by another route, because this layer removes it every time.

Grading each exam paper against its own answers rather than against the rest of the pile: the mark you get does not depend on who else sat that day.

saying these in an interview costs you the question

  • Says it normalizes each feature across the batch
  • Claims it keeps a running estimate for inference
  • Thinks batch size still shifts the output slightly
  • Forgets the learned gain and bias entirely
  • Believes it provides the same noise regularization as batch statistics

context