You have built the block, counted it, and stacked it. Now train it. With 12 layers, almost anything works. With 48 layers and a learning rate that worked at 12, the loss falls for a few hundred steps, then shoots up and never comes back. Or it trains for a week and then, at step 80,000, spikes by a full nat for no visible reason. Large-model logbooks are full of such events. PaLM's authors report around 20 loss spikes while training their largest model, each handled by rewinding to an earlier checkpoint and skipping the batches that followed, and in DeepNet's translation experiments an ordinary post-LN Transformer with 6 encoder and 6 decoder layers trained fine while the same model with 18 and 18 diverged.

This section is about why deep Transformers become unstable and what keeps them stable. Some causes we have met already, gradients that grow with depth (§5.1) and the post-LN gradient profile (§5.2). Here we add three more, an initialization that does not respect depth, the first steps of Adam, and attention that collapses onto single tokens. We finish with the practical question of how deep a model should be at all. The full operational toolkit for pretraining runs, including loss-spike recovery, z-loss, and monitoring, is the subject of §8.8. This section explains the mechanisms that toolkit is built on.

Initialization that respects depth

GPT-2 initializes every weight matrix from $\mathcal{N}(0, 0.02^2)$ (the value in OpenAI's released code, inherited from the first GPT paper), with one exception. The paper scales the weights of residual layers at initialization by one over the square root of the number of residual layers. In code (nanoGPT and many others) that means the output projection of each residual branch, attention's $\mathbf{W}_O$ and the FFN's $\mathbf{W}_2$, gets standard deviation

since a model with $L$ blocks has $2L$ residual branches.

Here is the derivation, which is §5.1's argument with numbers. In a pre-LN model each branch reads a normalized input, so its write at initialization is roughly a random vector whose per-coordinate variance is proportional to $\sigma_{\text{out}}^2$, about $c\,\sigma_{\text{out}}^2$ for a constant $c$ set by the branch's internal width and activation. The $2L$ writes are close to independent, so variances add,

With a fixed $\sigma_{\text{out}}$ the stream's variance grows linearly with depth. With $\sigma_{\text{out}} = \sigma_0/\sqrt{2L}$, the total added variance is $c\,\sigma_0^2$ no matter how deep the model is.

In plain words, if $2L$ branches each add a random push, make each push $\sqrt{2L}$ times smaller so that their sum has the same size for every depth.

In real numbers, GPT-2 small ($L = 12$) starts its residual projections at $0.02/\sqrt{24} \approx 0.0041$, and GPT-2 XL ($L = 48$) at $0.02/\sqrt{96} \approx 0.0020$. The rule touches only the two matrices whose output is added to the stream (both are called c_proj in GPT-2's code). Why those, rather than every matrix in the branch? Shrinking the last matrix shrinks the write in proportion and changes nothing else. Shrinking the input projections instead would also shrink the attention logits, which depend on $\mathbf{W}_Q\mathbf{W}_K^\top$ quadratically, and would move the FFN's pre-activations toward the nearly linear part of GELU, changing what the branch computes rather than only how loudly it writes. Section 0.7 previewed this rule, and §6.2 applies it when it initializes a whole GPT.

The same scaling keeps the backward pass bounded. In §5.1 we showed that a stack of $S$ residual branches with Jacobian gain $\alpha$ amplifies gradients by about $(1+\alpha^2)^{S/2}$, and that $\alpha^2 \propto 1/S$ keeps this near $e^{1/2}$. Shrinking each branch's output matrix by $\sqrt{2L}$ shrinks its gain the same way.

Unscaled, the stream ends up 1.9 times its starting size at 12 layers and about 4.55 times at 96. Scaled, it stays at 1.05 for every depth, matching the formula. Why is the base standard deviation 0.02? It is a fixed choice that GPT-2 inherited from the first GPT, close to $1/\sqrt{d}$ for widths in the low thousands. A fixed number cannot be right for every width, and principled width-dependent choices (including different scalings for the embedding, hidden, and output layers) are what μP derives in §9.5.

import numpy as np

def gpt2_init_block(d, n_layer, rng):
    raise NotImplementedError

The family of depth-aware schemes

GPT-2's rule is one of several ways to keep a deep residual stack near the identity at the start. They all shrink, gate, or rebalance the branches relative to the skip path.

DeepNorm's constants come from requiring that a single optimizer step changes the model's output by a bounded amount, whatever the depth. That is the right way to think about every entry in the table. Stability is about how much the function changes per update, and depth multiplies small per-layer changes into large ones unless something compensates.

Warmup, and why Adam needs it here

Almost every Transformer recipe starts with a learning-rate warmup. GPT-3 ramped linearly over its first 375 million tokens, and LLaMA 2 over 2,000 steps. Section 5.2 showed one reason, post-LN's large top-layer gradients. There is a second reason that applies to every architecture trained with Adam (§0.6), and it is visible in the very first step.

Adam updates each coordinate by $-\eta\,\hat m_t/(\sqrt{\hat v_t} + \varepsilon)$, where $\hat m_t$ and $\hat v_t$ are bias-corrected moving averages of the gradient and the squared gradient, and $\varepsilon$ is a small constant. As §0.6 showed, at step 1 bias correction makes $\hat m_1 = g$ and $\hat v_1 = g^2$ exactly, so

In plain words, Adam's first step moves every parameter by the full learning rate, in the direction of its gradient's sign, whether that gradient is meaningful or pure noise.

For the next few dozen steps $\hat v_t$ is an average over very few squared gradients, so it is a noisy estimate, and dividing by its square root keeps the steps large and erratic. Liu et al. (2019) quantified this, showing that the variance of Adam's adaptive step size is large early in training, and proposed RAdam to correct it. Warmup does the same job more crudely. It keeps $\eta$ small until $\hat v_t$ has seen enough gradients to be trusted, and it lets the model move away from its fragile initial state gently. The derivation of warmup length from Adam's variance is in §8.7.

A gradient of $10^{-4}$ and a gradient of $10^{2}$ produce the same first step, $-10^{-3}$. On a mostly-noise gradient stream, the first ten steps are on average more than twice as large as steps later in training. Warmup exists to tame exactly those early steps.

Attention entropy collapse

The third failure lives inside attention. Each row of an attention matrix is a probability distribution $\mathbf{a}$ over the $T$ keys, and its entropy,

measures how spread out it is. Uniform attention has entropy $\log T$, and attention on a single key has entropy 0. Zhai et al. (2023) tracked the entropy of every head during training across vision, speech, translation, and language models, and found a consistent pattern. Training instability, whether oscillating loss or outright divergence, was accompanied by pathologically low attention entropy. They called it entropy collapse.

Why would low entropy hurt? Look at what the softmax passes back. The Jacobian of $\mathbf{a} = \operatorname{softmax}(\mathbf{z})$ with respect to the logits is

which you derived in §4.3. When $\mathbf{a}$ approaches a one-hot vector $\mathbf{e}_i$, this becomes $\operatorname{diag}(\mathbf{e}_i) - \mathbf{e}_i\mathbf{e}_i^\top = \mathbf{0}$. The gradient with respect to the queries and keys vanishes, so the head cannot learn its way out, while a small change in the weights can still flip which key wins and change the output abruptly. A collapsed head is both stuck and brittle.

Where collapse comes from, and a bound you can derive

Attention logits are $z_n = \mathbf{x}_m\mathbf{W}_Q\mathbf{W}_K^\top\mathbf{x}_n^\top/\sqrt{d_k}$. Nothing in a pre-LN Transformer stops $\mathbf{W}_Q$ and $\mathbf{W}_K$ from growing during training, and if they grow, the logits grow in proportion to the product of their norms. Sharp logits mean low entropy. Here is a bound that makes the dependence exact.

Suppose every logit in a row lies in $[-s, s]$. The largest attention weight is at most what it would be with one logit at $+s$ and all others at $-s$,

Shannon entropy is never smaller than $-\log\max_n a_n$, because $\sum_n a_n\log(1/a_n) \ge \sum_n a_n\log(1/\max a) = \log(1/\max a)$. Therefore

In plain words, if the logits are confined to a range of size $s$, attention cannot collapse much. As the range grows, the guaranteed entropy falls off like $(T-1)e^{-2s}$, exponentially fast.

Zhai et al. prove a bound of the same exponential form stated through the weights, which falls exponentially fast as the spectral norm of the attention logits grows (they control that norm through $\lVert\mathbf{W}_K\mathbf{W}_Q^\top\rVert_2$ and the size of the inputs), and show it is tight. Both bounds say the same thing. Keep the effective scale of the logits under control and entropy cannot collapse.

The table tells the story. As the logit scale grows, entropy falls from near $\log 64 = 4.16$ toward zero and the average top weight rises past 0.9. The softmax Jacobian first grows, since attention becomes selective, and then shrinks as the rows harden into one-hot vectors and stop passing gradient back. The bound is loose for any single row, but it is always below the true entropy.

Two cures, σReparam and QK-norm

Zhai et al.'s fix, σReparam, reparameterizes every linear layer's weight as

where $\sigma_1(\mathbf{W})$ is the largest singular value of $\mathbf{W}$ (its spectral norm) and $\gamma$ is a learned scalar (the code below starts it at 1). The spectral norm is estimated cheaply by the power iteration of §0.10, run on $\mathbf{W}^\top\mathbf{W}$ without forming it, one step per training step with a persistent vector, as in spectral normalization (Miyato et al., 2018). Now the scale of every layer is controlled by a single number $\gamma$ that the optimizer moves deliberately, rather than by the drift of thousands of weights. With it, Zhai et al. trained Vision Transformers to competitive accuracy without warmup, weight decay, LayerNorm, or adaptive optimizers.

The more widely used cure is QK-norm, which normalizes each head's query and key vectors (with LayerNorm or RMSNorm) before the dot product. The logit $\mathbf{q}\cdot\mathbf{k}/\sqrt{d_k}$ then has a magnitude fixed by the norms' gains, so it cannot grow just because $\mathbf{W}_Q$ and $\mathbf{W}_K$ did. Dehghani et al. (2023) found that vision Transformers of around 8 billion parameters diverged because their attention logits grew so large that attention became almost one-hot, with near-zero entropy, which is entropy collapse in the wild. QK-norm fixed it for their 22-billion-parameter model. Wortsman et al. (2023) reproduced this "attention logit growth" in small models trained at high learning rates and showed the same fix works. QK-norm is now standard in many recent models, and §8.8 covers it with its relatives.

import numpy as np

def attention_entropy(scores):
    raise NotImplementedError

import numpy as np

def softmax_jacobian(z):
    raise NotImplementedError

import numpy as np

def spectral_norm(W, n_iter=50, u=None):
    raise NotImplementedError

def sigma_reparam(W, gamma, u, n_iter=1):
    raise NotImplementedError

import numpy as np

def entropy_floor(s, T):
    raise NotImplementedError

def worst_logits(s, T):
    raise NotImplementedError

Logit growth at the output

The final softmax over the vocabulary has its own version of the problem. Softmax is unchanged when a constant is added to every logit, so nothing in the cross-entropy loss pins down the overall level of the output logits. During long runs their normalizer $\log Z = \log\sum_v e^{z_v}$ can drift far from zero, and in 16-bit arithmetic that drift eventually costs precision and stability. PaLM added an auxiliary z-loss, $10^{-4}\cdot(\log Z)^2$, which gently pulls $\log Z$ toward zero, and Wortsman et al. (2023) showed that this output-logit divergence also appears in small models at high learning rates and that z-loss fixes it. Gemma 2 took a blunter approach, soft-capping both attention and final logits with $c\tanh(z/c)$. Both are covered in §8.8. The thread connecting them to attention entropy collapse is that a Transformer has several places where a scale is left free by the loss, and every free scale is a potential runaway.

Depth versus width

Stability is one reason not to make a model arbitrarily deep. It is worth knowing the others, because "how should I shape a model with $N$ parameters?" is a common design question.

Quality depends only weakly on shape. Kaplan et al. (2020) found that at a fixed non-embedding parameter count, loss varies by only a few percent across a wide range of shapes. The aspect ratio $d/L$ could change by a factor of 40 with little effect, and a 6-layer model with $d = 4288$ came within 3% of the loss of GPT-2's 48-layer, $d = 1600$ configuration. This is an empirical finding for the range they studied, not a law, and other work finds that very small models benefit from being deep and thin (Liu et al., 2024, MobileLLM, for models under a billion parameters).

Depth costs latency, width costs nothing extra in sequence. The $L$ blocks run one after another, so time per generated token grows with $L$, while a wider matrix multiply keeps a GPU busier for the same wall-clock time. Wide layers also split cleanly across GPUs with tensor parallelism (§10.4), and deep stacks split into pipeline stages that must be kept full (§10.5).

Depth is harder to train. Every mechanism in this section gets worse with depth, and so does the underuse of late layers in pre-LN models (§5.2).

In practice, large dense models settle on aspect ratios of roughly 100 to 130. GPT-3 175B has $d/L = 12{,}288/96 = 128$, LLaMA-2 7B has $4096/32 = 128$, and LLaMA-2 70B has $8192/80 \approx 102$. GPT-2 small, at $768/12 = 64$, is comparatively deep for its size.

A stability checklist

In PyTorch

Two pieces you will use in real code. GPT-2's initialization as a function applied over a module, and a σReparam linear layer with a persistent power-iteration buffer.

Read more

• Zhai et al. (2023). Stabilizing Transformer Training by Preventing Attention Entropy Collapse. Entropy collapse, the entropy lower bound, and σReparam.
• Wortsman et al. (2023). Small-scale proxies for large-scale Transformer training instabilities. Attention logit growth and output logit divergence, reproduced cheaply.
• Dehghani et al. (2023). Scaling Vision Transformers to 22 Billion Parameters. The QK-norm fix for attention logit growth at scale.
• Wang et al. (2022). DeepNet, Scaling Transformers to 1,000 Layers. DeepNorm, derived from bounding the size of model updates.
• Liu et al. (2019). On the Variance of the Adaptive Learning Rate and Beyond. Why Adam's early steps are unreliable, and RAdam.
• Bachlechner et al. (2020). ReZero is All You Need. Zero-initialized residual gates.
• Kaplan et al. (2020). Scaling Laws for Neural Language Models. Section 3.1 shows how little shape matters at fixed size.

Watch next

▶