Normalization in Transformers

In interviews related to roles where a lot is tinkering with LLMs and their architecture, asking “Why do Transformers use layer normalization instead of batch normalization?” is common. Interviewers use it to probe your deep understanding of neural network training dynamics and architectural design. It’s not just trivia—the question tests if you grasp how and why certain layers are chosen in cutting-edge models. Candidates who understand the nuances of normalization layers demonstrate they can reason about model behavior under different training and inference conditions rather than just memorizing facts. By discussing this question, you can think critically about model internals, trade-offs, and practical constraints in real-world settings (like training on many GPUs or doing autoregressive decoding).

Beyond simply naming transformers that use layer norm, the interviewer wants to see if you can explain the mechanics and motivations behind that design choice. Are you aware of how batch normalization (BN) and layer normalization (LN) work differently, especially when it comes to Transformers? Can you articulate why BN’s assumptions break in a Transformer’s context? The interviewer tests knowledge of normalization formulas, sequence-model training (including teacher-forcing and masking), and distributed training issues.

Ideally, you’ll discuss topics like when each norm computes statistics, how that affects training vs. inference, and why those differences matter for sequence data and parallel computation. A strong answer weaves together an understanding of batch vs. layer norm and practical concerns (batch size, variable sequence lengths, causal inference, etc.) in modern transformer-based systems.

What exactly is normalization for Transformers?

In machine learning, normalization broadly means rescaling data so its values fall into a common range or distribution. For inputs, normalization might map features into [0,1] or standardize them to mean 0 and variance 1, helping gradient descent converge faster. Inside neural networks, normalization layers (like BN or LN) perform a similar standardization on the layer’s activations. These layers typically subtract the mean and divide by the standard deviation of certain activations. The goal is to stabilize training by preventing internal covariate shift—i.e., large shifts in a layer’s input distribution from one batch to the next, which can slow training and cause gradients to explode or vanish.

Normalization comes in many flavors (batch, layer, instance, group norm, etc.), but we focus on batch vs. layer for Transformers. Batch normalization (BN) computes the mean and variance of each feature across the batch. Layer normalization (LN) computes the mean and variance across features of one example. A key difference is that BN’s behavior depends on having a meaningful batch of data, whereas LN works on each sample independently​. This difference will be crucial when we look at sequence models.

To make this distinction clearer, consider a toy example. Imagine two input vectors (e.g., token embeddings):

With BN, the mean and variance are computed per feature across the batch:

• Feature-wise mean = [(2+1)/2, (4+3)/2, (6+5)/2] = [1.5, 3.5, 5.5]

With LN, each vector (row) is normalized independently:

• x1 → mean = 4.0, std = 1.63 → x1' = [−1.225, 0, 1.225]

• x2 → mean = 3.0, std = 1.63 → x2' = [−1.225, 0, 1.225]

So BN normalizes each feature by comparing across tokens, while LN normalizes across features within a single token.

This distinction becomes crucial for Transformers, where we often process one token at a time or use variable-length sequences. In the next few sections, we’ll explore why LN fits these constraints better.

What is batch normalization in Transformers?

Batch normalization (BN) normalizes activations across the batch for each feature. It requires computing statistics from a mini-batch of examples. In contrast, the Transformer architecture often processes one token (or a few) at a time, making BN less applicable. Batch normalization became popular in computer vision: it computes the mean and variance of each feature (channel) over all samples in a mini-batch, then scales and shifts the normalized output. This is straightforward when all inputs (e.g., images) are the same shape and have a large batch. In practice, BN often improves training speed and allows higher learning rates, because each layer sees inputs with a stable distribution.

However, applying BN inside a Transformer has several challenges:

• Dependency on batch statistics: By design, BN needs multiple examples. It uses the batch ...