Practical Efficiency of Muon for Pretraining

TL;DR

This paper shows that Muon, a second-order optimizer, improves training efficiency and data retention at large batch sizes, enabling more economical training for large models.

Contribution

It introduces Muon as an effective optimizer that expands the Pareto frontier over AdamW and combines it with muP for efficient hyperparameter transfer.

Findings

Muon outperforms AdamW in data efficiency at large batch sizes.

The combination of Muon and muP enables efficient hyperparameter transfer.

Validated on models up to four billion parameters.

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.