Chapter 2.4 β€” Normalization and Stabilization

Contents

  1. The stack has converged; can we actually train it deep?
  2. Why deep networks need normalization at all
  3. LayerNorm mechanics, and why not BatchNorm
  4. Pre-LN vs. Post-LN: where you put the normalization changes everything
  5. RMSNorm: normalization with the mean-centering removed
  6. Interview angle
  7. Self-check questions
  8. Sources

1. The stack has converged; can we actually train it deep?

Chapter 2.3 closed by settling on decoder-only stacks of masked self-attention and feed-forward sublayers as the dominant shape for large-scale generative LLMs, and by pointing out that this stack isn't yet stable to train at real depth without help. That claim deserves to be taken literally rather than as a rhetorical flourish. The Transformer, like essentially every deep network architecture that came before it, does not simply get better the more layers you stack on top of each other by default β€” beyond a certain depth, naive stacking makes training harder, not easier, and can make it fail outright. Every Transformer block described in Chapters 2.1 through 2.3 was quietly drawn with a component this book hasn't examined yet: a normalization layer, wrapped around each sublayer, without which the residual stacking that makes deep Transformers possible at all doesn't actually work in practice.

This chapter is about that component. It's easy to treat normalization as a minor implementation detail β€” a line of code you add because the reference implementation has it β€” but the specific choices involved (what to normalize over, and where in the residual path to place the normalization) turn out to be responsible for whether a 96-layer Transformer trains stably at all, and they are exactly the kind of detail that separates an engineer who has internalized why the architecture works from one who has only used it. Once this is on the table, only one component of the block remains unexamined β€” the feed-forward sublayer, which Chapter 2.5 takes up to close out Part II.

2. Why deep networks need normalization at all

The core problem normalization addresses predates the Transformer by years and is a general property of deep networks trained by gradient descent: as data flows forward through many stacked layers, and as each layer's parameters are updated based on gradients that flowed backward through all the layers above it, the distribution of activations at any given layer tends to drift and shift over the course of training, partly because every layer's input distribution depends on the parameters of every layer before it, all of which are changing simultaneously. Ioffe and Szegedy, who introduced batch normalization for convolutional networks, described this phenomenon as "internal covariate shift" β€” each layer is effectively trying to learn a stable mapping while the statistics of its own input keep moving underneath it, which slows learning and forces the use of smaller, more conservative learning rates than you'd otherwise want to use.

In a residual architecture specifically β€” and the Transformer is a residual architecture through and through, since every sublayer's output is added back to its input rather than replacing it β€” there's a second, related failure mode worth naming explicitly. If you stack many residual blocks without any renormalization, the scale of the running sum can grow roughly with the number of blocks, since each block adds its own contribution on top of an ever-growing accumulated signal β€” $x_L = x_0 + \sum_{l=0}^{L-1}\text{Sublayer}_l(x_l)$, a sum of $L$ terms whose norm has no built-in ceiling; deep enough, and the activations reaching the later layers, and the gradients flowing back through them, can become numerically extreme β€” too large or too small to support stable, well-conditioned gradient updates. Normalization layers address both problems at once: by explicitly rescaling activations to have controlled statistics at various points in the network, they keep the signal each layer receives in a consistent, well-behaved range regardless of what's happening in the layers around it, and regardless of how deep the stack is, which is precisely the property that makes very deep stacking practical rather than merely theoretically possible.

3. LayerNorm mechanics, and why not BatchNorm

Batch normalization, the earlier and more established technique, normalizes each activation by the mean and variance computed across the batch dimension β€” for a given feature, you compute its mean and variance over all examples in the current minibatch, and use those statistics to normalize that feature for every example. This works well for the fixed-size, densely populated feature maps typical of convolutional networks, but it fits sequence data with variable length badly for a specific, concrete reason: a batch of variable-length sequences requires padding to a common length, and computing per-feature statistics across the batch dimension at a given sequence position means mixing together, in the same statistic, positions from real tokens in some sequences and padding tokens in others (or in the case of autoregressive training, statistics that would need to differ depending on how far into generation the model is), which is a poor match to what a per-position, per-token computation like attention actually needs. Batch statistics also behave differently at training time (computed from the current batch) versus inference time (typically using a running average), which introduces its own train/inference mismatch, and becomes even more awkward when inference happens one token at a time, as it does during autoregressive generation, where there may be no meaningful "batch" at all.

Layer normalization, introduced by Ba, Kiros, and Hinton, sidesteps all of this by normalizing across the feature dimension instead of the batch dimension, computed independently for each individual token. For a single token's activation vector $x \in \mathbb{R}^d$, LayerNorm computes the mean $\mu = \frac{1}{d}\sum_{k=1}^d x_k$ and variance $\sigma^2 = \frac{1}{d}\sum_{k=1}^d (x_k - \mu)^2$ across that token's own $d$ features, then normalizes: $$\hat{x} = \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}}, \qquad \text{LN}(x) = \gamma \odot \hat{x} + \beta$$ where $\gamma$ and $\beta$ are learned per-feature scale and shift parameters (allowing the network to undo the normalization partially or entirely if that's what's useful), and $\epsilon$ is a small constant added for numerical stability. Because this statistic is computed per token, independent of every other token in the sequence and independent of every other example in the batch, it requires no batch-level bookkeeping, behaves identically at training and inference time, and handles variable-length sequences and single-token autoregressive generation without any special-casing at all β€” exactly the properties a sequence model with variable-length inputs and one-token-at-a-time inference actually needs.

4. Pre-LN vs. Post-LN: where you put the normalization changes everything

Knowing that some form of normalization belongs in each Transformer sublayer still leaves an important design decision open: where, exactly, relative to the residual connection, does the normalization go? The original Transformer, as described in Vaswani et al., uses what's now called the Post-LN configuration: the sublayer (self-attention or feed-forward) is applied to the input, added back via the residual connection, and only then normalized β€” $$x_{l+1} = \text{LN}\big(x_l + \text{Sublayer}(x_l)\big)$$ so the normalization sits on the main residual path itself, after each addition. This has a real, empirically observed drawback in very deep stacks: because normalization is applied after the residual addition, the residual path itself is not left in its raw, undisturbed form β€” every block's output gets renormalized before being handed to the next block, which means the identity-like path that residual connections are supposed to provide (letting gradients flow essentially unchanged through arbitrarily many layers) is periodically disrupted by a normalization operation sitting directly in that path.

Xiong et al.'s analysis of this issue, published in 2020, makes the argument precisely: they show that in the Post-LN configuration, the expected magnitude of gradients at the layers nearest the input grows or is otherwise poorly controlled as the number of layers increases, which is exactly the kind of instability that makes very deep Post-LN Transformers require careful learning-rate warmup schedules to train at all, and can make them fail to train past a certain depth even with warmup. Pre-LN rearranges the same two components so the sublayer operates on a normalized copy of the input, and the residual addition happens after, on the raw (unnormalized) path: $$x_{l+1} = x_l + \text{Sublayer}\big(\text{LN}(x_l)\big)$$ Now the main residual path β€” the sequence of additions carrying signal from the very first layer to the very last β€” never passes through a normalization operation at all; only the input to each sublayer's internal computation is normalized. This preserves the property that made residual connections effective for very deep networks in the first place: an essentially unimpeded additive path from input to output, along which gradients can flow with a magnitude that doesn't depend on network depth the way Post-LN's does. Xiong et al.'s theoretical and empirical analysis shows Pre-LN Transformers can be trained stably without the learning-rate warmup schedule that Post-LN Transformers typically require, and can be trained successfully at depths where Post-LN configurations become unstable.

This stability doesn't come entirely for free. There is evidence, discussed in the literature on this comparison, that Pre-LN's improved trainability comes with some cost in final representational strength relative to a Post-LN model that can actually be trained to convergence at the same depth β€” the normalization applied before every sublayer, on every layer, constrains the range of functions each sublayer can express slightly more than applying it only once per block after the residual sum. In practice, this cost has been judged well worth paying at the scale of modern LLMs: a slightly less expressive but reliably trainable configuration beats a theoretically more expressive one that can't be trained past a moderate depth without extensive tuning, and Pre-LN (or close variants of it) is now the standard choice in essentially every large decoder-only LLM.

5. RMSNorm: normalization with the mean-centering removed

Zhang and Sennrich's root mean square normalization asks a different kind of question about LayerNorm: given that LayerNorm does two separate things β€” recentering the activations by subtracting the mean, and rescaling them by dividing by the standard deviation β€” is the recentering step actually necessary for the stabilizing benefit LayerNorm provides, or is the rescaling doing essentially all of the useful work on its own? Their answer, backed by both argument and experiment, is that the rescaling β€” controlling the overall magnitude, or root-mean-square, of the activation vector β€” is the component actually responsible for LayerNorm's stabilizing effect, and the mean-centering step can be dropped with little to no loss in model quality, while saving real computation.

Concretely, RMSNorm replaces LayerNorm's normalization with $$\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_{k=1}^d x_k^2}, \qquad \text{RMSNorm}(x) = \gamma \odot \frac{x}{\text{RMS}(x)}$$ which rescales the activation vector by its root-mean-square magnitude directly, with no mean subtraction anywhere, and typically no learned shift parameter $\beta$ either (only the learned per-feature scale $\gamma$ is kept). This is a small change on paper but a meaningful one computationally: dropping the mean computation and the associated subtraction removes a full reduction operation from every normalization call, in every sublayer, in every layer, for every token β€” at the scale of models with tens or hundreds of billions of parameters processed over trillions of training tokens, this adds up to a real, measurable reduction in training and inference cost, with Zhang and Sennrich reporting speedups on the order of 10-40% for the normalization computation itself depending on hardware, at essentially no cost to model quality on the tasks they evaluated. This efficiency-for-negligible-quality-cost tradeoff is exactly the kind of argument that tends to win out at LLM scale, which is why RMSNorm, typically paired with Pre-LN placement, is the normalization scheme used in most modern large open-weight language models.

A more recent, complementary technique targets a different failure mode than the ones RMSNorm and Pre-LN placement address: attention logits themselves growing large enough β€” especially at long context lengths or large scale β€” to destabilize training even when the residual stream itself is well-behaved. QK-norm, used in models like Gemma 2, normalizes the query and key vectors (typically with RMSNorm) before computing their dot product inside attention, which bounds the dot product's magnitude regardless of how large the individual query and key components might otherwise grow, targeting the attention score computation directly rather than the residual stream that ordinary normalization placement addresses. A related fix, logit soft-capping, caps attention logits (and, separately, the final output logits before the softmax that produces next-token probabilities) with a smooth bounded function instead of leaving them unbounded, keeping rare extreme values from ever reaching the numerically fragile parts of the softmax. Neither technique replaces RMSNorm or Pre-LN; both are narrowly targeted stabilization measures layered on top, and that's itself a pattern worth noticing β€” as models get pushed to greater depth, scale, and context length, new specific numerical failure modes keep surfacing that the original LayerNorm/Pre-LN toolkit didn't anticipate, and the field's response has consistently been to add a small, targeted fix rather than redesign normalization from scratch.

Taken together, the choices covered in this chapter β€” LayerNorm's per-token feature-dimension statistics in place of BatchNorm's batch statistics, Pre-LN placement to preserve a clean residual path through arbitrary depth, and RMSNorm's simplification to keep only the rescaling that matters β€” are what actually make it possible to stack the decoder-only architecture from Chapter 2.3 to the depths modern frontier models use, stably and efficiently. That leaves the block itself nearly complete: attention as the core mechanism, positional schemes to restore the ordering information attention discards, a choice of architectural shape suited to generative modeling, and normalization to make deep stacks of that shape actually trainable. The one piece still described only in passing is the other sublayer sitting inside every block β€” the position-wise feed-forward network, which despite getting a single sentence in most treatments holds the majority of a dense model's parameters. Chapter 2.5 gives it the attention it deserves, and with it Part II is done.

6. Interview angle

"Why doesn't batch normalization work well for Transformers, mechanistically?" A strong answer identifies that BatchNorm computes statistics across the batch dimension for each feature, which requires a well-defined, consistently populated batch of examples at that computation point. With variable-length sequences, padding contaminates those statistics, and single-token autoregressive inference at generation time may not have any meaningful batch structure at all. It should also mention the training/inference statistic mismatch (batch statistics vs. running averages) as a second, independent problem BatchNorm introduces that LayerNorm avoids by computing statistics per token, independent of the batch.

"Explain, with the actual equations, the difference between Pre-LN and Post-LN, and why one trains more stably at depth." A strong answer writes both update rules β€” $x_{l+1} = \text{LN}(x_l + \text{Sublayer}(x_l))$ for Post-LN versus $x_{l+1} = x_l + \text{Sublayer}(\text{LN}(x_l))$ for Pre-LN β€” and explains that in Pre-LN, the residual path itself (the running sum of $x_l$ terms) is never passed through a normalization operation, preserving the near-identity gradient path that residual connections are meant to provide. In Post-LN, normalization sits directly on that path after every addition, which Xiong et al. show leads to gradient magnitudes at early layers that scale poorly with depth, requiring careful warmup and limiting how deep the network can practically go.

"What specifically does RMSNorm remove from LayerNorm, and why is that safe to remove?" The answer should state precisely that RMSNorm removes the mean-subtraction (recentering) step, keeping only the rescaling by root-mean-square magnitude, and that Zhang and Sennrich's argument and experiments indicate the rescaling is what's actually responsible for the stabilizing effect on training, with recentering contributing comparatively little. The efficiency motivation should be stated concretely: removing a reduction operation from every normalization call, across every token, every layer, every training step, is a meaningful saving at LLM scale even if the effect per-call looks small.

"If Pre-LN is more stable, why wasn't it just used from the start? What's the tradeoff?" A strong answer notes that the original Transformer used Post-LN, and that Pre-LN's stability gain doesn't come entirely free β€” the literature has found some evidence that Pre-LN networks, layer for layer, are somewhat less representationally expressive than a Post-LN network of the same size that can actually be trained to convergence, because normalizing before every sublayer's computation constrains the functions each sublayer can express a bit more than a single normalization after the residual sum. The strong answer frames the actual industry choice correctly: at the depths and scales relevant to modern LLMs, trainability at all reliably wins over a theoretical expressiveness edge that's hard to realize in practice.

"Internal covariate shift was originally described for convolutional networks. Does the concept still apply cleanly to Transformers, and does that matter for how you'd explain normalization's role here?" A good answer treats this as a nuanced question rather than a yes/no: the general intuition (activation statistics drifting layer to layer and over training, making optimization harder) transfers conceptually, but the specific mechanism BatchNorm was designed around (batch statistics for convolutional feature maps) doesn't transfer at all to variable-length sequence data, which is exactly why LayerNorm's redesign around per-token statistics was necessary rather than simply reusing BatchNorm. It's fine, and arguably correct, to note that the field's understanding of exactly why normalization helps has evolved and is not fully settled, and to frame it as a genuinely open research question rather than pretending there's a single, complete theoretical account.

7. Self-check questions

  1. What two distinct problems does normalization address in a very deep residual network, and why does simply stacking residual blocks without normalization make both problems worse as depth increases?
  2. Why is BatchNorm's per-feature, per-batch statistic a poor fit for variable-length sequence data and single-token autoregressive inference, specifically?
  3. Write out LayerNorm's formula and explain, in terms of what it computes and over which dimension, why it avoids the specific problems you identified in question 2.
  4. Write the Post-LN and Pre-LN update equations side by side, and explain precisely which computational path differs between them.
  5. Why does keeping the residual path free of normalization operations (as Pre-LN does) help gradients flow better through very deep stacks, connecting this back to what residual connections are supposed to provide in the first place?
  6. What specific component of LayerNorm does RMSNorm remove, and what's the argument for why removing it doesn't meaningfully hurt model quality?
  7. Looking back across Chapters 2.1 through 2.4, list the specific problem each chapter's central idea solved (recurrence's sequential bottleneck, attention's permutation invariance, architectural shape and objective mismatch, and deep-stack trainability), and explain how each solution's necessity followed causally from the previous chapter's design choice.

8. Sources

  • Ba, J. L., Kiros, J. R., & Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450. https://arxiv.org/abs/1607.06450
  • Xiong, R., Yang, Y., He, D., Zheng, K., Zheng, S., Xing, C., Zhang, H., Lan, Y., Wang, L., & Liu, T.-Y. (2020). On Layer Normalization in the Transformer Architecture. ICML 2020. arXiv:2002.04745. https://arxiv.org/abs/2002.04745
  • Zhang, B., & Sennrich, R. (2019). Root Mean Square Layer Normalization. NeurIPS 2019. arXiv:1910.07467. https://arxiv.org/abs/1910.07467
  • Henry, A., Dachapally, P. R., Pawar, S. S., & Chen, Y. (2020). Query-Key Normalization for Transformers. Findings of EMNLP 2020. arXiv:2010.04245. https://arxiv.org/abs/2010.04245
  • Gemma Team, Riviere, M., et al. (2024). Gemma 2: Improving Open Language Models at a Practical Size (QK-norm and logit soft-capping in production). arXiv:2408.00118. https://arxiv.org/abs/2408.00118

Self-check quiz

Test your understanding of this chapter with a short quiz β€” a mix of concept questions and quick calculations.

Why is BatchNorm a poor fit for Transformers processing variable-length sequences?

Explanation: Batch statistics mix real-token and padding-token positions across variable-length sequences, and behave differently at training time (batch stats) versus inference time (running average) β€” a mismatch LayerNorm avoids entirely.

What does LayerNorm normalize over, in contrast to BatchNorm?

Explanation: LayerNorm computes mean and variance over one token's d features, so it requires no batch-level bookkeeping and behaves identically at training and inference time.

What is the key structural difference between Pre-LN and Post-LN placement?

Explanation: Post-LN: x_{l+1} = LN(x_l + Sublayer(x_l)). Pre-LN: x_{l+1} = x_l + Sublayer(LN(x_l)). Pre-LN keeps the running residual sum free of normalization, preserving a near-identity gradient path.

What does RMSNorm remove from LayerNorm's computation?

Explanation: RMSNorm keeps only the rescaling β€” dividing by RMS(x) β€” dropping mean subtraction (and typically the shift parameter Ξ²), saving a full reduction operation per call at negligible quality cost.

For a single token's activation vector $x = (1, 2, 3, 4)$, what is the LayerNorm mean $\mu = \frac{1}{d}\sum_k x_k$ used in the computation?

Explanation: ΞΌ = (1+2+3+4)/4 = 10/4 = 2.5.

For the same vector $x = (1, 2, 3, 4)$ (mean 2.5), what is the LayerNorm variance $\sigma^2 = \frac{1}{d}\sum_k (x_k - \mu)^2$?

Explanation: σ² = [(1-2.5)Β² + (2-2.5)Β² + (3-2.5)Β² + (4-2.5)Β²] / 4 = [2.25+0.25+0.25+2.25]/4 = 5/4 = 1.25.

For a vector $x = (3, 4)$, what is $\text{RMS}(x) = \sqrt{\frac{1}{d}\sum_k x_k^2}$? (Round to 2 decimal places.)

Explanation: RMS(x) = √((9+16)/2) = √12.5 β‰ˆ 3.54.

Without any normalization, suppose 10 stacked residual blocks each add a contribution of magnitude 2 to a running sum starting at 0. What is the magnitude of the sum after all 10 blocks?

Explanation: 10 Γ— 2 = 20 β€” illustrating how the running sum's scale grows roughly with the number of blocks when nothing renormalizes it.