# 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

> Jointly tuning learning rate and weight decay eliminates about 90% of Muon-SW's reported advantage, leaving a small residual that grows with model scale.

- **Source:** [arXiv](https://arxiv.org/abs/2611.02529)
- **Published:** 2026-11-07
- **Permalink:** https://picx.dev/p/B4EuSR
- **Whiteboard:** https://picx.dev/p/B4EuSR/image

## Summary

## 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:

| Hypothesis | Prediction | Rationale |
|------------|------------|-----------|
| H_none | r = 1 | Schedule doesn't change optimal product |
| H_int | r = 1/$\bar{s}$ (1.364 cosine, 1.046 WSD) | Optimum matches total decay applied |
| H_end | r = 10 | Optimum matches final-step coefficient |

### 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**

| Rung | Active (M) | Untuned gap | Tuned gap | Survival | p*_const | p*_SW | Ratio | Seeds |
|------|------------|-------------|-----------|----------|-----------|--------|-------|-------|
| L3 | 1.89 | 0.03499 | 0.00359 | 10.3% | 0.192 | 0.269 | 1.40 | 4 |
| L4 | 3.89 | 0.03105 | 0.00403 | 13.0% | 0.253 | 0.414† | 1.63 | 1 |
| L6 | 15.04 | 0.02407 | 0.00514 | 21.3% | 0.331† | 0.631† | 1.91 | 1 |
| **L3 full 2D** | 1.89 | 0.03499 | **0.00261** | **7.4%** | 0.186 | 0.260 | 1.41 | 4 |

*†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**

| Schedule | Measured ratio | 95% CI | Consistent with | Closest |
|----------|---------------|--------|-----------------|---------|
| cosine | 1.41 | [0.88, 2.25] | H_none, H_int | H_int |
| WSD | 0.29 | [0.19, 0.43] | none | (H_none) |

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.

---

_Markdown view of https://picx.dev/p/B4EuSR, served by PicX — AI-generated visual whiteboard summaries of research papers._