Elamur Research⌕
☰
home / papers / 2023 / A Theory on Adam Instability in Large-Scale Machine Learning

A Theory on Adam Instability in Large-Scale Machine Learning

Paper · 2023 · Igor Molybog, Peter Albert, Moya Chen et al.

Summary

Very large language models sometimes show sudden loss spikes during training, and the cause has not been explained. Chowdhery et al. reported this on TPUs. This Meta AI paper reproduces the effect on GPUs with a different codebase, different data and some architecture changes: ReLU instead of SwiGLU, and learned absolute positional embeddings instead of RoPE. The authors trained decoder-only transformers of 7b, 30b, 65b and 546b parameters with AdamW.

They argue that the spikes come from Adam itself rather than from bad data. Their claim is that gradient estimates in early layers can become correlated over time, which breaks a key assumption behind Adam. In that state, the per-coordinate update ratios take on a bimodal distribution. The update vector then grows large while pointing almost nowhere along the true descent direction, and training diverges. This regime is more likely in deep models trained with big batches, which is exactly the frontier LLM setting. The paper concludes with mitigation options, noting that no single fix works for every setup.

Key points

  • The four runs grow in depth and batch size, from 32 layers with batch 2048 (7b) to 108 layers with batch 65,536 (546b). All used β1=0.9, β2=0.95 and ε=10⁻⁸.
  • Instability tracked scale:
    • The 546b run had prominent spikes, and most of them recovered within roughly a dozen iterations without batch skipping.
    • The 65b run was moderately unstable.
    • The 7b and 30b runs never diverged.
  • Results at 65b were cross-checked with a separate codebase and a LLaMa-like dataset, with no significant differences.
  • Proposed remedies include:
    • skipping batches;
    • shrinking the batch size, although this slows data-parallel training;
    • adjusting the averaging constants;
    • lowering ε or resetting the optimizer state when the Hartigan dip statistic detects that the update-ratio distribution has become bimodal.
  • Unlike earlier Adam divergence studies such as AMSGrad and Yogi, which address slow drifting divergence, this work targets abrupt spike-type divergence.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 works

Cite

Cite as Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, Binh Tang, Diana Liskovich, Puxin Xu, Yuchen Zhang, Melanie Kambadur, Stephen Roller and Susan Zhang. 2023. A Theory on Adam Instability in Large-Scale Machine Learning. arXiv:2304.09871.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2304.09871 · arXiv 2304.09871

@misc{molybog2023theory,
  title = {{A Theory on Adam Instability in Large-Scale Machine Learning}},
  author = {Igor Molybog, Peter Albert, Moya Chen, Zachary DeVito, David Esiobu, Naman Goyal, Punit Singh Koura, Sharan Narang, Andrew Poulton, Ruan Silva, Binh Tang, Diana Liskovich, Puxin Xu, Yuchen Zhang, Melanie Kambadur, Stephen Roller and Susan Zhang},
  year = {2023},
  doi = {10.48550/arxiv.2304.09871},
  eprint = {2304.09871},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2304.09871}
}

References

27 references, 9 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Patrick Billingsley. Probability and measure. John Wiley & Sons, 2008. find
  2. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
  3. Xiangyi Chen, Sijia Liu, Ruoyu Sun, and Mingyi Hong. On the convergence of a class of adam-type algorithms for non-convex optimization. arXiv preprint arXiv:1808.02941, 2018. — [arXiv:1808.02941](https://arxiv.org/abs/1808.02941) find
  4. Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
  5. John A Hartigan and Pamela M Hartigan. The dip test of unimodality. The annals of Statistics, pages 70–84, 1985. find
  6. Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies, 2001. find
  7. Jack Kiefer and Jacob Wolfowitz. Stochastic estimation of the maximum of a regression function. The Annals of Mathematical Statistics, pages 462–466, 1952. find
  8. Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
  9. Zhenyu Liao and Michael W Mahoney. Hessian eigenspectra of more realistic nonlinear models. Advances in Neural Information Processing Systems, 34:20104–20117, 2021. find
  10. Y Liu. Statistical behavior of the eigenvalues of random matrices, 2000. find
  11. Ilya Loshchilov and Frank Hutter. arXiv:1711.05101, 2017. find
  12. Decoupled weight decay regularization. arXiv preprint Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. Understanding the exploding gradient problem. CoRR, abs/1211.5063, 2(417):1, 2012. find
  13. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. find
  14. Sashank J Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv preprint arXiv:1904.09237, 2019.
  15. Adepu Ravi Sankar, Yash Khasbage, Rahul Vigneswaran, and Vineeth N Balasubramanian. A deeper look at the hessian eigenspectrum of deep neural networks and its applications to regularization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 9481–9488, 2021. find
  16. Noam Shazeer. Glu variants improve transformer. CoRR, abs/2002.05202, 2020. URL https://arxiv.org/abs/2002.05202. ↗
  17. Noam Shazeer and Mitchell Stern. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pages 4596–4604. PMLR, 2018. find
  18. Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model. arXiv preprint arXiv:22…
  19. Jianlin Su, Yu Lu, Shengfeng Pan, Bo Wen, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. CoRR, abs/2104.09864, 2021. URL https://arxiv.org/abs/2104. 09864. ↗
  20. Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
  21. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
  22. Ruiqi Wang and Diego Klabjan. Divergence results and convergence of a variance reduced version of adam. arXiv preprint arXiv:2210.05607, 2022. — [arXiv:2210.05607](https://arxiv.org/abs/2210.05607) find
  23. Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pages 581–590. find
  24. Manzil Zaheer, Sashank Reddi, Devendra Sachan, Satyen Kale, and Sanjiv Kumar. Adaptive methods for nonconvex optimization. Advances in neural information processing systems, 31, 2018. find
  25. Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022.
  26. Zhiming Zhou, Qingru Zhang, Guansong Lu, Hongwei Wang, Weinan Zhang, and Yong Yu. Adashift: Decorrelation and convergence of adaptive learning rate methods. arXiv preprint arXiv:1810.00143, 2018. — [arXiv:1810.00143](https://arxiv.org/abs/1810.00143) find
  27. Fangyu Zou, Li Shen, Zequn Jie, Weizhong Zhang, and Wei Liu. A sufficient condition for convergences of adam and rmsprop. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 11127–11135, 2019. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/molybog-2023-theory-adam-instability-large-scale/