PaperlayerView exact revision on arXiv

Source-linked AI summary

Fantastic Pretraining Optimizers and Where to Find Them

Kaiyue Wen, David Hall, Tengyu Ma, Percy Liang

arXiv:2509.02046v2cs.LGcs.AIstat.ML

TL;DR

The paper addresses whether reported optimizer speedups over AdamW survive fair tuning and informative evaluation. It benchmarks modern optimizers across model scales and data-to-model ratios, finding that gains are smaller and decline with scale, while intermediate checkpoints can mislead rankings.

  • Problem

    Optimizer comparisons often use unequal hyperparameter tuning and limited or potentially misleading evaluation settings.

  • Method

    The paper benchmarks eleven optimizers with rigorous hyperparameter sweeps across model scales, Chinchilla ratios, and end-of-training evaluations.

  • Results

    Alternative optimizers achieve at most 1.4× speedup over well-tuned AdamW, with gains declining to around 1.1× for 1.2B-parameter models.

  • Takeaways & Limitations

    Matrix-based optimizers are fastest for smaller models, but optimizer rankings depend on data-to-model ratio and can change during learning-rate decay.

Abstract

from arXiv · show

AdamW has long been the dominant optimizer in language model pretraining, despite numerous claims that alternative optimizers offer 1.4 to 2x speedup. We posit that two methodological shortcomings have obscured fair comparisons and hindered practical adoption: (i) unequal hyperparameter tuning and (ii) limited or misleading evaluation setups. To address these two issues, we conduct a systematic study of ten deep learning optimizers across four model scales (0.1B-1.2B parameters) and data-to-model ratios (1-8x the Chinchilla optimum). We find that fair and informative comparisons require rigorous hyperparameter tuning and evaluations across a range of model scales and data-to-model ratios, performed at the end of training. First, optimal hyperparameters for one optimizer may be suboptimal for another, making blind hyperparameter transfer unfair. Second, the actual speedup of many proposed optimizers over well-tuned baselines is lower than claimed and decreases with model size to only 1.1x for 1.2B parameter models. Thirdly, comparing intermediate checkpoints before reaching the target training budgets can be misleading, as rankings between two optimizers can flip during training due to learning rate decay. Through our thorough investigation, we find that all the fastest optimizers such as Muon and Soap, use matrices as preconditioners -- multiplying gradients with matrices rather than entry-wise scalars. However, the speedup of matrix-based optimizers is inversely proportional to model scale, decreasing from 1.4x over AdamW for 0.1B parameter models to merely 1.1x for 1.2B parameter models.

1 Introduction

The paper argues that optimizer comparisons require equal hyperparameter tuning and evaluations across model scales, data-to-model ratios, and completed training budgets. Under these conditions, reported speedups shrink with scale, while matrix-based optimizers remain strongest for smaller models.

  • Motivation: Pretraining optimizer studies can be unfair when baselines are under-tuned or experiments remain limited to smaller-scale settings.The paper identifies unequal hyperparameter tuning and narrow evaluation regimes as two problems undermining confidence in new methods.
  • Hyperparameter tuning: Up to 2× speedup is achievable by tuning only the learning rate of the GPT-3 AdamW recipe for a 100M model.This result demonstrates that baseline tuning can substantially alter apparent optimizer comparisons.
  • Hyperparameter tuning: Lion’s optimal weight decay is approximately 0.6 versus approximately 0.1 for AdamW, so transferring hyperparameters can produce unfair comparisons.Conceptually similar optimizers may nevertheless require different optimal settings.
  • Scaling results: Against well-tuned AdamW, alternative optimizers do not exceed 1.4× speedup, falling from 1.3× at 0.1B parameters to around 1.1× at 1.2B parameters.The paper evaluates optimizer behavior across model scales and data-to-model ratios rather than relying only on small-model results.
  • Evaluation methodology: Intermediate checkpoints can mis-rank optimizers because their loss curves may cross multiple times during learning-rate decay.The paper therefore emphasizes comparisons at the target training budget.
  • Optimizer design: Matrix-based optimizers achieve approximately 1.3× speedup over AdamW for models under 520M parameters, whereas scalar-based optimizers average less than 1.2×.Matrix-based methods precondition gradients through matrix multiplication, while scalar-based methods update parameters individually.
  • Scaling results: Muon is strongest in smaller Chinchilla-ratio regimes, but Kron and Soap outperform it when the data-to-model ratio reaches 8× or more.The preferred optimizer changes with the training regime.

2 Related Works

Related work spans optimizer design, pretraining-specific acceleration, and methodological re-evaluation. Prior re-evaluations similarly emphasize problem-specific optimizer choices and rigorous tuning, while concurrent work reports differing matrix-optimizer rankings under different experimental settings.

  • Optimizers for Deep Learning: Deep-learning optimizer research evolved from classical optimization and adaptive learning-rate methods to Adam, AdamW, and newer variants.Examples include Nesterov Adam, decoupled weight decay, convergence-focused methods, and Lion.
  • Optimization for Pretraining: Pretraining-specific optimizers have claimed 1.4× to 2× speedups over AdamW, motivating systematic comparison of these methods.The paper studies optimizers including Sophia, Soap, Muon, MARS, Cautious AdamW, FOCUS, and SWAN.
  • Optimization for Pretraining: Prior work uses tensor programs and power-law fits to extrapolate hyperparameters across model scales, an approach incorporated by this paper.These methods address the cost of selecting hyperparameters for large-model pretraining.
  • Re-evaluation Methodology: Earlier optimizer re-evaluation found that the best optimizer is problem-specific and that unequal hyperparameter tuning can explain most claimed speedups.This paper extends that re-evaluation perspective to LLM pretraining and modern optimizers.
  • Re-evaluation Methodology: Related pretraining studies also observed that optimizer rankings can flip during learning-rate decay and that properly tuned AdamW remains competitive.These observations align with the paper’s emphasis on end-of-training evaluation.
  • Comparison with concurrent work: Concurrent work agrees that non-zero weight decay and learning-rate decay matter, but reports different relative performance for matrix-level optimizers.The discrepancy is associated with differences in batch size and hyperparameter-tuning methodology.

3 Methodology

The study evaluates eleven optimizers with a multi-phase tuning framework across model scales and data-to-model ratios. It emphasizes fair comparisons by tuning optimizer-specific hyperparameters and tracking final validation loss and downstream benchmarks.

  • Experimental design: Experiments use Llama-derived models from approximately 130M to 1.2B parameters, with fixed 4,096-token sequences and 32 transformer layers.
  • Evaluation: The primary metric is final language-modeling loss on C4-EN, supplemented by downstream accuracy and bits-per-byte across ten benchmarks.
  • Experimental design: Eleven optimizers are compared, including AdamW and Lion baselines alongside recently proposed methods grouped into five general classes.
  • Hyperparameter tuning: Phase I exhaustively sweeps each optimizer’s hyperparameters across six scale and Chinchilla-ratio settings using coordinate descent until no significant validation-loss improvement remains.The sweep includes learning rate, weight decay, warmup steps, momentum coefficients, epsilon, gradient clipping, and batch size for AdamW.
  • Hyperparameter tuning: Phase II retains scaling-sensitive hyperparameters and obtains near-optimal configurations for twelve model-size and data-budget settings.The approximate-optimal set includes configurations within Δ2 = 6.4e-3 of the best observed regime loss.
  • Empirical findings: Matrix-based optimizers outperform scalar-based optimizers, but speedups over AdamW remain at or below 1.5× and can fall below 1.2× at larger scale.Muon performs best at 1–4× Chinchilla ratio, while Soap and Kron overtake it at higher ratios; Phase III reports under 1.2× speedup up to 1.2B parameters.

4 Empirical Findings

Across model scales and data budgets, carefully tuned matrix-based optimizers generally outperform scalar-based methods, but their speedups shrink as models grow. Final evaluations also reveal optimizer rankings that early checkpoints, transferred hyperparameters, or narrow settings can misrepresent.

  • Small- and medium-scale models: 1.4× is the highest observed speedup over the AdamW baseline, below several prior claims of 2× acceleration.This result comes from benchmarks on 0.1B–0.5B-parameter models across Chinchilla ratios of 1 to 8.
  • Small- and medium-scale models: Matrix-based optimizers consistently outperform scalar-based methods, with their advantage increasing at larger data budgets but decreasing with model size.Muon performs best at 1× Chinchilla, while Soap and Kron become strongest at 8× Chinchilla.
  • Scaling to larger models: 1.3× speedups for 0.1B models diminish to roughly 1.1× for 1.2B models at 8× Chinchilla.At 1.2B parameters, NAdamW, Muon, and Soap still speed up over AdamW but no longer improve downstream benchmarks.
  • Optimizer families: Variance-reduced Adam variants consistently surpass vanilla AdamW, while Lion and Adam-mini remain within a 5% slowdown of AdamW.Muon combines matrix-based updates with Nesterov momentum, whereas memory-efficient variants can sometimes outperform AdamW.
  • High data-to-model ratios: At 16× Chinchilla, Soap outperforms Muon on 300M models, indicating that the preferred optimizer changes as the data-to-model ratio increases.The authors conjecture that Soap and Kron’s second-order momentum becomes more effective at higher data budgets.
  • Evaluation methodology: Transferred hyperparameters and early checkpoints can reverse optimizer rankings, so comparisons require optimizer-specific tuning and end-of-training evaluation.Small deviations from ideal hyperparameters can substantially degrade validation loss, while early rapid descent may be followed by plateauing.

5 Conclusion and Limitations

The study finds that optimizer comparisons require careful tuning and end-of-training evaluation, while reported gains over AdamW are smaller and decline with model scale. Its main scope limitation is that models larger than 1.2B parameters were not evaluated.

  • Fair hyperparameter sweeps eliminate many apparent optimizer advantages over AdamW.
  • Optimizer rankings can change across the full training trajectory, making early or inconsistent evaluations misleading.
  • 1.1× is the reported speedup of the best-performing alternatives at 1.2B parameters.
  • The benchmark does not scale beyond 1.2B parameters, leaving frontier-scale persistence of diminishing speedups untested.

B Omitted Experiments

The omitted scaling-law analysis compares Muon and AdamW across parameter and data scales. Its extrapolation predicts that Muon may lose its loss advantage at larger scale in the 1× Chinchilla regime.

  • Scaling laws were fitted for Muon, NAdamW, and AdamW using the 1.2B runs to select grid-aligned hyperparameters.
  • Muon’s fitted scaling law predicts higher loss than AdamW at 7B parameters in the 1× Chinchilla regime.

B.2 Sophia Experiments

The Sophia experiments compare Sophia with AdamW across smaller compute, model-size, and data-size settings. Sophia underperforms in smaller compute regimes but eventually slightly outperforms AdamW as model or data size increases.

  • Sophia tends to underperform AdamW in smaller compute regimes.
  • Sophia eventually slightly outperforms AdamW when either model size or data size increases.
  • The 1× Chinchilla and 130M-model experiments plot Sophia and AdamW loss curves.

B.3 High Data-to-model Ratio

The high-data-to-model-ratio case study evaluates a 130M model at 16× Chinchilla. In this overtraining setting, SOAP outperforms Muon.

  • SOAP outperforms Muon in the 130M, 16× Chinchilla overtraining setting.

B.4 Evaluation Performance

This section lists evaluation-performance tables for multiple optimizers across model sizes from 130M to 520M parameters, with final validation loss used in the reported results.

  • The tables report evaluation performance separately for Mars, Muon, Lion, NAdamW, Kron, Scion, SOAP, AdamW, and Adam-Mini.
  • The reported model sizes include 130M, 300M, and 520M parameters.
  • The loss presented for the Phase I sweeps is the final loss on the C4/EN validation set.

C.1 Sweeping Results for AdamW

This section reports Phase I hyperparameter-ablation results for the evaluated optimizers, organized across model sizes and Chinchilla data ratios.

  • The Phase I results cover hyperparameter-ablation tables for multiple optimizers, including AdamW, Cautious, Lion, Mars, NAdamW, Adam-Mini, Kron, Soap, Muon, Scion, and Sophia.
  • The reported ablations span 130M, 300M, and 520M model sizes and Chinchilla ratios from 1× to 8× where listed.

D.1 Sweeping Results for AdamW

This section reports Phase II hyperparameter-ablation results across optimizers, model sizes, and Chinchilla data ratios using the same presentation format as Phase I.

  • The Phase II results include hyperparameter-ablation tables for AdamW, Cautious, Lion, Mars, NAdamW, Adam-Mini, Muon, Kron, Soap, and Scion.
  • The tables cover 300M and 520M models at Chinchilla ratios from 2× to 8×.

E.1 Sweeping Results for AdamW

This section reports hyperparameter-ablation results for 1.2B-parameter models across AdamW, NAdamW, Soap, and Muon.

  • The tables cover AdamW, NAdamW, Soap, and Muon hyperparameter ablations.
  • The ablations are conducted on 1.2B-parameter models.

F.1 Sweeping Results for AdamW

The section lists hyperparameter-ablation tables for AdamW, NAdamW, Soap, and Muon at 130M and 300M parameters on 16x Chinchilla data, alongside comparisons with prior optimizer reports.

  • Prior Comparisons: Reported speedups vary across prior studies, including 1.4× for Soap, 2× for Muon, and 2–3× for SWAN and DION over AdamW.The cited works use different experimental setups and baselines.
  • Prior Comparisons: The reported Muon speedup ranges from 1.3× in the original Nanogpt speedrun to 2× in a later Kimi reproduction.The Kimi scaling experiments used a notably low AdamW learning rate of 8e-4 to 9e-4 for 400M–1.5B models.
  • Prior Comparisons: SWAN and DION comparisons use less than 1× Chinchilla data, a regime where matrix-based optimizer speedups may be larger.The comparison also omits communication cost, which is central to DION.
  • Prior Comparisons: SPlus reports a 2× AdamW speedup using a constant learning rate, an atypical training setup.This differs from experiments using learning-rate schedules.
Loading 2509.02046v2…