Small-scale proxies for large-scale Transformer training instabilities
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.
65125ab90cd5dc7c2ca224a0
Mitchell Wortsman
admin_assigned
2023-09-26T12:57:52.466Z
638e42f96251c8bd7ab84089
/avatars/cb7a7f902846bd9710e2c48a50ce3527.svg
Mitchell Wortsman
user
mitchellw
65125ab90cd5dc7c2ca224a1
Peter J. Liu
admin_assigned
2023-09-26T12:57:59.926Z
62ded294dcdc9043efaa4fc3
/avatars/1294078b8cd5c1901e4021253d877357.svg
Peter J Liu
user
peterjliu
65125ab90cd5dc7c2ca224a2
Lechao Xiao
65125ab90cd5dc7c2ca224a3
Katie Everett
65125ab90cd5dc7c2ca224a4
Alex Alemi
admin_assigned
2023-09-26T12:58:30.903Z
6303ff84fc783bfc74441afc
/avatars/bc295e0d94c9c8be0c6a8264587ad4a5.svg
Alex Alemi
user
knowflake
65125ab90cd5dc7c2ca224a5
Ben Adlam
65125ab90cd5dc7c2ca224a6
John D. Co-Reyes
65125ab90cd5dc7c2ca224a7
Izzeddin Gur
admin_assigned
2023-09-26T12:58:58.627Z
64c0ccc19832a455bc2a2069
/avatars/9d4fdf13cd62c6c258b5f50d67c8a1b1.svg
IZZEDDIN GUR
user
izzeddingur
65125ab90cd5dc7c2ca224a8
Abhishek Kumar
65125ab90cd5dc7c2ca224a9
Roman Novak
admin_assigned
2023-09-26T12:59:31.618Z
63e2769bf0740bec2bfff1a1
/avatars/e2841a836319d43035bcddfbfe3344d0.svg
Roman Novak
user
kulpims
65125ab90cd5dc7c2ca224aa
Jeffrey Pennington
65125ab90cd5dc7c2ca224ab
Jascha Sohl-dickstein
admin_assigned
2023-09-26T12:59:46.787Z
63d209ae8eaa4831006b3b0d
/avatars/51cc0f6dfab2233bddd73bfcd0c8aee3.svg
Jascha Sohl-Dickstein
user
jaschasd
65125ab90cd5dc7c2ca224ac
Kelvin Xu
65125ab90cd5dc7c2ca224ad
Jaehoon Lee
65125ab90cd5dc7c2ca224ae
Justin Gilmer
65125ab90cd5dc7c2ca224af
Simon Kornblith
admin_assigned
2023-09-26T13:01:08.745Z
6443ffc2af034cdfd69c5c00
/avatars/0c12cf9baaf3b24a02ffb42238cdc6df.svg
Simon Kornblith
user
simonsterk
Small-scale proxies for large-scale Transformer training instabilities
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.
Training instabilities in large Transformer models are reproducible in smaller models at high learning rates, with similar mitigations effective; scaling behaviors of activations and gradients predict instabilities.
Transformer-based models logits attention layers output logits learning rate loss warm-up weight decay μParam model activation gradient norms