Depth became trainable through a handful of unglamorous fixes

The paper describing PaLM, Google's 540-billion-parameter language model, includes a section that reads more like an operations log than a research result. During training, the loss spiked roughly twenty times, even with gradient clipping switched on, at irregular intervals. Smaller models in the same family did not show the spikes. The team's remedy was procedural. They restarted from a checkpoint taken about 100 steps before each spike and skipped somewhere between 200 and 500 batches of data covering the spike, and the loss did not spike again at that point. They did not think bad data was the cause, since replaying the same batches from an earlier checkpoint produced no spike. The trouble arose from a particular combination of a batch and the state of the parameters.

In 2022, in other words, a team with some of the best training infrastructure in the world managed the stability of a deep network by watching a graph and rolling back. That tells you something about how settled the problem of training deep networks is. Depth is the defining feature of deep learning, and for most of the field's history it was also the feature that made networks fail. The reasons are mathematical and simple. The fixes that overcame them are modest. Because the same fixes sit inside every large model, anyone who trains or fine-tunes one, or who pays someone to, benefits from knowing what they are and which of them is likely to be the culprit when a run goes wrong.

Gain compounds, and it compounds both ways

A network with L layers passes a signal forward through L transformations, and passes a gradient backward through the same L. If each layer changes the typical magnitude of the signal by a factor g, the output after L layers is scaled by g^L. For a network of a hundred layers the arithmetic is merciless:

Per-layer gain	After 50 layers	After 100 layers
0.90	0.005	0.00003
0.95	0.08	0.006
1.00	1	1
1.05	11	130
1.10	117	13,800

A layer that shrinks signals by only 10% leaves about 0.003% of the signal at the bottom of a hundred-layer stack. A layer that amplifies by 10% produces a signal fourteen thousand times too large. Training adjusts weights in proportion to gradients, so vanishing gradients stop the early layers from learning and exploding ones throw the whole network off a cliff. The knife-edge is narrow: only gains very close to 1 work at depth.

This is the old vanishing-gradient problem, analyzed by Sepp Hochreiter in 1991 and by Yoshua Bengio, Patrice Simard and Paolo Frasconi in 1994, in the context of recurrent networks, where depth comes from time. The same compounding applies to a feed-forward stack, and for the better part of two decades it capped how deep a network could usefully be. The solution, when it arrived, was not one idea. It was four, each addressing a different place where the gain can drift.

Initialization sets the starting gain

The first lever is the scale of the random weights at the start. If the weights are too small, every layer shrinks the signal, and if too large, it grows. Xavier Glorot and Yoshua Bengio derived in 2010 a scale that keeps the variance of activations and gradients roughly constant for symmetric activation functions. Rectified linear units, which zero out half their inputs, need a different one. Kaiming He and colleagues derived it in 2015 for rectifiers: weights with variance 2 divided by the number of inputs, so that the factor per layer is 1 on average. With that initialization they trained extremely deep rectified networks from scratch, and a network built on it reached 4.94% top-5 error on ImageNet, which they reported as the first result to surpass the 5.1% human benchmark.

Compare the table above. An initialization whose variance is off by 10% in each layer is a gain of about 1.05 in amplitude per layer, and by layer 100 it has the signal 130 times too large. The right constant, 2 and not 1.8 or 2.2, is the difference between a network that trains and one that returns NaN in the first step.

Normalization keeps the gain from drifting

Correct initialization fixes the start, but the weights change during training and the gain drifts. Sergey Ioffe and Christian Szegedy's batch normalization (2015) renormalizes each layer's inputs to zero mean and unit variance across the mini-batch, with learned scale and shift. The headline result was practical: a state-of-the-art image classifier reached the same accuracy in 14 times fewer training steps, with much higher learning rates tolerated and less sensitivity to initialization. An ensemble of batch-normalized networks reached 4.9% top-5 validation error on ImageNet.

The original explanation was that normalization reduces "internal covariate shift," the drift in each layer's input distribution. That account has been contested. Shibani Santurkar and colleagues argued in 2018 that the benefit comes from making the loss landscape smoother, and that the distribution-shift story does not hold up in experiments. Whatever the mechanism, the effect is real, and the lesson is that engineering results can run ahead of theory by years.

Batch normalization depends on batch statistics, which is awkward for sequences of varying length and for small batches. Layer normalization, by Jimmy Ba, Jamie Kiros and Geoffrey Hinton (2016), computes the statistics over the features of a single example, and that is what Transformers use. Its simplified form, RMS normalization (Zhang and Sennrich, 2019), drops the mean subtraction and is common in recent language models.

Where the normalization sits matters as much as whether it exists. The original Transformer placed it after each residual addition ("post-LN"). Ruibin Xiong and colleagues showed at ICML 2020 that this arrangement leaves large expected gradients near the output at initialization, which makes training unstable unless the learning rate is warmed up slowly from near zero. Moving the normalization inside the residual branch ("pre-LN") gives well-behaved gradients at initialization, and in their experiments the warm-up stage could be dropped with comparable results and much less tuning. GPT-2 had already placed the normalization at the input of each sub-block, and Xiong's analysis explained why that helps. Most large language models since use pre-LN or a close variant, and this small move of one operation is the "normalization moved to the front" that appears in any list of Transformer refinements.

Skip connections turn a product into a sum

The fix that made hundred-layer networks routine was the residual connection. Kaiming He and colleagues observed in 2015 that deeper plain networks gave worse results than shallower ones even on the training set, which meant the problem was optimization and not overfitting. Their remedy was to let each block learn a correction to its input, so that the layer computes x + f(x) instead of f(x). The paper's abstract reports residual networks of up to 152 layers on ImageNet, eight times deeper than the VGG networks of the time with lower complexity, an ensemble error of 3.57% that won the 2015 ILSVRC classification task, and analysis on CIFAR-10 with 100 and 1000 layers.

The reason it works is visible in the arithmetic of the table. Without the skip, the Jacobian of each layer is some matrix whose size can be anything, and the product of a hundred of them is the problem. With the skip, each layer's Jacobian is the identity plus a correction, and a gradient can travel from the output to any earlier layer along the identity path with a gain of exactly 1. The deep network behaves, for the purposes of gradient flow, like a sum of paths of many different lengths, including very short ones. Andreas Veit and colleagues showed that this has a striking consequence: removing a single block from a trained residual network barely hurts it, unlike deleting a layer from a plain network, and residual networks behave like ensembles of relatively shallow ones.

A related trick makes the point sharply. Priya Goyal and colleagues, in their 2017 work on training ImageNet models with very large batches, initialized the last normalization scale in every residual block to zero. At the start of training each block then outputs nothing, so the whole hundred-layer network is exactly the identity function, with a gain of 1 by construction. The blocks learn to contribute as training proceeds, and depth is used only where the data asks for it. Starting shallow in effect and growing deeper in practice is a good summary of what the whole family of fixes achieves.

A residual stream also has a quieter problem. If each of L blocks adds a contribution with variance 1, the stream's variance after 100 blocks is about 100, so its standard deviation is ten times larger than at the start. The GPT-2 paper addressed this by scaling the weights of residual layers at initialization by 1 over the square root of the number of layers, which keeps the stream's magnitude roughly constant. Anthropic researchers later popularized a way of seeing the whole Transformer around this object: every attention and feed-forward block reads from the residual stream and writes back into it, and the stream is the channel through which layers communicate.

The optimizer and the number format

Two more pieces complete the picture. Adam, by Diederik Kingma and Jimmy Ba (2014), keeps a running estimate of each parameter's gradient mean and variance and scales its step accordingly, so parameters whose gradients are consistently tiny are not left behind by those whose gradients are large. This makes training far less sensitive to the leftover differences in gradient scale between layers, and AdamW, which separates weight decay from the gradient update, became the default for Transformers. The bfloat16 number format keeps the 8 exponent bits of 32-bit floating point and shortens the mantissa, which gives it the same dynamic range at half the memory and avoids most of the overflow problems of 16-bit floats.

What is still hard

PaLM's loss spikes show that the stack of fixes narrows the problem and leaves a remainder. At very large scale, with enormous learning rates, huge batches and low-precision arithmetic, rare events occur often enough to matter, and a run that costs millions of dollars cannot wait to find out what they were. The practical response has been operational: frequent checkpoints, monitoring of loss and gradient norms, the ability to roll back and skip data, and a team on call.

Two cheap instruments catch most trouble early. Gradient clipping, introduced by Razvan Pascanu, Tomas Mikolov and Yoshua Bengio in 2013 for recurrent networks, rescales the whole gradient whenever its norm exceeds a threshold, which turns a rare explosion into one ordinary step. And a plot of gradient norm per layer, taken every few hundred steps, shows the compounding directly: as a rule of thumb, if the first layer's gradient norm is a thousand times smaller than the last layer's, or larger, the gain per layer is off and no amount of learning-rate tuning will rescue the run.

The same fixes explain much of what goes wrong in small training runs. When a fine-tuning job diverges, the list of suspects follows from the sections above: a learning rate too high for the batch size, no warm-up, no gradient clipping, a precision setting that overflows, an optimizer state restored incorrectly, a data order that put a pathological batch at the wrong moment. Teams that fine-tune models in-house can build a short routine around those: log the loss and the gradient norm at every step, checkpoint often enough that a rollback costs minutes, and start every new configuration at a lower learning rate than the one that last worked.

What the history suggests for leaders

The history carries two lessons for a leader who is not training models. One is about where capability comes from. The jump from 19-layer to 152-layer networks owed little to a new theory of intelligence and a great deal to arithmetic about gain and a few lines of code. Progress in this field often looks like plumbing until it is seen from outside. The people who make training stable are therefore worth more than their job titles suggest, and a vendor that can describe how its runs stay stable, including the failures, is giving better evidence than one that shows only a benchmark table.

The other is about the pace of understanding. Batch normalization worked for years before anyone could explain why, and the first explanation turned out to be at best incomplete. A good deal of what is deployed rests on techniques that practitioners trust because they work, not because they are fully understood. That is not a reason to distrust them, since the empirical record is strong, but it is a reason to measure outcomes directly instead of reasoning from first principles about what a model must be doing.

Sources

1. Chowdhery, A., Narang, S., Devlin, J., et al. (2022). PaLM: Scaling Language Modeling with Pathways. arXiv:2204.02311, section on training instability.
2. Bengio, Y., Simard, P., Frasconi, P. (1994). Learning long-term dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks 5(2).
3. Glorot, X., Bengio, Y. (2010). Understanding the difficulty of training deep feedforward neural networks. AISTATS 2010.
4. He, K., Zhang, X., Ren, S., Sun, J. (2015). Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. ICCV 2015. arXiv:1502.01852.
5. Ioffe, S., Szegedy, C. (2015). Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. ICML 2015. arXiv:1502.03167.
6. Santurkar, S., Tsipras, D., Ilyas, A., Madry, A. (2018). How Does Batch Normalization Help Optimization? NeurIPS 2018. arXiv:1805.11604.
7. Ba, J. L., Kiros, J. R., Hinton, G. E. (2016). Layer Normalization. arXiv:1607.06450.
8. Xiong, R., Yang, Y., He, D., et al. (2020). On Layer Normalization in the Transformer Architecture. ICML 2020. arXiv:2002.04745.
9. He, K., Zhang, X., Ren, S., Sun, J. (2016). Deep Residual Learning for Image Recognition. CVPR 2016. arXiv:1512.03385.
10. Veit, A., Wilber, M., Belongie, S. (2016). Residual Networks Behave Like Ensembles of Relatively Shallow Networks. NeurIPS 2016. arXiv:1605.06431.
11. Radford, A., Wu, J., Child, R., et al. (2019). Language Models are Unsupervised Multitask Learners (GPT-2). OpenAI technical report.
12. Elhage, N., Nanda, N., Olsson, C., et al. (2021). A Mathematical Framework for Transformer Circuits. Anthropic, transformer-circuits.pub.
13. Kingma, D. P., Ba, J. (2015). Adam: A Method for Stochastic Optimization. ICLR 2015. arXiv:1412.6980.
14. Hochreiter, S. (1991). Untersuchungen zu dynamischen neuronalen Netzen. Diploma thesis, Technische Universität München.
15. Zhang, B., Sennrich, R. (2019). Root Mean Square Layer Normalization. NeurIPS 2019. arXiv:1910.07467.
16. Loshchilov, I., Hutter, F. (2019). Decoupled Weight Decay Regularization. ICLR 2019. arXiv:1711.05101.
17. Pascanu, R., Mikolov, T., Bengio, Y. (2013). On the difficulty of training recurrent neural networks. ICML 2013. arXiv:1211.5063.
18. Goyal, P., Dollár, P., Girshick, R., et al. (2017). Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour. arXiv:1706.02677.