cs.CLNov 2, 2026

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

Authors: Esther Xin

Organizations: Independent Researcher

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.

Figures & tables

Table 1: Main result. Untuned gap : constant-WD BPB minus Muon-SW BPB, both at the MAI-law setting ( p=1p=1 ). Tuned gap : the same difference after each method is tuned over the LR ⋅\cdot WD multiplier pp (quadratic fit in log⁡2p\log_{2}p ). Survival = tuned gap / untuned gap. p⋆p^{\star} is the fitted optimal multiplier; Ratio is pSW⋆/pconst⋆p^{\star}_{\rm SW}/p^{\star}_{\rm const} . Upper block: one-dimensional sweep at learning-rate multiplier 1 for every rung. Lower block: the full two-dimensional (LR ×\times LR ⋅\cdot WD) L3 analysis, confirmed by four paired seeds at the fitted optima. † Fitted optimum lies outside the sampled range; the tuned BPB is the better of the fitted and best observed value. All BPB values are on held-out data.
Table 2: Measured ratio of optimal LR ⋅\cdot WD products (Muon-SW over constant decay) at L3, with bootstrap 95% intervals, against the three preregistered predictions (§ 3.4 ). Under the cosine schedule the data are consistent with no change and with matched integrated decay, and rule out endpoint matching. Under WSD the ratio falls below every prediction.
Figure 1: Held-out BPB against the LR ⋅\cdot WD multiplier pp (log scale) for each rung. Points are runs; curves are quadratic fits in log⁡2p\log_{2}p ; dotted lines mark fitted optima; the dashed line marks the MAI-law setting p=1p=1 . At p=1p=1 Muon-SW (orange) is far better, but constant decay (blue) recovers most of the gap by moving to a smaller pp .
Figure 2: Tuned gap at L3 for each of four paired seeds at the 2D-fitted optima. Black line: mean; grey band: 95% CI; dashed: the preregistered 0.0020.002 BPB threshold.
Figure 3: Left: untuned and tuned gaps against active parameters; the diamond is the four-seed L3 estimate with its 95% CI. Right: share of the untuned gap that survives tuning.
Figure 4: Left: fitted optimal LR ⋅\cdot WD multiplier per method (circles) and the L6 values predicted from L3 and L4 (crosses). Right: ratio of the two optima with the HnoneH_{\rm none} and HintH_{\rm int} predictions for the cosine schedule.
Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Table 3: Scaling-ladder rungs. Each rung is defined by its layer count LL , which fixes the rest of the architecture. Tokens are for 120 training tokens per active parameter. Throughput is measured on one NVIDIA T4 (§ 3.6 ). L9 was planned but not trained.
Table 4: Transfer of the LR ⋅\cdot WD law to L6. The optimal multiplier predicted by extrapolating the L3 and L4 optima overshoots the fitted L6 optimum by 1.31.3 – 1.5×1.5\times for both methods. Regret is the BPB lost by using the predicted rather than the fitted value.
More in this topic

Weight Decay

Weight decay is a parameter-shrinkage method studied for its effects on generalization, optimization, model structure, and training stability.

See all 28 papers

Explore similar work

CardsList
Jul 26, 2026

Scale Weight Decay and Train Better

The discovery of scaling laws has motivated training neural networks on ever increasing quantities of data. This is typically done with a constant decoupled weight decay which causes the network weights to shrink steadily over the course of training. Taking inspiration from the Robbins--Monro conditions, we propose to scale weight decay by the fraction of the peak learning rate η/ηmax⁡η/η_{\max}. We prove that this scaled weight decay preserves the asymptotic stationarity guarantees of the corresponding unregularized methods for both stochastic gradient descent and the non-Euclidean spectral optimizer Muon, thereby avoiding the additional asymptotic bias introduced by constant decoupled weight decay. This retains the stability benefits of weight decay without changing the asymptotic optimization target. Using a steady-state analysis, we explain why under standard weight decay the weight norm shrinks steadily as training proceeds, whereas under scaled weight decay it settles to a roughly constant value. When applied to the training of mixture-of-experts models, Muon with scaled weight decay (Muon-SW) consistently outpaces Muon with identical hyperparameters, reaching the same validation loss 30%\mathbf{30\%} faster at our largest scale across models from 72−93072 - 930 million parameters trained at ∼600\sim 600 tokens per active parameter. If this trend continues to hold, the method promises to substantially accelerate the pre-training of frontier models while requiring only a few lines of code to implement.

Anuj Apte

Global Technology Applied Research, JPMorganChase, New York, NY 10001

Weight DecayNeural Network OptimizationStochastic Gradient Descent
AskPDF
Nov 2, 2026

Pretraining Data Scale Reverses the Muon-AdamW Fine-Tuning Gap

Muon has recently emerged as a promising optimizer for pretraining large language models, offering a speedup over AdamW in minimizing pretraining loss. However, it remains unclear whether Muon's advantage persists after fine-tuning. In this paper, we study the downstream behavior of Muon and AdamW and uncover a surprising trend: beyond a sufficiently large pretraining data budget, AdamW outperforms Muon after fine-tuning. Moreover, as model size increases, AdamW overtakes Muon at progressively smaller pretraining token-parameter ratios. To explain this reversal, we view fine-tuning as a perturbation of the outputs and introduce the margin-to-spread ratio, which measures how far the correct logit lies above the wrong ones relative to how widely the wrong logits are spread. In language models, Muon has the larger margin-to-spread ratio at small data budgets and AdamW at large ones, mirroring the fine-tuning trend. In a linear associative memory, we prove that after one update the margin-to-spread ratio of gradient descent gains on that of Muon as the data budget grows, and simulations show that this trend persists after extended training.

Catherine Li, Jingchu Gai, Andrej Risteski +2

Carnegie Mellon University

Muon OptimizerLLM Fine-TuningFine-Tuning
AskPDF
May 11, 2026

Muown: Row-Norm Control for Muon Optimization

Muon has emerged as a strong competitor to AdamW for language model pre-training, yet its behavior at scale is sensitive to weight decay. Recent work has observed that, for Muon without decoupled weight decay, the spectral norm of weight matrices drifts upward over training. Through a decomposition of the spectral norm into a row-magnitude factor and a row-coherence factor, we identify the former as the empirical driver of this drift under Muon, while the latter remains well-behaved along the trajectory. Motivated by this diagnosis, we introduce Muown, a drop-in replacement for Muon that treats the row-magnitude vector as an explicit optimizer variable, updating it under the ℓ∞\ell_\infty geometry induced by the decomposition, while applying Muon unchanged to the remaining direction component. We prove that Muown attains the optimal non-convex rates in both deterministic and stochastic regimes under a dual norm aligned with the underlying geometries and with a stochastic noise coefficient that empirically remains below that of Muon throughout training. Across GPT-style pre-training on FineWeb-Edu with model sizes from 124M up to 2.7B parameters, Muown improves perplexity over Muon, SOAP, AdamW, and Lion. It also widens the plateau of near-optimal learning rates across model scales, reduces sensitivity to weight decay, and avoids the spectral norm drift at negligible step-time overhead when appropriately sharded.

Kai Lion, Florian Hübler, Bingcong Li +2

Department of Computer Science, ETH Zurich, Switzerland · Department of Mathematics, School of Computation, Information and Technology; Technical University of Munich, Germany · ELLIS Institute Tübingen, MPI-IS, Tübingen AI Center, Germany

Language Model PretrainingMuon Optimizer
AskPDF