[Paper Review] Small-scale proxies for large-scale Transformer training instabilities

Mitchell Wortsman, Peter J. Liu+14 more|arXiv (Cornell University)|Sep 25, 2023

Explainable Artificial Intelligence (XAI)4 citations

TL;DR

This paper introduces small-scale proxy experiments to study training instabilities in large-scale Transformers, demonstrating that key instabilities—logit growth in attention layers and divergence from log probabilities—can be reproduced and mitigated in small models using techniques like qk-layernorm and z-loss. The authors show that learning rate sensitivity is a reliable diagnostic, and that scaling behavior of activation and gradient norms can predict instability before it emerges.

ABSTRACT

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.

Motivation & Objective

• Investigate training instabilities in large-scale Transformers that do not appear at smaller scales, despite identical hyperparameters.
• Develop small-scale proxy models to study these instabilities without requiring large compute resources.
• Identify and validate effective interventions—such as qk-layernorm and z-loss—that stabilize training across learning rate variations.
• Analyze how optimizer and model modifications (e.g., warm-up, weight decay, μ-Param) affect learning rate sensitivity in small models.
• Explore whether scaling trends in activation and gradient norms can predict instability before it occurs during training.

Proposed method

• Train small decoder-only Transformer models with GPT-2-style architecture using Flax and AdamW optimizer with default hyperparameters.
• Systematically vary the learning rate across three orders of magnitude and measure final evaluation loss to compute learning rate (LR) sensitivity as a summary statistic.
• Introduce and evaluate qk-layernorm and z-loss regularization to reduce LR sensitivity and stabilize training across wide LR ranges.
• Study the impact of warm-up duration, weight decay, and μ-Param on LR sensitivity and training stability in small models.
• Analyze scaling behavior of activation norms and gradient norms across model scale and learning rate to detect early signs of instability.
• Use the observation that default AdamW epsilon is too large for large-scale training, leading to overly small updates, to link to instability mechanisms.

Experimental results

Research questions

• RQ1 Can training instabilities observed in large-scale Transformers—such as attention logit growth and output logit divergence—be reproduced in small-scale models?
• RQ2 To what extent do known large-scale mitigation techniques (e.g., qk-layernorm, z-loss) remain effective in small-scale proxy settings?
• RQ3 How do standard optimizer and model interventions (e.g., warm-up, weight decay, μ-Param) affect learning rate sensitivity in small models?
• RQ4 Can scaling trends in activation and gradient norms predict training instability before it manifests?
• RQ5 What role does the AdamW epsilon hyperparameter play in instability, particularly in relation to gradient norm scaling?

Key findings

• Training instabilities such as attention logit growth and output logit divergence are reproducible in small models when trained at high learning rates.
• qk-layernorm and z-loss regularization reduce learning rate sensitivity, enabling stable training across three orders of magnitude of learning rate variation in small models.
• Scaling depth increases learning rate sensitivity faster than scaling width, indicating deeper models are more sensitive to learning rate choice.
• Gradient norms decrease with both model scale and learning rate, indicating that the default AdamW epsilon (1e-8) is too large, leading to overly small parameter updates.
• The interaction between progressive sharpening and adaptive optimizer dynamics (e.g., Adam with β₂) explains periodic loss spikes, and reducing β₂ can prevent them by allowing faster adaptation of the preconditioner.
• Scaling behavior of activation and gradient norms can predict instability before it emerges, offering a diagnostic tool for training stability without large-scale runs.