How to Scale Second-Order Optimization
Shikai Qiu, Charlie Chen, Viet Hoang Phan, Qi Lei, Andrew Wilson
In deep learning, the heavy computational overhead of second-order optimizers has traditionally limited their practical adoption in favor of simpler first-order methods like SGD and Adam. Recent empirical successes, however, suggest that some second-order methods may unlock significant computational efficiency in training large-scale machine learning models. In this work, we systematically investigate several key open questions. Do second-order optimizers outperform first-order methods under a fixed compute budget? How can we optimally scale the hyperparameters as we scale up the model? Do efficiency gains in small-scale experiments to transfer to larger-scale systems, and if so, how does second-order optimization affect nerual scaling laws?  Through theoretical and empirical analysis, we derive scaling rules for hyperparameters such as learning rate as we jointly scale up width, depth, and batch size for a wide range of optimizers including Shampoo, SOAP, and Muon. We find that, under the correct scaling, second-order optimizers consistently outperform Adam across scales in the compute-optimal setting for training transformers. Furthermore, we show second-order optimizers alter the Chinchilla scaling law, requiring up to $2\times$ less data per parameter for compute-optimality.