Small-scale proxies for large-scale Transformer training instabilities
Authors: name: Mitchell Wortsman, name: Peter J. Liu, name: Lechao Xiao, name: Katie Everett, name: Alex Alemi, name: Ben Adlam, name: John D. Co-Reyes, name: Izzeddin Gur, name: Abhishek Kumar, name: Roman Novak, name: Jeffrey Pennington, name: Jascha Sohl-dickstein, name: Kelvin Xu, name: Jaehoon Lee, name: Justin Gilmer, name: Simon Kornblith

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 muParam (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.