Source-linked AI summary
Small-scale proxies for large-scale Transformer training instabilities
Mitchell Wortsman, Peter J. Liu, Lechao Xiao, Katie Everett, Alex Alemi, Ben Adlam, John D. Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, Simon Kornblith
TL;DR
Large-scale Transformer training instabilities are difficult to investigate because reproducing them requires substantial resources. This paper studies whether small Transformers can reproduce, mitigate, and predict such instabilities using learning-rate sensitivity and scaling behavior, finding that known instabilities appear at high learning rates and that several interventions reduce sensitivity. Its scope is limited to instabilities that cause slow divergence rather than loss spikes.
Problem
Large-Transformer training can exhibit instabilities that do not appear at smaller scales, while reproducing them requires substantial computational resources.
Method
The paper measures learning-rate versus loss behavior across scales, studies optimizer and model interventions, and examines activation and gradient norm scaling to predict instability.
Results
Previously reported attention-logit growth and output-logit divergence instabilities can be reproduced in small Transformers at high learning rates, and qk-layernorm and z-loss regularization mitigate them.
Takeaways & Limitations
Learning-rate sensitivity and model-characteristic scaling provide tools for studying Transformer training stability without access to large resource pools.
Takeaways & Limitations
The paper focuses on instabilities that lead to slow divergence, not loss spikes.
Abstract
from arXiv · showhide
Teams that have trained large Transformer-based models have reported training instabilities at large scale that did not appear when training with the same hyperparameters at smaller scales. Although the causes of such instabilities are of scientific interest, the amount of resources required to reproduce them has made investigation difficult. In this work, we seek ways to reproduce and study training stability and instability at smaller scales. First, we focus on two sources of training instability described in previous work: the growth of logits in attention layers (Dehghani et al., 2023) and divergence of the output logits from the log probabilities (Chowdhery et al., 2022). By measuring the relationship between learning rate and loss across scales, we show that these instabilities also appear in small models when training at high learning rates, and that mitigations previously employed at large scales are equally effective in this regime. This prompts us to investigate the extent to which other known optimizer and model interventions influence the sensitivity of the final loss to changes in the learning rate. To this end, we study methods such as warm-up, weight decay, and the $μ$Param (Yang et al., 2022), and combine techniques to train small models that achieve similar losses across orders of magnitude of learning rate variation. Finally, to conclude our exploration we study two cases where instabilities can be predicted before they emerge by examining the scaling behavior of model activation and gradient norms.
1 Introduction
The paper shows that training instabilities reported at large scale can be reproduced and studied in small Transformers using high learning rates, while learning-rate sensitivity and scaling trends reveal mitigations and emerging risks.
- Motivation: Large-Transformer training can become unstable, making stability important to understand as the resources required for large runs increase.The paper motivates small-scale study because large training runs can slow or destabilize learning and are increasingly resource-intensive.
- Approach: The paper introduces LR sensitivity, measuring deviation from optimal performance as learning rate varies across orders of magnitude.It uses the relationship between learning rate and loss across scales to summarize training sensitivity.
- Small-scale reproduction: Two previously reported instabilities—attention-logit growth and output-logit divergence from log probabilities—also appear in small Transformers at high learning rates.This reproduces the instabilities without requiring large-scale resources.
- Mitigations: Qk-layernorm and z-loss regularization reduce LR sensitivity and enable successful training across three orders of magnitude of learning-rate variation.These interventions had previously been used at scale and are also effective in the high-learning-rate small-model regime.
- Other interventions: Longer warm-up and independently scaled learning rate and weight decay reduce learning-rate sensitivity, while depth scaling increases it faster than width scaling.With qk-layernorm and z-loss regularization, these techniques generally affect sensitivity within the stable range more than the range itself.
- Prediction and broader implications: Scaling activation and gradient norms can predict instability before it emerges and indicate that default AdamW epsilon may make updates too small.The paper connects this behavior and attention-logit growth to parameter-norm growth, supporting small-scale investigation of stability.
2 Experimental methodology
The study uses small decoder-only Transformers and analyzes instability through learning-rate/loss relationships and scaling trends in activation and gradient norms. It summarizes learning-rate robustness with LR sensitivity, defined from validation-loss deviations across a broad learning-rate sweep.
- 2.2 LR vs. loss curves and learning rate sensitivity: Learning-rate versus loss curves reveal how the lowest unstable learning rate changes with model size and often show instability as a loss explosion at high learning rates.The analysis compares models of different sizes to study how instability emerges with scale.
- 2.2 LR vs. loss curves and learning rate sensitivity: LR sensitivity measures the expected deviation of final validation loss from the optimum while sweeping learning rates across three orders of magnitude.It is computed using the best loss over the range, with losses capped by the initialization loss when evaluating each learning rate.
- 2.2 LR vs. loss curves and learning rate sensitivity: The default LR-sensitivity sweep uses AdamW learning rates from 3e-4 to 3e-1, where LR is the maximum value in a warm-up plus cosine-decay schedule.The evaluated rates are 3e-4, 1e-3, 3e-3, 1e-2, 3e-2, 1e-1, and 3e-1.
- 2.3 Scaling trends for model characteristics: Scaling trends for gradient and activation norms are examined as a way to predict instabilities before they emerge.This complements prior scaling analyses focused primarily on relationships between model scale and loss.
3 Results
The study reproduces two previously reported Transformer instabilities in small models at high learning rates, then evaluates interventions and predicts instability from model characteristics. Learning-rate sensitivity reveals how warm-up, weight decay, scaling choices, and µParam affect stability and loss across scales.
- 3.1 Reproducing two known instabilities at small scale: Two instabilities previously described at scale—attention-logit growth and output-logit divergence—also appear in small Transformers at high learning rates.Mitigations used at scale are effective in the small-model regime.
- 3.1.1 Attention logit growth: qk-layernorm raises the maximum trainable learning rate and lowers LR sensitivity, allowing a 1.2B-parameter model to train at learning rate 0.3.Without qk-layernorm, divergence occurs at smaller learning rates as model size increases.
- 3.1.2 Output logit divergence: Z-loss resolves output-logit divergence in models without weight decay across scales, while weight decay also mitigates the instability in the larger models tested.The instability occurs toward the end of training when logits become very negative.
- 3.2 Measuring the effect of other known interventions: Longer warm-up reduces LR sensitivity and loss, with the clearest effect in larger models that are unstable at learning rate 3e-1 without long warm-up.The experiment fixes total training steps at 1e5 and uses qk-layernorm.
- 3.2 Optimizer and scaling interventions: Independent weight decay reduces LR sensitivity, while independently scaling depth increases LR sensitivity faster than width but yields lower validation loss at the largest tested scale.Joint scaling performs best at the largest scale and extrapolates more reliably.
- 3.2.4 µParam: µParam stabilizes the optimal learning rate but does not improve loss or reduce LR sensitivity in these experiments.The results indicate that µParam does not remove the need for qk-layernorm at high learning rates.
- 3.3 Predicting attention logit growth instability: Scaling behavior can predict attention-logit divergence before it emerges: a fit forecast divergence for a new 4.8B-parameter model at learning rate 1e-2.The fit also closely extrapolated the maximum attention-logit value.
4 Related work
Related work covers documented Transformer instabilities, optimizer dynamics, and alternative approaches to learning-rate sensitivity. It also situates this paper alongside work on loss spikes, parameter-free methods, and other parameterizations.
- Attention and loss instabilities: Attention logit growth can collapse attention entropy, while qk-layernorm was proposed to normalize queries and keys before computing logits.Related work reports the issue in large ViT models and describes its one-hot attention consequence.
- Attention and loss instabilities: Fast loss spikes differ from the slow divergences emphasized here because they quickly increase loss and often eventually recover.Edge-of-stability work attributes such behavior to evolving curvature, parameter oscillation, and adaptive optimizer state.
- Attention and loss instabilities: Adaptive edge-of-stability analyses relate periodic loss spikes to progressive sharpening and optimizer-specific preconditioned sharpness thresholds.For Adam with β1=0.9, the cited threshold is 38/LR, and reducing β2 can help the adaptive state adjust more quickly.
- Optimizer behavior: A gradient-norm investigation finds that norms decrease with scale and learning rate, making default AdamW ϵ too large and updates too small.The associated figure illustrates reduced updates when gradient RMS falls below the AdamW ϵ hyperparameter.
- Alternative optimization and parameterizations: Other research seeks to remove explicit learning-rate selection or develops parameterizations beyond µParam.One cited method sets each step size using distance from initialization divided by the root sum of historical gradient squares.
5 Conclusion
The paper argues that small Transformers can provide useful insights into training stability as the compute required for the largest models increases. It presents this as an opportunity for stability research without access to large resource pools.
- 5 Conclusion: Small Transformers can reveal useful insights into training stability, supporting research without access to large resource pools.The conclusion links this opportunity to the increasing compute required for training the largest models.
A Additional infrastructure details
The infrastructure uses Flax, Jax, and TPUs, with sharded model and optimizer states compiled using JIT. Data loading is deterministic and sequences are packed without padding.
- A Additional infrastructure details: NanoDO is built on Flax, Jax, and TPUs, and shards model and optimizer states before JIT compilation to enable larger model training.Orbax provides checkpointing, while Grain provides deterministic data loading.
- A Additional infrastructure details: Packed sequences omit padding by appending an end-of-sequence token followed by the beginning of a new sequence when needed.This applies when a sequence is shorter than the context-length hyperparameter.
B When is learning rate sensitivity a useful metric
The paper identifies cases where learning-rate sensitivity is no longer a useful metric and examines those scenarios to justify its use for the studied interventions.
- B When is learning rate sensitivity a useful metric: Learning-rate sensitivity can become uninformative when an intervention changes the meaning of learning rate.The paper directs readers to an appendix discussion of these cases.
Interventions which change the meaning of learning rate
Learning-rate sensitivity is not useful when an intervention changes what the learning rate means. The paper therefore excludes or treats cautiously such manipulations, while retaining cases where the transformation is consistent across the learning-rate curve.
- Interventions which change the meaning of learning rate: Taking the square root of the learning rate before optimization makes LR sensitivity comparisons uninformative.The paper identifies this as a clear example of changing the meaning of learning rate.
- Interventions which change the meaning of learning rate: The authors generally avoid interventions that meaningfully change the learning-rate meaning.
- Interventions which change the meaning of learning rate: qk-layernorm, z-loss, and warm-up preserve comparable LR vs. loss curves up to a critical learning rate, supporting LR-sensitivity comparisons.
- Interventions which change the meaning of learning rate: µParam remains analyzable because its per-layer learning-rate modification is linear and unchanged across points on the LR vs. loss curve.
- Interventions which change the meaning of learning rate: Scaling learning rate by parameter root mean square is the paper’s example where LR sensitivity is likely not useful, so it is not measured.
Shifting of the optimal LR
The LR-sensitivity definition does not account for shifts in the optimal learning rate when selecting the learning-rate range. The authors recommend shifting the three-order-of-magnitude range accordingly, although their main experiments did not require it.
- Shifting of the optimal LR: LR sensitivity does not account for the optimal learning rate shifting within the specified range [a, b].
- Shifting of the optimal LR: The recommended response is to shift the three-order-of-magnitude range [a, b] to match the optimal-learning-rate shift.The paper gives Section 3.2.2 as an example and notes that Figure 1 did not use a sufficiently large scale to require this adjustment.
LR sensitivity is invariant to loss
LR sensitivity is invariant to the absolute scale of the loss, so it can be zero even when models consistently perform randomly. The authors recommend interpreting it together with LR-versus-loss curves rather than optimizing it alone.
- LR sensitivity is invariant to loss: LR sensitivity is invariant to the scale of the loss.
- LR sensitivity is invariant to loss: If a network achieves random performance across learning rates, LR sensitivity can still be zero.
- LR sensitivity is invariant to loss: The paper recommends examining LR sensitivity alongside LR vs. loss curves instead of treating it as an isolated optimization target.
C Output norm growth
The paper relates output-norm growth to attention-logit and AdamW-epsilon instabilities through parameter-norm growth. It proposes that quadratic dependence of attention logits on parameter RMS and shrinking gradients from growing residual-stream RMS help explain these instabilities.
- C Output norm growth: Output-norm growth is studied in relation to attention-logit growth and AdamW-epsilon instabilities through parameter-norm growth.
- C Output norm growth: The RMS of Transformer-block outputs is mainly determined by learning rate in the reported empirical evidence.
- C Output norm growth: Attention logits can become large first because their dot-product computation depends quadratically on parameter RMS, unlike typical matrix-multiply features.
- C Output norm growth: Residual connections increase RMS with depth, so the output layer norm scales gradients by a shrinking reciprocal input RMS.
- C Output norm growth: As the RMS leaving the final Transformer block grows, the received gradient shrinks.
E Additional figures
The additional figures examine how model scale, learning rate, normalization, optimizer settings, and architectural choices relate to training loss, instability, and scaling predictions. They also document implementation details and robustness checks for the paper’s experiments.
- Activation behavior: Transformer block-output RMS is roughly consistent across scale but increases with learning rate and becomes larger at greater depth because of residual connections.The measurements average RMS over the final 500 training steps.
- Instabilities: Logit growth is associated with increasing query and key norms rather than increasing cosine similarity, and it also occurs without softmax attention.The pointwise attention variant replaces softmax with squared-relu and uses inverse-square-root sequence-length scaling.
- Scaling behavior: Jointly scaling width and depth lowers loss at the largest tested scale and improves scaling predictions from models below 1e8 parameters.Best loss is reported from a learning-rate sweep.
- Parameterization and normalization: µParam stabilizes the optimal learning rate, but qk-layernorm remains necessary; with qk-layernorm, µParam (full) trains without diverging at learning rate 1.At the best learning rate, µParam (full) provides no measurable improvement over µParam (simple) at the largest tested scale.