Summary (Overview)

• Core finding: Approximately 90% of Muon's scaled weight decay (Muon-SW) advantage over constant decoupled weight decay disappears when both methods are jointly tuned for learning rate (LR) and weight decay (WD) through their product.
• Residual advantage: A small but real advantage (0.0026 BPB at smallest scale, 95% CI [0.0018, 0.0034]) remains after tuning, with all four paired seeds favoring Muon-SW.
• Scale dependence: The surviving advantage grows with model size—from ~10% to 21% of the untuned gap across an 8× range of active parameters (1.9M to 15M).
• Optimal hyperparameter drift: The optimal LR·WD product increases with scale for both methods, faster for Muon-SW (ratio 1.40 → 1.91), meaning baselines tuned at small scale become increasingly mis-tuned at larger scales.
• Practical significance: At small scale, tuning constant decay's LR·WD product captures most of Muon-SW's benefit; the untuned comparison overstates gains roughly tenfold.

Introduction and Theoretical Foundation

Background and Motivation

New optimizer variants are typically evaluated by training a model with the variant and a baseline at one shared hyperparameter setting. This creates a tuning gap problem: when a variant changes how strongly a hyperparameter acts (as a weight-decay schedule does), the shared setting may be well-tuned for one method but poorly tuned for the other.

The concern is particularly acute for weight decay because:

• Its effect is coupled to the learning rate through their product (Loshchilov and Hutter, 2019; Kosson et al., 2023)
• Its optimum moves with model and data size (Wang and Aitchison, 2024)

Theoretical Foundation

Muon (Jordan et al., 2024) is an optimizer used for hidden-layer matrices of large language models, including the 2.8T parameter MAI-Base-2. Muon-SW (Apte, 2026) scales Muon's weight-decay coefficient over training, reporting gains over constant decay.

The MAI-Base-2 protocol enables cleaner testing through two practices:

1. Joint LR·WD tuning: Sweeping hyperparameters on a grid in (LR, LR·WD) because good settings lie on a narrow diagonal of the (LR, WD) plane
2. Source scaling: Comparing models on a scaling ladder with fixed tokens per active parameter (TPP), keeping data repetition fixed across rungs

Key quantities:

• Untuned gap = constant-decay BPB − Muon-SW BPB at p = 1
• Tuned gap = difference between methods' tuned BPB
• Survival = tuned gap / untuned gap

Methodology

Methods Compared

Both methods use Muon for non-embedding 2D parameters and AdamW for all other parameters:

• Constant decay: Applies decoupled weight decay with fixed coefficient λ: θ ← θ(1 − η_t λ)
• Muon-SW: Multiplies the coefficient at step t by s(t) = η_t/η_max = f(t), so decay is applied at full strength at peak LR and weakens as LR anneals: θ ← θ(1 − η²_t λ/η_max)

The ratio of total decay applied by Muon-SW to constant decay over a run is:

$$\bar{s} = \frac{\sum_t f(t) s(t)}{\sum_t f(t)}$$

• Under cosine schedule: $\bar{s} = 0.733$
• Under WSD schedule: $\bar{s} = 0.956$

Joint LR·WD Tuning Protocol

Following MAI-Base-2, each rung has a law-predicted LR η₀ and product (ηλ)₀. Runs use:

$$\eta = m_\eta \eta_0, \quad \eta\lambda = p(\eta\lambda)_0$$

• L3: Full 2D grid (m_η ∈ {0.5, 1, 2}, p on factor-of-2 grid), 2D quadratic fit per method
• L4 and L6: 1D sweep of p at m_η = 1, 1D quadratic fit in log₂p

Preregistered Hypotheses for Ratio of Optima

Three predictions for r = p_SW / p_const:

Ladder, Data and Training

• Three rungs: L3 (1.89M active), L4 (3.89M active), L6 (15.04M active parameters)
• Data: FineWeb-Edu (60%), CodeParrot-clean (25%), OpenWebMath (15%); 8,192-token byte-level BPE tokenizer
• Training: 120 tokens per active parameter; source scaling holds data repetition fixed (web seen once, code/math four times)
• Metric: Held-out bits per byte (BPB), pooled over three sources
• Compute: Single NVIDIA T4 GPU with pre-emption-safe pipeline (checkpointing every 15 minutes)

Empirical Validation / Results

Main Results

Table 1: Main Results Across Rungs

†Fitted optimum lies outside sampled range; †tuned gap 95% CI [0.00183, 0.00338]; Muon-SW better on 4/4 seeds

Key Findings

1. Untuned comparison: Muon-SW wins by 0.035 BPB at shared setting (1.52× compute saving), reproducing reported gains.

2. After joint tuning: Gap collapses to 0.0026 BPB (7.4% survival in 2D analysis, 10.3% in 1D slice), a 1.03× saving—about 90% of advantage disappears.

3. Residual is real: All four paired seeds favor Muon-SW at L3 tuned optima; 95% CI excludes zero but doesn't clear the 0.002 BPB practical significance threshold.

4. Scale trend: Untuned gap shrinks (0.0350 → 0.0311 → 0.0241) while tuned gap grows (0.0036 → 0.0040 → 0.0051); survival rises from ~10% to 21% over 8× range.

Hypothesis Testing Results

Table 2: Ratio of Optimal LR·WD Products at L3

Under cosine, data are consistent with no change and with matched integrated decay; endpoint matching (H_end, r = 10) is ruled out. Under WSD, the ratio falls below every prediction—an unresolved anomaly.

Extrapolation Test

Extrapolating L3–L4 optima to L6 predicts multipliers of 0.44 (constant) and 0.95 (Muon-SW), vs. fitted values of 0.33 and 0.63. Using predicted values costs 0.0014 and 0.0020 BPB respectively—comparable to the tuned gap itself.

Theoretical and Practical Implications

For Optimizer Comparison Methodology

• Tuning gap is pervasive: Optimizer comparisons that hold LR·WD product fixed overstate gains roughly tenfold in this setting
• Scale matters for tuning: The optimal LR·WD product moves with scale (faster for Muon-SW), so baselines tuned at small scale become increasingly mis-tuned at larger scales
• Direction of error depends on method: Constant decay is more penalized by over-large product, explaining why untuned gaps shrink with scale

For Muon-SW Adoption

• At small scale, tuning constant decay's LR·WD product captures most of Muon-SW's benefit
• At larger scales, Muon-SW shows a small genuine improvement whose importance grows with size
• The data are consistent with Muon-SW being a real but modest improvement, not the dramatic gain reported in untuned comparisons

Practical Recommendations

Practitioners should:

1. Tune the baseline at every scale reported in optimizer comparisons
2. Use tuned comparisons at the target scale when making claims about scaled weight decay
3. Be aware that LR·WD product optima move with scale and differ between methods

Conclusion

Most of Muon-SW's reported advantage is a tuning gap: when constant decay's LR·WD product is tuned as carefully as the method's, about 90% of the gain disappears at small scale. What remains is small but consistent, growing from about 10% to 21% of the untuned gap across an 8× range of active parameters, while the optimal products of the two methods move apart.

Key takeaways:

• Optimizer comparisons should tune the baseline at every scale they report
• Claims about scaled weight decay should be made from tuned comparisons at the target scale
• The residual advantage of Muon-SW grows with scale, warranting further investigation at production scales

Future directions: The WSD anomaly (ratio 0.29, below all predictions) remains unresolved and may reflect an interaction between the long constant-LR phase of WSD and the decay schedule. The planned L9 rung (51M active) was not run for compute reasons; extrapolating the growth of the residual beyond L6 is not supported by the data.