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