Elamur Research⌕
☰
home / papers / 2023 / Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training

Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training

Paper · 2023 · Hong Liu, Zhiyuan Li, David Hall et al.

Summary

Pre-training large language models is very expensive, and Adam-style optimizers have dominated for years. Earlier second-order methods were often too costly per iteration to pay off. The authors propose Sophia (Second-order Clipped Stochastic Optimization). It preconditions a momentum-averaged gradient with a cheap, periodically refreshed estimate of the diagonal Hessian, then clips each coordinate of the resulting update. The clipping limits step size where the curvature estimate is unreliable, negative, or changing fast. This allows the Hessian to be re-estimated only every few iterations, keeping time and memory close to AdamW.

Sophia was tested by pre-training GPT-2 models on OpenWebText and a GPT NeoX model on the Pile, with sizes from 125M to 1.5B parameters. It reached AdamW's validation loss in roughly half the steps, and the savings carried over to total compute and wall-clock time. A simplified theoretical analysis shows its convergence bound does not depend on the loss's condition number. Since pre-training budgets are enormous, halving optimizer cost is a substantial practical gain.

Key points

  • There are two diagonal Hessian estimators: Hutchinson (via Hessian-vector products, the Sophia-H variant) and Gauss-Newton-Bartlett (gradients on labels sampled from the model, the Sophia-G variant). Both add only about 5% average overhead per step.
  • The Hessian estimate is refreshed every k = 10 steps. The update is the EMA of gradients divided by the EMA of the Hessian estimate, with element-wise clipping applied.
  • Sophia gives about a 2x speedup over AdamW in steps, compute, and wall-clock time on GPT-2 medium (355M), GPT-2 large (770M), and GPT NeoX 1.5B.
  • The advantage grows with model size. At 100K steps, a 540M model trained with Sophia matches a 770M model trained with AdamW, which needs about 40% more training and inference cost.
  • No special architecture or hardware is required. It runs in standard frameworks such as PyTorch and JAX.

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

 Institutions

🎓Stanford UniversityStanford · United States567 works · 335,189 citations

 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

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 4,145 works

Cite

Cite as Hong Liu, Zhiyuan Li, David Hall, Percy Liang and Tengyu Ma. 2023. Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training. arXiv:2305.14342.

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

@misc{liu2023sophia,
  title = {{Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training}},
  author = {Hong Liu, Zhiyuan Li, David Hall, Percy Liang and Tengyu Ma},
  year = {2023},
  doi = {10.48550/arxiv.2305.14342},
  eprint = {2305.14342},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2305.14342}
}

References

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

  1. Anil, R., Gupta, V., Koren, T., and Singer, Y. Memory efficient adaptive optimization. Advances in Neural Information Processing Systems, 32, 2019. find
  2. Anil, R., Gupta, V., Koren, T., Regan, K., and Singer, Y. Scalable second order optimization for deep learning. arXiv preprint arXiv:2002.09018, 2020. — [arXiv:2002.09018](https://arxiv.org/abs/2002.09018) find
  3. Ba, J., Grosse, R., and Martens, J. Distributed second-order optimization using kronecker-factored approximations. In International Conference on Learning Representations, 2017. find
  4. Balles, L. and Hennig, P. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In International Conference on Machine Learning, pp. 404–413. PMLR, 2018. find
  5. Bartlett, M. Approximate confidence intervals. Biometrika, 40(1/2):12–19, 1953. find
  6. Becker, S. and Le Cun, Y. Improving the convergence of back-propagation learning with. 1988. find
  7. Bernstein, J., Wang, Y.-X., Azizzadenesheli, K., and Anandkumar, A. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 560–569. PMLR, 2018. find
  8. Black, S., Biderman, S., Hallahan, E., Anthony, Q., Gao, L., Golding, L., He, H., Leahy, C., McDonell, K., Phang, J., et al. Gpt-neox-20b: An open-source autoregressive language model. arXiv preprint arXiv:2204.06745, 2022.
  9. Botev, A., Ritter, H., and Barber, D. Practical gauss-newton optimisation for deep learning. In International Conference on Machine Learning, pp. 557–565. PMLR, 2017. find
  10. Boyd, S. P. and Vandenberghe, L. Convex optimization. Cambridge university press, 2004. find
  11. Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax. ↗
  12. Braun, H. and Riedmiller, M. Rprop: a fast adaptive learning algorithm. In Proceedings of the International Symposium on Computer and Information Science VII, 1992. find
  13. Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901, 2020.
  14. Broyden, C. G. The convergence of a class of double-rank minimization algorithms 1. general considerations. IMA Journal of Applied Mathematics, 6(1):76–90, 1970. find
  15. Chapelle, O., Erhan, D., et al. Improved preconditioner for hessian free optimization. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning, volume 201. Citeseer, 2011. find
  16. Chen, P. Hessian matrix vs. gauss–newton hessian matrix. SIAM Journal on Numerical Analysis, 49(4): 1417–1435, 2011. find
  17. Chen, X., Liang, C., Huang, D., Real, E., Wang, K., Liu, Y., Pham, H., Dong, X., Luong, T., Hsieh, C.-J., et al. Symbolic discovery of optimization algorithms. arXiv preprint arXiv:2302.06675, 2023. find
  18. Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
  19. Conn, A. R., Gould, N., and Toint, P. L. Trust-region methods, siam. MPS, Philadelphia, 2000. find
  20. Crawshaw, M., Liu, M., Orabona, F., Zhang, W., and Zhuang, Z. Robustness to unbounded smoothness of generalized signsgd. arXiv preprint arXiv:2208.11195, 2022. — [arXiv:2208.11195](https://arxiv.org/abs/2208.11195) find
  21. Dennis Jr, J. E. and Schnabel, R. B. Numerical methods for unconstrained optimization and nonlinear equations. SIAM, 1996. find
  22. Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
  23. Dozat, T. Incorporating nesterov momentum into adam. 2016. find
  24. Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of Machine Learning Research, 12(Jul):2121–2159, 2011. find
  25. Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S., and Leahy, C. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. find
  26. Gargiani, M., Zanelli, A., Diehl, M., and Hutter, F. On the promise of the stochastic generalized gauss-newton method for training dnns. arXiv preprint arXiv:2006.02409, 2020. — [arXiv:2006.02409](https://arxiv.org/abs/2006.02409) find
  27. George, T., Laurent, C., Bouthillier, X., Ballas, N., and Vincent, P. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018. find
  28. Ghorbani, B., Krishnan, S., and Xiao, Y. An investigation into neural net optimization via hessian eigenvalue density. In International Conference on Machine Learning, pp. 2232–2241. PMLR, 2019. find
  29. Gokaslan, A. and Cohen, V. Openwebtext corpus, 2019. find
  30. Grosse, R. Neural Network Training Dynamics. 2022. find
  31. Grosse, R. and Martens, J. A kronecker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pp. 573–582. PMLR, 2016. find
  32. Gupta, V., Koren, T., and Singer, Y. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning, pp. 1842–1850. find
  33. He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770–778, 2016.
  34. Hinton, G., Srivastava, N., and Swersky, K. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. Cited on, 14(8):2, 2012. find
  35. Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556, 2022.
  36. Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics-Simulation and Computation, 18(3):1059–1076, 1989. find
  37. Izsak, P., Berchansky, M., and Levy, O. How to train bert with an academic budget. arXiv preprint arXiv:2104.07705, 2021. find
  38. Jahani, M., Rusakov, S., Shi, Z., Richtárik, P., Mahoney, M. W., and Takáč, M. Doubly adaptive scaled algorithm for machine learning using second-order information. arXiv preprint arXiv:2109.05198, 2021. — [arXiv:2109.05198](https://arxiv.org/abs/2109.05198) find
  39. Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
  40. Karamcheti, S., Orr, L., Bolton, J., Zhang, T., Goel, K., Narayan, A., Bommasani, R., Narayanan, D., Hashimoto, T., Jurafsky, D., Manning, C. D., Potts, C., Ré, C., and Liang, P. Mistral – a journey towards reproducible language model training. https://crfm.stanford.edu/2021/08/ 26/mistral.html, 202… ↗
  41. Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
  42. Kunstner, F., Hennig, P., and Balles, L. Limitations of the empirical fisher approximation for natural gradient descent. Advances in neural information processing systems, 32, 2019. find
  43. Kunstner, F., Chen, J., Lavington, J. W., and Schmidt, M. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. arXiv preprint arXiv:2304.13960, 2023. — [arXiv:2304.13960](https://arxiv.org/abs/2304.13960) find
  44. Liu, L., Liu, X., Gao, J., Chen, W., and Han, J. Understanding the difficulty of training transformers. arXiv preprint arXiv:2004.08249, 2020.
  45. Loshchilov, I. and Hutter, F. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983, 2016. find
  46. Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101, 2017. find
  47. Mai, V. V. and Johansson, M. Stability and convergence of stochastic gradient clipping: Beyond lipschitz continuity and smoothness. In International Conference on Machine Learning, pp. 7325–7335. PMLR, 2021. find
  48. Martens, J. New insights and perspectives on the natural gradient method. The Journal of Machine Learning Research, 21(1):5776–5851, 2020. find
  49. Martens, J. and Grosse, R. Optimizing neural networks with kronecker-factored approximate curvature. In International conference on machine learning, pp. 2408–2417. PMLR, 2015. find
  50. Martens, J., Ba, J., and Johnson, M. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations, 2018. find
  51. Martens, J. et al. Deep learning via hessian-free optimization. In ICML, volume 27, pp. 735–742, 2010. find
  52. Merity, S., Keskar, N. S., and Socher, R. Regularizing and optimizing lstm language models. arXiv preprint arXiv:1708.02182, 2017.
  53. Nesterov, Y. and Polyak, B. T. Cubic regularization of newton method and its global performance. Mathematical Programming, 108(1):177–205, 2006. find
  54. OpenAI. Gpt-4 technical report. arXiv, 2023.
  55. Ortega, J. M. and Rheinboldt, W. C. Iterative solution of nonlinear equations in several variables. SIAM, 2000. find
  56. Pascanu, R. and Bengio, Y. arXiv:1301.3584, 2013. — [arXiv:1301.3584](https://arxiv.org/abs/1301.3584) find
  57. Revisiting natural gradient for deep networks. arXiv preprint Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai,… — PyTorch: An Imperative Style, High-Performance Deep Learning Library
  58. Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019. find
  59. Rae, J. W., Borgeaud, S., Cai, T., Millican, K., Hoffmann, J., Song, F., Aslanides, J., Henderson, S., Ring, R., Young, S., et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11446, 2021. find
  60. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21:1–67, 2020.
  61. Reddi, S. J., Kale, S., and Kumar, S. On the convergence of adam and beyond. arXiv preprint arXiv:1904.09237, 2019.
  62. Roosta-Khorasani, F. and Ascher, U. Improved bounds on sample size for implicit matrix trace estimators. Foundations of Computational Mathematics, 15(5):1187–1212, 2015. find
  63. Sagun, L., Bottou, L., and LeCun, Y. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07476, 2016. — [arXiv:1611.07476](https://arxiv.org/abs/1611.07476) find
  64. Sankar, A. R., Khasbage, Y., Vigneswaran, R., and Balasubramanian, V. N. 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, pp. 9481–9488, 2021. find
  65. Schaul, T., Zhang, S., and LeCun, Y. No more pesky learning rates. In International conference on machine learning, pp. 343–351. PMLR, 2013. find
  66. Schraudolph, N. N. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 14(7):1723–1738, 2002. find
  67. Shazeer, N. and Stern, M. Adafactor: Adaptive learning rates with sublinear memory cost. In International Conference on Machine Learning, pp. 4596–4604. PMLR, 2018. find
  68. Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov, R. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1): 1929–1958, 2014.
  69. Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971, 2023.
  70. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I. Attention is all you need. arXiv preprint arXiv:1706.03762, 2017.
  71. Wang, A., Pruksachatkun, Y., Nangia, N., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. Superglue: A stickier benchmark for general-purpose language understanding systems. Advances in neural information processing systems, 32, 2019.
  72. Wei, C., Kakade, S., and Ma, T. The implicit and explicit regularization effects of dropout. arXiv preprint arXiv:2002.12915, 2020. — [arXiv:2002.12915](https://arxiv.org/abs/2002.12915) find
  73. Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Scao, T. L., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. Transformers: State-of-the-art natural… — Transformers: State-of-the-Art Natural Language Processing
  74. Yao, Z., Gholami, A., Keutzer, K., and Mahoney, M. W. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), pp. 581–590. find
  75. Yao, Z., Gholami, A., Shen, S., Mustafa, M., Keutzer, K., and Mahoney, M. Adahessian: An adaptive second order optimizer for machine learning. In proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 10665–10673, 2021. find
  76. You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song, X., Demmel, J., Keutzer, K., and Hsieh, C.-J. Large batch optimization for deep learning: Training bert in 76 minutes. arXiv preprint arXiv:1904.00962, 2019. find
  77. Zhang, J., He, T., Sra, S., and Jadbabaie, A. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint arXiv:1905.11881, 2019. — [arXiv:1905.11881](https://arxiv.org/abs/1905.11881) find
  78. Zhang, J., Karimireddy, S. P., Veit, A., Kim, S., Reddi, S., Kumar, S., and Sra, S. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383– 15393, 2020. find
  79. Zhang, L., Shi, S., and Li, B. Eva: Practical second-order optimization with kronecker-vectorized approximation. In The Eleventh International Conference on Learning Representations, 2022a. find
  80. Zhang, S., Roller, S., Goyal, N., Artetxe, M., Chen, M., Chen, S., Dewan, C., Diab, M., Li, X., Lin, X. V., et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01068, 2022b.
  81. Zhuang, J., Tang, T., Ding, Y., Tatikonda, S. C., Dvornek, N., Papademetris, X., and Duncan, J. Adabelief optimizer: Adapting stepsizes by the belief in observed gradients. Advances in neural information processing systems, 33:18795–18806, 2020. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/liu-2023-sophia-scalable-stochastic-second-order/