Elamur Research⌕
☰
home / papers / 2024 / Scaling Exponents Across Parameterizations and Optimizers

Scaling Exponents Across Parameterizations and Optimizers

Paper · 2024 · Katie Everett, Lechao Xiao, Mitchell Wortsman et al.

Summary

Scaling a Transformer from small to large width requires choosing how the initialization, parameter multipliers and learning rates change with width. Earlier theory (Yang & Hu; Yang & Littwin) derived these exponents by assuming strong correlation, called "alignment," between parameter updates and the data. This paper tests that assumption. It extends the theory to arbitrary alignment levels and to Adafactor-style optimizers, and introduces an alignment ratio metric to measure alignment during training.

The authors trained tens of thousands of models. These covered standard, NTK, muP and Mean-Field parameterizations combined with SGD, Adam and Adafactor. The measured alignment changes over training and generally falls between the two extremes assumed in prior work. As a result, the best learning-rate exponents were often ones that earlier theory would have ruled out.

With per-layer learning-rate prescriptions and tuned constants, every parameterization supports hyperparameter transfer. A new per-layer recipe for standard parameterization outperforms muP. The work matters for large "train once" runs, where silent scaling mismatches are costly.

Key points

  • The sweep covered 3 optimizers, 4 parameterizations, several alignment assumptions, more than 12 learning rates and 14 model sizes, up to 26.8B parameters. The vocabulary was 32,000 and widths reached 16,384. The work appeared at ICML 2024 and was done at Google DeepMind and MIT.
  • Alignment depends mainly on the training step, the layer and the parameterization, and much less on the optimizer.
  • The best Adam setup used standard parameterization with an embedding learning rate of O(1) and hidden and readout learning rates of O(1/n). Constant factors tuned on small models carried over to the largest ones.
  • Adam's epsilon can cause gradient underflow at realistic model sizes, and Mean-Field parameterization is the most affected. The proposed Adam-atan2 is a scale-invariant variant that removes epsilon entirely.
  • The study covers width scaling only. Co-scaling depth, batch size and training horizon is left for future work.

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 2The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 works

Cite

Cite as Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A. Alemi, Roman Novak, Peter J. Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee and Jeffrey Pennington. 2024. Scaling Exponents Across Parameterizations and Optimizers. arXiv:2407.05872.

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

@misc{everett2024scaling,
  title = {{Scaling Exponents Across Parameterizations and Optimizers}},
  author = {Katie Everett, Lechao Xiao, Mitchell Wortsman, Alexander A. Alemi, Roman Novak, Peter J. Liu, Izzeddin Gur, Jascha Sohl-Dickstein, Leslie Pack Kaelbling, Jaehoon Lee and Jeffrey Pennington},
  year = {2024},
  doi = {10.48550/arxiv.2407.05872},
  eprint = {2407.05872},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2407.05872}
}

References

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

  1. Araújo, D., Oliveira, R. I., and Yukimura, D. A mean-field limit for certain deep neural networks. arXiv preprint arXiv:1906.00193, 2019. — [arXiv:1906.00193](https://arxiv.org/abs/1906.00193) find
  2. Choi, D., Shallue, C. J., Nado, Z., Lee, J., Maddison, C. J., and Dahl, G. E. On empirical comparisons of optimizers for deep learning. arXiv preprint arXiv:1910.05446, 2019. — [arXiv:1910.05446](https://arxiv.org/abs/1910.05446) find
  3. Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
  4. 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. Journal of Machine Learning Research, 24(240):1–113, 2023.
  5. Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C., Hemsley, R., Hennigan, T., Hessel, M., Hou, S., Kapturowski, S., Keck, T., Kemaev, I., King, M., Kunesch, M., Martens, L.,… find
  6. Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A. P., Caron, M., Geirhos, R., Alabdulmohsin, I., et al. Scaling vision transformers to 22 billion parameters. In International Conference on Machine Learning, pp. 7480–7512. PMLR, 2023. find
  7. Dey, N., Gosal, G., Khachane, H., Marshall, W., Pathria, R., Tom, M., Hestness, J., et al. Cerebras-gpt: Open compute-optimal language models trained on the cerebras wafer-scale cluster. arXiv preprint arXiv:2304.03208, 2023. — [arXiv:2304.03208](https://arxiv.org/abs/2304.03208) find
  8. Bi, X., Chen, D., Chen, G., Chen, S., Dai, D., Deng, C., Ding, H., Dong, K., Du, Q., Fu, Z., et al. Deepseek llm: Scaling open-source language models with longtermism. arXiv preprint arXiv:2401.02954, 2024. find
  9. Dinan, E., Yaida, S., and Zhang, S. Effective theory of transformers at initialization. arXiv preprint arXiv:2304.02034, 2023. — [arXiv:2304.02034](https://arxiv.org/abs/2304.02034) find
  10. Blake, C., Orr, D., and Luschi, C. Unit scaling: Out-of-the-box low-precision training. arXiv preprint arXiv:2303.11257, 2023. — [arXiv:2303.11257](https://arxiv.org/abs/2303.11257) find
  11. Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., et al. Glam: Efficient scaling of language models with mixture-of-experts. In International Conference on Machine Learning, pp. 5547–5569. PMLR, 2022.
  12. Bordelon, B. and Pehlevan, C. Self-consistent dynamical field theory of kernel evolution in wide neural networks. Advances in Neural Information Processing Systems, 35: 32240–32256, 2022. find
  13. Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011. find
  14. Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018. find
  15. Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1–39, 2022. find
  16. Kaplan, J. Notes on contemporary machine learning for physicists, 2019. find
  17. 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.
  18. Geiger, M., Spigler, S., Jacot, A., and Wyart, M. Disentangling feature and lazy training in deep neural networks. Journal of Statistical Mechanics: Theory and Experiment, 2020(11):113301, 2020. find
  19. Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
  20. Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 249–256. JMLR Workshop and Conference Proceedings, 2010. find
  21. Kudo, T. and Richardson, J. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. arXiv preprint arXiv:1808.06226, 2018.
  22. Golovin, D., Solnik, B., Moitra, S., Kochanski, G., Karro, J., and Sculley, D. Google vizier: A service for black-box optimization. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, Halifax, NS, Canada, August 13 - 17, 2017, pp. 1487–1495. find
  23. ACM, 2017. doi: 10. 1145/3097983.3098043. URL https://doi.org/ 10.1145/3097983.3098043. — [doi:10.1145/3097983.3098043](https://doi.org/10.1145/3097983.3098043) find
  24. Large, T., Liu, Y., Huh, M., Bahng, H., Isola, P., and Bernstein, J. Scalable optimization in the modular norm. arXiv preprint arXiv:2405.14813, 2024. — [arXiv:2405.14813](https://arxiv.org/abs/2405.14813) find
  25. He, K., Zhang, X., Ren, S., and Sun, J. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026–1034, 2015. find
  26. Lingle, L. A large-scale exploration of µ-transfer. arXiv preprint arXiv:2404.05728, 2024. — [arXiv:2404.05728](https://arxiv.org/abs/2404.05728) find
  27. Lee, J., Bahri, Y., Novak, R., Schoenholz, S. S., Pennington, J., and Sohl-Dickstein, J. Deep neural networks as gaussian processes. arXiv preprint arXiv:1711.00165, 2017. — [arXiv:1711.00165](https://arxiv.org/abs/1711.00165) find
  28. Liu, P. J., Novak, R., Lee, J., Wortsman, M., Xiao, L., Everett, K., Alemi, A. A., Kurzeja, M., Marcenac, P., Gur, I., Kornblith, S., Xu, K., Elsayed, G., Fischer, I., Pennington, J., Adlam, B., and Dickstein, J.-S. Nanodo: A minimal transformer decoder-only language model implementation in JAX., 20… find
  29. Heek, J., Levskaya, A., Oliver, A., Ritter, M., Rondepierre, B., Steiner, A., and van Zee, M. Flax: A neural network library and ecosystem for JAX, 2023. URL http://github.com/google/flax. ↗
  30. Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In International Conference on Learning Representations, 2018. find
  31. Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415, 2016. find
  32. 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
  33. Matthews, A. G. d. G., Rowland, M., Hron, J., Turner, R. E., and Ghahramani, Z. Gaussian process behaviour in wide deep neural networks. arXiv preprint arXiv:1804.11271, 2018. — [arXiv:1804.11271](https://arxiv.org/abs/1804.11271) find
  34. 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.
  35. McCandlish, S., Kaplan, J., Amodei, D., and Team, O. D. An empirical model of large-batch training. arXiv preprint arXiv:1812.06162, 2018. find
  36. Mei, S., Montanari, A., and Nguyen, P.-M. A mean field view of the landscape of two-layer neural networks. Proceedings of the National Academy of Sciences, 115(33): E7665–E7671, 2018. find
  37. Ishikawa, S. and Karakida, R. On the parameterization of second-order optimization effective towards the infinite width. arXiv preprint arXiv:2312.12226, 2023. — [arXiv:2312.12226](https://arxiv.org/abs/2312.12226) find
  38. Neal, R. M. Priors for infinite networks. Bayesian learning for neural networks, pp. 29–53, 1996. mesh/blob/master/mesh_tensorflow/ layers.py\#L48-L98. find
  39. Numpy. Atan2 function. URL https: Shazeer, N. and Stern, M. Adafactor: Adaptive learning //github.com/numpy/numpy/blob/ rates with sublinear memory cost. In International Con-5c7b7b69cc3a6ea0bbbae2976445a169484e0a85/ ference on Machine Learning, pp. 4596–4604. PMLR, numpy/_core/src/npymath/npy_math_… — [arXiv:2001.07301](https://arxiv.org/abs/2001.07301) find
  40. IEEE, 2020. 10533. PMLR, 2020. Rotskoff, G. and Vanden-Eijnden, E. Parameters as inter-Yaida, S. Meta-principled family of hyperparameter scaling acting particles: long time convergence and asymptotic strategies. arXiv preprint arXiv:2210.04909, 2022. error scaling of neural networks. Advances in ne… find
  41. BigScience Workshop, 2022. Yang, G. and Littwin, E. Tensor programs ivb: Adaptive optimization in the infinite-width limit. arXiv preprint Shallue, C. J., Lee, J., Antognini, J., Sohl-Dickstein, J., arXiv:2308.01814, 2023. Frostig, R., and Dahl, G. E. Measuring the effects of data parallelism on neu… — [arXiv:2308.01814](https://arxiv.org/abs/2308.01814) find
  42. Yang, G., Simon, J. B., and Bernstein, J. A spectral condition for feature learning. arXiv preprint arXiv:2310.17813, 2023a. find
  43. Yang, G., Yu, D., Zhu, C., and Hayou, S. Tensor programs vi: Feature learning in infinite-depth neural networks. arXiv preprint arXiv:2310.02244, 2023b. find
  44. Zhai, S., Likhomanenko, T., Littwin, E., Busbridge, D., Ramapuram, J., Zhang, Y., Gu, J., and Susskind, J. M. Stabilizing transformer training by preventing attention entropy collapse. In International Conference on Machine Learning, pp. 40770–40803. PMLR, 2023. find
  45. Zhai, X., Kolesnikov, A., Houlsby, N., and Beyer, L. Scaling vision transformers, 2021. https://arxiv.org/ abs/2106.04560. ↗
  46. Zhang, G., Li, L., Nado, Z., Martens, J., Sachdeva, S., Dahl, G., Shallue, C., and Grosse, R. B. Which algorithmic choices matter at which batch sizes? insights from a noisy quadratic model. Advances in neural information processing systems, 32, 2019. find
  47. Zoph, B., Bello, I., Kumar, S., Du, N., Huang, Y., Dean, J., Shazeer, N., and Fedus, W. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08906, 2022. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/everett-2024-scaling-exponents-across-parameterizations-optimizers/