TLDR: A systematic study of ten deep learning optimizers found that previously claimed 1.4x to 2x speedups over AdamW were often due to unequal hyperparameter tuning and limited evaluation. When rigorously benchmarked, the actual speedup of alternative optimizers is lower (max 1.4x) and diminishes with model size (to 1.1x for 1.2B models). Matrix-based optimizers generally outperform scalar-based ones, and the best optimizer choice can depend on the data-to-model ratio. The study emphasizes the critical need for thorough hyperparameter tuning and end-of-training evaluations.
In the rapidly evolving landscape of large language model pretraining, optimizers play a crucial role in determining training efficiency and cost. For a long time, AdamW has been the go-to optimizer, despite numerous claims that alternative methods could offer significant speedups, ranging from 1.4 to 2 times faster. However, a recent systematic study by researchers Kaiyue Wen, David Hall, Tengyu Ma, and Percy Liang from Stanford University suggests that these claims might be overstated due to methodological shortcomings in previous comparisons.
The research paper, titled “Fantastic Pretraining Optimizers and Where to Find Them,” delves into two primary issues that have obscured fair comparisons and hindered the practical adoption of newer optimizers: inconsistent hyperparameter tuning and limited or misleading evaluation setups. To address these, the Stanford team conducted a comprehensive study of ten different deep learning optimizers across four model scales, ranging from 0.1 billion to 1.2 billion parameters, and varying data-to-model ratios (1 to 8 times the Chinchilla optimum).
The Critical Role of Hyperparameter Tuning
One of the most significant findings is the absolute necessity of rigorous hyperparameter tuning. The study reveals that optimal settings for one optimizer are often suboptimal for another, making direct comparisons without proper tuning inherently unfair. For instance, simply tuning the learning rate for AdamW, a commonly used baseline, can yield up to a 2x speedup in pretraining. This highlights that many reported speedups for new optimizers might simply reflect a poorly tuned AdamW baseline rather than a true advantage of the new method. The researchers found that even conceptually similar optimizers, like Lion and AdamW, require vastly different optimal weight decay values (around 0.6 for Lion versus 0.1 for AdamW).
Actual Speedups Are More Modest and Diminish with Scale
Against a meticulously tuned AdamW baseline, the study found that the actual speedup of alternative optimizers is considerably lower than previously claimed, never exceeding 1.4 times. Furthermore, this speedup tends to decrease as model size increases. While optimizers like Muon and Soap showed a 1.3x speedup for smaller models (0.1 billion parameters), this advantage dwindled to merely 1.1x for larger 1.2 billion parameter models. This suggests that the efficiency gains from novel optimizers might not scale as effectively to the massive models used in frontier AI research.
Matrix-Based Optimizers Lead the Pack
The investigation also shed light on different classes of optimizers. Scalar-based optimizers, which update each parameter individually using scalar operations (e.g., AdamW, Lion, Mars), generally achieved similar optimization speeds to AdamW after proper tuning, with an average speedup of less than 1.2x. In contrast, matrix-based optimizers (e.g., Kron, Muon, Soap) leverage the inherent matrix structure of neural networks and precondition gradients using matrix multiplication. These optimizers consistently delivered approximately a 1.3x speedup over AdamW for models under 520 million parameters. This indicates a promising direction for optimizer design, though their speedup still diminishes with increasing model scale.
Data-to-Model Ratio Matters
Another intriguing discovery is that the optimal choice of optimizer can shift depending on the data-to-model ratio. An optimizer that performs best in a 1x Chinchilla regime (a standard data-to-model ratio) might be suboptimal when the data-to-model ratio increases. For example, while Muon was consistently the top performer in smaller Chinchilla ratio regimes, it was outperformed by Kron and Soap when the data-to-model ratio increased to 8x or even 16x. This suggests that different optimizers might be better suited for different training budgets and data availability scenarios.
Also Read:
- Unveiling the Silent Thought Processes of Large Language Models
- MLP-Offload Accelerates Large Language Model Training by Breaking the GPU Memory Wall
Avoiding Misleading Early-Stage Evaluations
The study also cautions against relying on early-stage loss curves for optimizer evaluation. During learning rate decay, the performance rankings of different optimizers can flip multiple times. Judging optimizers based on intermediate checkpoints can therefore lead to inaccurate conclusions compared to evaluating models at their target training budget, after the full training trajectory has been completed.
In conclusion, this rigorous benchmarking study underscores that while new optimizers offer some improvements over AdamW, their true gains are often more modest than widely reported. The key takeaways emphasize the critical importance of thorough hyperparameter tuning, evaluating across a range of model scales and data-to-model ratios, and assessing performance at the end of training. These insights are crucial for guiding future research in optimizer design, particularly for developing methods whose efficiency remains stable as models and datasets continue to grow. For more detailed information, you can read the full research paper here.
