Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Training Stability › Stability Techniques

Michael BrenndoerferJanuary 28, 202610 min read

Language AI HandbookAll chapters

Part 32 of 67Training Optimization

Chapter 237Training Stability › Stability Techniques

Reading section 6 of 15 · Stability Techniques · 10 minSection 6 of 15 · 10 min

Section 6 of 15 · 10 min read

Training Stability › Stability TechniquesOne section · Sections in this chapter6 / 15 ↓

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Beyond gradient clipping, several architectural and algorithmic choices improve training stability. Understanding these techniques and why they work allows you to apply them appropriately rather than blindly copying configuration choices from other projects.

Layer Normalization PlacementLink Copied

Normalization layers are powerful stabilizers because they prevent the magnitude of activations from drifting unboundedly through forward passes. As covered in the layer normalization chapter, LayerNorm operates independently per sample, making it compatible with variable-length sequences and small batch sizes, unlike BatchNorm.

The placement of layer normalization has a substantial effect on stability. Two conventions exist:

Post-norm: LayerNorm applied after the residual connection, as in the original "Attention Is All You Need" paper. The residual path passes through unnormalized, which can lead to large activation magnitudes in deep models.
Pre-norm: LayerNorm applied before the self-attention or FFN sub-layer, inside the residual path. The residual connection accumulates clean (unnormalized) values, giving a stable gradient highway through depth.

Pre-norm transformers are substantially more stable for deep architectures. GPT-2 and later models use pre-norm by default. With post-norm, it is common to need aggressive learning rate warmup (thousands of steps) to avoid early instability. Pre-norm tolerates larger learning rates and shorter warmup schedules.

The mathematical reason is that with pre-norm, gradient flow through the residual connection is direct and unobstructed. In a pre-norm transformer, each layer computes:

\mathbf{x}^{(l+1)} = \mathbf{x}^{(l)} + f\bigl(\text{LN}(\mathbf{x}^{(l)})\bigr)

where $\mathbf{x}^{(l)}$ is the residual stream at layer $l$ , $\text{LN}(\cdot)$ is layer normalization, and $f(\cdot)$ is the attention or FFN sub-layer. The gradient of the loss with respect to $\mathbf{x}^{(l)}$ is:

\frac{\partial \mathcal{L}}{\partial \mathbf{x}^{(l)}} = \frac{\partial \mathcal{L}}{\partial \mathbf{x}^{(l+1)}} \cdot \left(1 + \frac{\partial f(\text{LN}(\mathbf{x}^{(l)}))}{\partial \mathbf{x}^{(l)}}\right)

where:

$\mathbf{x}^{(l)}$ : input to layer $l$ (the residual stream value)
$f(\cdot)$ : the attention or FFN sub-layer function applied after LayerNorm
$\text{LN}(\cdot)$ : layer normalization applied before the sub-layer
The "1 +" term: the direct residual path that always passes gradient through unchanged

The "1 +" in the gradient equation means the residual connection provides an additive path that is always present regardless of the sub-layer's behavior. In pre-norm, this path carries the unchanged residual stream values, creating a gradient highway from output to input that is independent of depth. In post-norm, this path goes through an additional normalization operation that can attenuate gradients in early training.

A recent variant, RMSNorm (root mean square normalization), provides similar stability benefits to LayerNorm but with reduced computational cost. It removes the mean subtraction step of standard LayerNorm, relying only on scale normalization. LLaMA and many subsequent models use RMSNorm with pre-norm placement, combining the stability benefits of both techniques.

Residual Connections and Initialization ScalingLink Copied

Residual connections, introduced to enable very deep network training, also contribute directly to stability by providing gradient highways through depth. But they interact with initialization in a subtle way that large model practitioners have learned to address carefully.

In a transformer with $L$ layers, if each sub-layer adds a residual contribution with unscaled weights, the variance of the residual stream grows with depth. Specifically, if each sub-layer output has variance $\sigma_f^2$ , the residual stream after $L$ layers has variance approximately $L \sigma_f^2$ , growing linearly with depth. For models with 24, 48, or 96 layers, this can produce activations with large magnitudes even at initialization, before any training.

The solution adopted in GPT-2 and variants is to scale down the initialization of the final projection in each sub-layer by a factor of $\frac{1}{\sqrt{2L}}$ , where $L$ is the number of transformer layers. Specifically, the output projection of the attention block and the second linear layer of the FFN block are initialized with standard deviation $\frac{\sigma}{\sqrt{2L}}$ , where $\sigma$ is the default initialization standard deviation. This ensures that the total variance contributed to the residual stream across all layers remains approximately constant:

\text{Var}[\mathbf{x}^{(L)}] \approx \text{Var}[\mathbf{x}^{(0)}] + L \cdot \frac{\sigma_f^2}{2L} = \text{Var}[\mathbf{x}^{(0)}] + \frac{\sigma_f^2}{2}

This is initialization scale invariance with respect to depth: adding more layers does not change the activation magnitudes at initialization, which is a desirable property for numerical stability.

Learning Rate Scheduling and WarmupLink Copied

As discussed in the learning rate warmup and cosine schedule chapters, the learning rate schedule has a direct effect on stability. The warmup phase improves performance and is also required for stable training in many architectures.

Early training is the most fragile phase. Model weights are far from their final values, the loss landscape is poorly characterized, and the optimizer has no momentum history to guide it. A large learning rate at initialization sends the optimizer on large random walks that can easily land in regions of extreme curvature.

Learning rate warmup addresses this by starting with an extremely small learning rate (often 1-10% of the peak value) and linearly increasing it over the first few thousand steps. This allows the optimizer to take small, cautious steps while building up accurate momentum estimates. By the time the learning rate reaches its peak, the model is in a better-conditioned region of parameter space and the optimizer's momentum reflects the actual loss landscape geometry.

The warmup duration is typically expressed as a fixed number of steps (1000-2000 for small models, 4000-10000 for large models) rather than as a fraction of total training. The reason is that the required warmup duration scales with model depth and initial instability, not with the total number of training steps. A 175B parameter model requires roughly the same warmup duration regardless of whether you train it for 100B or 300B tokens.

Post-warmup learning rate decay also contributes to long-term stability. As the model approaches a good region of parameter space, reducing the learning rate prevents the optimizer from overshooting and allows fine-grained convergence. Cosine decay is preferred because it is smooth (no discontinuities that could cause sudden instability) and naturally brings the learning rate to near-zero at the end of training.

Adam Epsilon and Numerical StabilityLink Copied

The Adam optimizer has a small constant $\epsilon$ in its denominator to prevent division by zero:

\theta_t = \theta_{t-1} - \frac{\eta}{\sqrt{\hat{v}_t} + \epsilon} \cdot \hat{m}_t

where:

$\hat{v}_t$ : the bias-corrected second moment estimate (variance of recent gradients)
$\hat{m}_t$ : the bias-corrected first moment estimate (mean of recent gradients)
$\epsilon$ : a small constant for numerical stability (default $10^{-8}$ in most frameworks)
$\eta$ : the learning rate

The default $\epsilon = 10^{-8}$ works well for standard float32 training. But for mixed-precision training with float16, this value is often too small relative to the precision of float16 arithmetic. Float16 can represent values as small as approximately $6 \times 10^{-5}$ before underflow to zero. If $\hat{v}_t$ is small (which happens early in training or for parameters that receive weak gradient signals), then $\sqrt{\hat{v}_t}$ can be close to or below the float16 precision floor, making the $\epsilon$ in the denominator effectively invisible. The result is effective division by near-zero, which causes enormous parameter updates.

A common fix for instability in mixed-precision training is to increase $\epsilon$ to $10^{-6}$ or even $10^{-5}$ . This makes the effective step size slightly less adaptive for parameters with very small second moments, but significantly more numerically stable. Google's Adafactor optimizer, designed for extreme-scale models, uses an even larger effective epsilon to ensure stability.

In[15]:

Code

def run_training_with_epsilon(epsilon, n_steps=200, lr=1e-3, use_fp16=False):
    """Train a model with a specific Adam epsilon and return loss trajectory."""
    torch.manual_seed(42)
    model_eps = SimpleTransformerBlock(d_model=64, n_heads=4)
    dtype = torch.float16 if use_fp16 else torch.float32
    model_eps = model_eps.to(dtype=dtype)
    opt = optim.Adam(model_eps.parameters(), lr=lr, eps=epsilon)
    loss_history = []

    for step in range(n_steps):
        x = torch.randn(8, 16, 64, dtype=dtype)
        target = torch.randn(8, 16, 64, dtype=dtype)
        opt.zero_grad()
        try:
            out = model_eps(x)
            loss = criterion(out.float(), target.float())
            loss.backward()
            opt.step()
            loss_history.append(loss.item())
        except RuntimeError:
            loss_history.append(float("nan"))

    return loss_history


epsilon_values = [1e-8, 1e-7, 1e-6]
epsilon_results = {}
for eps in epsilon_values:
    epsilon_results[eps] = run_training_with_epsilon(eps, use_fp16=False)

Out[16]:

Visualization

Line plot showing training loss for three epsilon settings (1e-8, 1e-7, 1e-6) converging at similar rates over 200 steps. — Training loss curves for Adam with three different epsilon values over 200 steps in standard float32 mode. All three settings converge similarly under float32, showing that the choice of epsilon matters less in full precision. The critical distinction emerges in float16 mixed-precision training, where smaller epsilon values can cause numerical underflow in the second moment estimates, leading to instability that does not manifest in this float32 comparison.

Weight Decay and RegularizationLink Copied

Weight decay, as covered in the AdamW chapter, constrains parameter magnitudes by penalizing large weights. From a stability perspective, this has an important secondary effect: it prevents any individual weight from growing to extreme values that would dominate the gradient signal in subsequent steps. Bounded weights mean bounded activation magnitudes (given normalized inputs), which in turn means bounded gradients.

AdamW's decoupled weight decay is more effective for stability than the L2 regularization form of Adam. In standard Adam with L2 regularization, the regularization signal is mixed with the gradient signal and then adaptively scaled, which weakens its effect. When a parameter has a small second moment estimate (that is, it has received weak gradients recently), Adam applies a very large adaptive scaling factor. If L2 regularization is included in the gradient, this large factor also applies to the regularization term, which can cause the regularization to suddenly push parameters far from their current values. AdamW applies weight decay directly to parameters before the gradient update, avoiding this interaction entirely.

QK Normalization in AttentionLink Copied

A more recent stability technique, adopted in models like PaLM 2 and some variants of GPT-4, is query-key normalization in the attention mechanism. The attention computation involves dot products between query and key vectors:

\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\!\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d_k}}\right)\mathbf{V}

For very large models with large head dimensions $d_k$ , or for long sequences where many tokens contribute to the dot product sum, the pre-softmax logits $\mathbf{Q}\mathbf{K}^T / \sqrt{d_k}$ can still reach large magnitudes. When these logits are very large, the softmax becomes extremely peaked (a near-one-hot distribution), the gradient of the softmax with respect to the logits approaches zero everywhere, and training slows dramatically or becomes unstable.

QK normalization applies LayerNorm or L2 normalization to the query and key vectors before computing the dot product. This bounds the range of attention logits independent of sequence length or model scale, giving a structural guarantee against attention logit explosion. The technique adds almost no computational overhead but provides meaningful stability benefits for large models trained on long sequences.