Muon optimizer improves loss with limited data and tuning budgets
When Does Muon Actually Win? A Paired Data-, Time-, and Tuning-Budget Benchmark
Machine LearningArtificial Intelligence
Summary
Machine learning models learn by adjusting their settings with methods called optimizers. This paper compares a newer optimizer called Muon to a common one named AdamW using small language and vision models on a single GPU. The authors find that Muon generally reduces errors more efficiently when given the same amount of training data and tuning effort, especially for vision models. However, Muon takes more time per training step, and AdamW sometimes performs better on language models when considering time. These results apply to smaller models and specific setups, so conclusions for larger models or different implementations are not certain.
What this means in practice
- •For machine learning engineers: Choose Muon when training small vision transformer models under strict data and tuning budgets to reduce validation loss more efficiently than AdamW.
- •For ai model deployment teams: Adjust optimizer choice between Muon and AdamW based on hardware time constraints and the model type to optimize training efficiency on single GPUs.
Authors
Mushir Akhtar, Hazem N. Nounou, Majdi Mansouri, Mohamed N. Nounou
Abstract
Muon's gains in data efficiency motivate evaluating its performance under explicit data, time, and tuning budgets. On a single A100 GPU, we compare a reference single-device Muon implementation with AdamW across 17.7M- and 33.6M-parameter causal language models and 2.70M-parameter vision transformers. A common 23-candidate development budget selects configurations evaluated with five matched seeds, full validation trajectories, optimizer-step timing, and a frozen one-shot test. Muon reduces validation-loss area under the curve by 3.25--22.44\% at equal data and lowers test loss in all 20 primary pairs. In this configuration, Muon improves time-normalized loss area} by 19.68--22.31\% on vision, while AdamW is better by 3.38--6.96\% on language. The full Muon optimizer call, including auxiliary AdamW updates, consumes 56.1--69.6\% of training time, compared with 6.1--9.1\% for the AdamW baseline. Muon also has higher expected regret under restricted access to these candidate pools. These results characterize small models with reference kernels; their wall-clock ordering does not establish a ranking for larger models or optimized implementations.