Building the Science of Scaling: Improving the Efficiency of Deep Learning Optimizers
A profound regime change in the field of optimization may be around the corner. For a decade, the Adam optimizer has overwhelmingly dominated as the way to train deep learning models. Recently, however, Adam is being threatened by the rise of different optimizers such as Muon. Though it hasn’t been clear how to properly configure these optimizers for consistently good results. Indeed, the behavior of deep learning models has often been treated as a guessing game, where practitioners sweep through countless configurations in hopes of stumbling upon the most efficient setup. Now, researchers working on a new science of scaling have shown that the optimal settings for these models actually follow predictable mathematical rules as they scale, providing a reliable blueprint for faster training. New work develops this science for modern optimizers, paving the way for a sea change in the way we train neural networks.
The findings are detailed in the paper “Hyperparameter Transfer Enables Consistent Gains of Matrix-Preconditioned Optimizers Across Scales.” The work was authored by Courant PhD students Shikai Qiu and Zixi (Charlie) Chen, CDS PhD student Hoang Phan, CDS Assistant Professor of Mathematics and Data Science Qi Lei, and CDS Professor of Computer Science and Data Science Andrew Gordon Wilson.
Training large language models demands immense computational power, prompting the field of machine learning to constantly seek more efficient optimizers. Recent matrix-preconditioned algorithms, such as Muon, SOAP, and Shampoo, demonstrated improvements over the industry-standard AdamW optimizer in small-scale settings. However, previous attempts to validate these new optimizers at scale yielded contradictory results, with at least one study (“Fantastic Pretraining Optimizers and Where to Find Them”) suggesting their advantages vanished as models grew larger.
The research team suspected that these diminishing returns were not a flaw in the optimizers themselves, but rather a failure to properly scale their hyperparameters.
“People have been really trying to empirically figure out more efficient ways of training neural networks,” Qiu said. “And then in parallel, there have been very exciting developments in the theoretical understanding of how to scale optimizers. But people haven’t done enough work to bridge the gap in theory and practice.”
Translating this theory into practical application presented its own set of challenges. Matrix-preconditioned optimizers speed up training by significantly increasing the numerical rank of the parameter updates, thereby amplifying the learning signal. The team discovered that this behavior broke the assumptions of the original µP theory for scaling prescriptions, causing the optimal learning rate to drift unpredictably.
“We found that in our experiments this discrepancy actually caused trouble,” Chen said. “The optimal learning rate shifts. To understand this behavior, we went back to the theory, and we realized we need to account for the numerical rank changing across widths as training progresses, breaking an assumption in the original theory that happens to hold well for conventional optimizers like AdamW.”
After revising their approach to account for these observations, the researchers tested their framework on models ranging from 190 million to 1.4 billion parameters. By applying their optimized scaling rules, they demonstrated that Muon, SOAP, and Shampoo consistently achieved a 1.4x speedup over well-tuned AdamW.
“I’m sure a hundred years from now, when people train their neural networks, they won’t just be sweeping all their hyperparameters,” Qiu said. “They will follow mathematical rules of how to choose them based on works like this.”
By Stephen Thomas
