Practical Efficiency of Muon for Pretraining

Abstract: We demonstrate that Muon, the simplest instantiation of a second-order optimizer, explicitly expands the Pareto frontier over AdamW on the compute-time tradeoff. We find that Muon is more effective than AdamW in retaining data efficiency at large batch sizes, far beyond the so-called critical batch size, while remaining computationally efficient, thus enabling more economical training. We study the combination of Muon and the maximal update parameterization (muP) for efficient hyperparameter transfer and present a simple telescoping algorithm that accounts for all sources of error in muP while introducing only a modest overhead in resources. We validate our findings through extensive experiments with model sizes up to four billion parameters and ablations on the data distribution and architecture.

• The paper introduces Muon, a second-order optimizer that expands the Pareto frontier over AdamW by optimizing compute-time tradeoffs and data efficiency.
• Muon employs matrix-structured steepest descent with Newton-Schulz iteration and Nesterov momentum, integrated with muP scaling for robust hyperparameter transfer.
• Experiments on transformer models show that Muon achieves target loss faster and allocates resources more effectively, validated by telescoping strategies across model widths.

Practical Efficiency of Muon for Pretraining

Introduction

The paper "Practical Efficiency of Muon for Pretraining" introduces Muon as an efficient second-order optimizer expanding the Pareto frontier over AdamW by enhancing data efficiency at large batch sizes without sacrificing computational advantage. While AdamW has dominated the landscape due to its data-efficiency and marginal compute cost increase, Muon challenges this status quo by maintaining superior performance in the compute-time tradeoff. This tradeoff enables practitioners to optimize resource allocation effectively, necessitating studies that characterize the modification in this tradeoff when adopting different optimizers.

Figure 1: Muon expands the Pareto frontier over AdamW on the compute-time tradeoff, maintaining data efficiency at large batch sizes.

Muon Improves the Compute-Time Tradeoff

Review of Muon

Muon employs matrix-structured steepest descent with spectral norm regularization:

$$O_t = \argmin_{ O \in \mathbb{R}^{m \times n}:\; ||O||_2 \leq 1} tr(G_t^\top O)$$

Where $G_t = U \Sigma V^\top$ is the singular value decomposition (SVD) of gradients. Muon approximates the SVD using Newton-Schulz iteration, combines it with Nesterov momentum, learning rate scaling, and weight decay. At each step, given the gradient $G_t$, it updates weights $W_{t+1}$ using momentum hyperparameters.

Experimental Setup

The study employs modern decoder-only transformer models and a mixture of high-compression text and code data from DCLM and Python to validate Muon's performance. The authors conduct extensive experiments by varying batch size, achieving close to 50% model FLOP utility on TPU v5p chips. Hyperparameter tuning reveals that Muon can achieve target loss significantly faster than AdamW in wall time.

Figure 2: Muon expands the Pareto frontier over AdamW at various loss thresholds on Python and DCLM datasets.

Characterization of Compute-Time Tradeoff

On the compute-time plane, optimizers are represented as iso-loss curves, comparing the total training time against the number of devices. Muon demonstrates superior flexibility in resource allocation compared to AdamW by maintaining data-efficiency with large batch sizes, confirmed by observing non-decreasing token ratio even in large-batch regimes.

Figure 3: Relative data efficiency of Muon over AdamW at varying batch sizes across different model sizes.

Choosing Hyperparameters for Muon

The Maximal Update Parameterization (muP)

Muon is compatible with muP scaling, which permits consistent predictive hyperparameter transfer across model scales by regulating initialization and learning rate scaling. This resolves how hyperparameters calibrated on smaller models transfer efficiently to larger models, overcoming estimation errors by a telescoping algorithm.

Figure 4: Token ratios to loss portray Muon's practical advantages over AdamW for large batch training.

Telescoping Strategy for Hyperparameter Transfer

The telescoping protocol ensures efficient hyperparameter tuning by reducing sweep points systematically as width increases. This approach bounds hyperparameter tuning overhead while maintaining access to optimal hyperparameters due to the controlled drift of optimum values as model width increases. The telescoping protocol was validated through family-size transformer models demonstrating efficient hyperparameter transfer to pretraining with Muon.

Figure 5: Telescoping algorithm visualization applied to weight decay and learning rate across model widths.

Related Work

Comparative studies on optimizers like AdamW and emerging second-order techniques such as Shampoo and Muon highlight the advantages offered by Muon. These findings are consistent with efforts in understanding optimizer batch size dynamics, with Muon expanding upon these foundations with increased compute-time flexibility demonstrated by its data efficiency and automated parameter coverage.

Figure 6: Comparison of the best 1B model training runs under AdamW and Muon demonstrating wall-time efficiency.

Conclusion

Muon emerges as a robust alternative for large-scale pretraining, delivering superior compute-time tradeoffs and flexible resource allocations compared to AdamW. This advantage is sustained throughout large batch sizes due to Muon's efficient token consumption. Hyperparameter tuning with muP further enhances Muon's applicability by ensuring scaling efficiency, thus establishing Muon as a practical replacement for AdamW in extensive pretraining tasks. With broadened empirical validation across multiple setups, Muon stands out as a practically efficient optimizer for distributed neural network training.

Overall, Muon optimizes resource strategies more effectively, rendering it a preferred choice for practitioners focusing on scale-efficient pretraining within the landscape of modern compute infrastructures.