Tuning weight decay cuts most muon efficiency gains but advantage grows with size
Tune the Baseline First: Most of Muon's Scaled Weight Decay Gain Is a Tuning Gap, but Its Small Real Advantage Grows With Model Size
Computation and LanguageMachine Learning
Summary
Training language models efficiently often involves adjusting settings like learning rate and weight decay. The authors show that most of Muon-SW's apparent advantage over standard weight decay is due to not properly tuning these settings for both methods. After careful tuning, Muon-SW still wins slightly, especially as model size grows larger. This means that while tuning is crucial for fair comparisons, Muon-SW does offer a small real benefit that increases with bigger models.
What this means in practice
- •For machine learning engineers: Adjust training schedules for MoE language models to improve efficiency by properly tuning weight decay and learning rate jointly.
- •For ai service operators: Optimize deployment configurations for large language models by leveraging Muon-SW’s advantages to reduce compute costs at scale.
Authors
Esther Xin
Abstract
Scaled weight decay for Muon (Muon-SW) has been reported to train language models more efficiently than constant decoupled weight decay, but such comparisons usually reuse one learning-rate (LR) and weight-decay (WD) setting for both methods. We re-run the comparison under the MAI-Base-2 protocol: LR and WD are tuned jointly through their product for each method, on a source-scaled mixture-of-experts (MoE) scaling ladder (1.9M to 15M active parameters) that holds tokens per active parameter and data repetition fixed across sizes. We find: (1) Most of the gain is a tuning gap. At the shared setting Muon-SW wins by 0.035 bits per byte (BPB), a 1.52x compute saving; after both methods are tuned, the gap is 0.0026 BPB (95% CI [0.0018, 0.0034]), a 1.03x saving, so about 90% of the advantage disappears. (2) The small advantage that remains is real: at the smallest size, all four paired seeds favour Muon-SW at the tuned settings. (3) That advantage grows with model size: across an 8x range of active parameters, the tuned gap rises from 0.0036 to 0.0051 BPB while the untuned gap shrinks from 0.035 to 0.024, so the surviving share rises from about 10% to 21%. (4) The best weight decay moves with scale: the optimal LRWD product grows for both methods, faster for Muon-SW (ratio 1.40 to 1.91), and a law fitted on the two smallest sizes overshoots the third by 1.3-1.5x, costing up to 0.002 BPB. The whole study runs on one NVIDIA T4 GPU; we release the code and all run results.