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

概要

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.

参考リンク

• Paper ↗

リンク先のライセンス・再現性・データ条件をご確認ください。自動分類や実装の目安は保証ではありません。

実装を始めるプロンプト

参考リンクを確認したうえで、次の内容をコピーして利用できます。

# 論文・記事
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

# 概要
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…

# 原典
https://arxiv.org/abs/2611.02529

上記の手法を実装してください。実行前に依存関係とライセンスを確認してください。