Training Stability › Why Training Becomes Unstable

Before reaching for solutions, it is worth understanding the underlying mechanics. Training instability is not a single phenomenon but a family of related problems, each with its own signature and cause.

Neural network loss landscapes are high-dimensional and highly non-convex. Most of the time, gradient descent moves through these landscapes smoothly, making incremental progress. But occasionally, the optimizer wanders into regions where curvature is extremely high: a small step in parameter space produces a large change in loss. These high-curvature regions cause loss spikes.

The curvature of the loss landscape is described mathematically by the Hessian matrix $\mathbf{H}$, where each entry $H_{ij} = \frac{\partial^2 \mathcal{L}}{\partial \theta_i \partial \theta_j}$ captures how the gradient of the loss with respect to parameter $i$ changes as parameter $j$ moves. The largest eigenvalue of $\mathbf{H}$, often called the sharpness $\lambda_{\max}$, determines how large a learning rate is safe. Specifically, gradient descent is guaranteed to decrease loss only when:

$$\eta < \frac{2}{\lambda_{\max}}$$

where $\eta$ is the learning rate and $\lambda_{\max}$ is the largest eigenvalue of the Hessian. When the optimizer encounters a region where $\lambda_{\max}$ is much larger than expected, the effective learning rate becomes too large for that region, and a single update can push parameters past the minimum and far up the other side of the curvature.

This explains why loss spikes often happen suddenly after many stable steps. The optimizer can spend thousands of iterations in well-behaved regions, then stumble into a sharp curvature cliff that sends loss skyrocketing. The transition is not gradual because the loss landscape itself has abrupt changes in curvature that have no visible warning in the loss curve.

One insight from recent research on sharpness-aware minimization (SAM) is that the loss landscapes of large language models tend to contain sharp minima, meaning regions where the loss is low but the curvature is high. These sharp minima are not just bad for generalization; they are also fragile during training because the same property that makes them sharp makes them prone to instability. Flat minima, where a wide basin surrounds the optimal point, are more stable because the region of safe learning rates is much larger.

In deep networks, gradients flow backward through many layers. As covered in earlier chapters on backpropagation and vanishing gradients, this can lead to gradients either shrinking exponentially (vanishing) or growing exponentially (exploding). Exploding gradients are the more acute stability threat because they cause immediate, dramatic parameter updates.

The fundamental issue is that in a network with $L$ layers, each backward pass involves $L$ matrix multiplications. If the singular values of the weight matrices are consistently above 1, gradients compound multiplicatively and can reach astronomically large values before a single parameter update is applied. For a simple linear network with weight matrices $\mathbf{W}_1, \mathbf{W}_2, \ldots, \mathbf{W}_L$, the gradient of the loss with respect to the first layer involves a product of the form:

$$\frac{\partial \mathcal{L}}{\partial \mathbf{W}_1} \propto \mathbf{W}_L^T \mathbf{W}_{L-1}^T \cdots \mathbf{W}_2^T \frac{\partial \mathcal{L}}{\partial \mathbf{h}_L}$$

where $\mathbf{h}_L$ is the output of the final layer. If the spectral norm of each weight matrix is $\sigma > 1$, then this product can grow as $\sigma^{L-1}$, which becomes enormous for deep networks. This exponential amplification is precisely the exploding gradient problem.

Recurrent neural networks are especially vulnerable because the same weight matrix is applied repeatedly for each time step. As covered in the BPTT chapter, gradients scale as $(\mathbf{W}^T)^T$ where $T$ is the sequence length, and if $\|\mathbf{W}\| > 1$, the gradient norm grows exponentially with sequence length. This is one historical reason why RNNs were notoriously hard to train on long sequences before gradient clipping became standard practice.

Modern transformers are less susceptible to the pure exponential blowup because residual connections provide gradient highways that bypass layers, and layer normalization prevents activation magnitudes from drifting. But they are not immune. Very deep transformers (48+ layers), long sequence lengths, and the presence of attention softmax can still produce gradient explosions under the right conditions.

A less discussed but important source of instability is data heterogeneity. When a batch contains outlier samples with very high loss (rare but extreme examples), the gradient from that batch carries an unusually large signal. If the learning rate is tuned for typical batches, this outsized gradient can cause a destabilizing update.

This is particularly relevant in language modeling, where certain sequences (very long documents, repeated tokens, unusual characters) can produce loss values far outside the typical distribution. Consider a batch that happens to sample a document consisting of thousands of repetitions of a single rare token. The model will assign very low probability to each token given its context (because the pattern is unusual), resulting in a very high cross-entropy loss. The gradient from this batch is dominated by this single anomalous document and is not representative of the general distribution the model should learn.

These outlier batches are often the direct trigger for loss spikes that practitioners observe when training on large, diverse web-scraped corpora. The spike is not a fundamental property of the model or optimizer; it is an artifact of sampling a particularly difficult batch. Data quality filtering therefore improves training stability rather than merely making the training process easier to manage.

Modern optimizers like Adam maintain momentum estimates for each parameter. These estimates are exponential moving averages of past gradients, and they allow the optimizer to move smoothly even when individual gradient estimates are noisy. But they also carry history: a period of large gradients can corrupt momentum estimates in ways that take many steps to decay.

Adam's first and second moment estimates are defined as:

$$\begin{aligned} m_t &= \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ v_t &= \beta_2 v_{t-1} + (1 - \beta_2) g_t^2 \end{aligned}$$

where $g_t$ is the gradient at step $t$, $\beta_1 \approx 0.9$ is the first moment decay, and $\beta_2 \approx 0.999$ is the second moment decay. The second moment decay $\beta_2 = 0.999$ means that the influence of a gradient from 1000 steps ago is reduced by a factor of $0.999^{1000} \approx 0.37$. So a corrupt gradient signal persists in the optimizer state for roughly 1000 steps before it decays substantially.

When a loss spike occurs, the Adam optimizer's momentum buffers absorb large gradient values. Even after the spike resolves, the inflated momentum can continue causing larger-than-appropriate updates for many subsequent steps. This creates a cascading instability where a single bad batch can degrade training for hundreds of iterations. Practitioners who simply continue training after a spike often see the loss partially recover and then remain high for a long time, which is exactly the signature of corrupted optimizer state slowly decaying.

A fourth, often underappreciated cause of instability is numerical precision. Modern large model training uses mixed-precision arithmetic, combining float32 for some computations with float16 or bfloat16 for others. Float16 has a maximum representable value of 65504, and values larger than this cause overflow to infinity. The underflow boundary is even more constraining for small values, with subnormal numbers starting below approximately $6 \times 10^{-5}$.

In practice, this means that attention logits for very long sequences (where the dot-product values can be large), or softmax outputs for peaked distributions, or intermediate activations in very deep networks, can easily overflow float16 representation. The result is NaN (not a number) values that propagate forward through the network and produce NaN gradients, which corrupt optimizer state catastrophically.

The bfloat16 format, used by default in many modern TPU and GPU training setups, has the same exponent range as float32 (so no overflow risk for typical values) but only 7 bits of mantissa instead of float32's 23 bits. This makes bfloat16 much more resistant to overflow, but it can still produce instability through accumulated rounding errors in numerically sensitive operations like layer normalization.

Understanding these four sources of instability, sharp loss landscape curvature, exploding gradients through depth, outlier batches, and numerical precision limitations, gives you a principled framework for diagnosing problems and choosing appropriate interventions.