Summary of "Muon Makes the First Mixture-of-Experts Layer Harder to Balance"

Summary (Overview)

• The paper investigates how the choice of optimizer (Muon vs. AdamW) affects the difficulty of balancing expert loads in mixture-of-experts (MoE) language models, specifically measuring the "span" of routing biases (largest minus smallest bias in a layer).
• At tuned learning rates, Muon raised the span of the first MoE layer at every seed tested (roughly doubling it on average from 0.084 to 0.194), while other layers changed much less.
• The rise in span is attributed to Muon training the backbone (all non-router weights), which caused router inputs to share a common direction (input cosine ~0.4 vs. ~0.07 under AdamW), leading to a systematic expert preference that biases must counteract.
• Subtracting the batch mean from the first router's inputs removed the overload in 3,000-step probes with negligible loss change (<0.003), confirming the shared direction as the causal mechanism.
• At this scale, keeping the router on Muon is favored over moving it to AdamW, as the latter increased warmup load churn and slightly raised loss, while halving Muon's learning rate removed most of the span rise and improved loss slightly.

Introduction and Theoretical Foundation

The paper addresses a practical question in modern MoE language model training: should Muon train the router weights? Leading open models of 2026 disagree (DeepSeek-V4, Kimi K3, GLM-5.3 use Muon for routers; Qwen3.8-Flash-Next, MiMo-V2.6, IQuest-Q1 do not), yet no report provides comparative data.

The MoE Layer and Bias Balancing:

• MoE layers route tokens to a subset of experts. Left alone, routers tend to concentrate tokens on few experts, creating a feedback loop where busy experts get busier.
• Modern models use auxiliary-loss-free balancing (bias), where each expert has a bias added to its routing score, updated after every step:

$$b_i \leftarrow b_i + u \operatorname{sign}(1/E - \ell_i), \quad i = 1, \dots, E$$

with fixed update rate $u$ and load $\ell_i$ (fraction of tokens routed to expert $i$).

The Span Metric:

The paper's key measurement is the span: the largest minus smallest bias in a layer. Since adding a constant to all biases leaves routing unchanged, only differences matter:

• Span near zero: router scores naturally spread tokens evenly
• Large span: bias must compensate for systematic expert preferences
• Ceiling: span cannot exceed $1 + 2u$ since scores lie in $(0,1)$ and differences grow by at most $u$ per step

Muon Optimizer:

Muon orthogonalizes weight updates via Newton–Schulz iterations:

$$\mathbf{M}_t = \mu \mathbf{M}_{t-1} + \mathbf{G}_t, \quad \mathbf{W}_t = \mathbf{W}_{t-1} - \eta_t (0.2\sqrt{\max(p,q)}\operatorname{Ortho}(\mu\mathbf{M}_t + \mathbf{G}_t) + \lambda\mathbf{W}_{t-1})$$

Under Muon, every expert receives updates of similar spectral norm regardless of load, weakening the "busy gets busier" feedback—leading the authors to expect Muon would decrease span, contrary to what they found.

Methodology

Model Architecture: Two decoder-only Transformer MoE models with:

• Small: 8 blocks, width 512, 55M active / 303M total parameters
• Large: 12 blocks, width 768, 108M active / 496M total parameters
• $E = 32$ routed experts, $k = 4$ selected per token (sigmoid router), shared expert, dense FFN in block 0
• Trained on FineWeb-Edu (sample-10BT), sequences of 1,024 tokens, batch size 64

Parameter Partitions: Parameters divided into six groups (routers, routed experts, shared experts, attention, dense FFN, and embeddings/output head). The last group always uses AdamW. Key partitions:

Learning Rate Tuning: Each optimizer tuned separately on 0.5B tokens (small) or 1B tokens (large), selecting:

• Small: $2\times10^{-3}$ (AdamW), $4\times10^{-3}$ (Muon)
• Large: $10^{-3}$ (AdamW), $2\times10^{-3}$ (Muon), scaled by 0.7 for 3B token runs

Key Metrics:

• Input cosine $\rho$: mean cosine similarity between router inputs, measuring anisotropy
• Score offset: spread of mean expert scores (largest minus smallest offset)
• Overload $E \max_i \ell_i$: busiest expert's load as multiple of fair share (ceiling = 8)
• Load churn: total variation distance between successive loads
• Common share $\kappa$: fraction of router update along the shared direction

Empirical Validation / Results

Main Comparison (Small Model, 1B tokens, seeds 0-2):

The span difference was largest in the first MoE layer, especially early in training (0.74 vs. 0.19 over steps 501-1,000).

Causal Mechanism via Shared Direction:

The router input splits as:

$$\mathbf{w}_i^\top \mathbf{x}_n = \mathbf{w}_i^\top \mathbf{c} + \mathbf{w}_i^\top (\mathbf{x}_n - \mathbf{c})$$

where $\mathbf{c}$ is the batch-mean input. The first term creates a token-independent expert preference that biases must cancel.

• With backbone on Muon: input cosine ≈ 0.4 (small) and 0.61-0.73 (large), vs. ~0.07 under AdamW
• Centering the first router's input (subtracting batch mean) removed the overload: early overload dropped from 3.30-4.66 to 1.11-1.16
• Over steps 501-3,000, the span rise fell by 79% with centering, while loss changed by <0.003

Learning Rate Effects:

• Halving Muon's rate (to AdamW's rate) reduced the span excess from +0.103 to +0.024 (75% reduction)
• Half-rate Muon trained better at 1B tokens (loss 3.166 vs. 3.175), consistent with optimal rate falling with training length
• The shared direction persisted at half rate, but the router row spread along it decreased

Router Optimizer Comparison (Backbone on Muon):

Moving the router to AdamW consistently increased warmup load churn and slightly raised loss. The common share explains this: AdamW preserves the gradient component along the shared direction (0.497 vs. 0.160 for Muon), amplifying load shifts.

Additional Findings:

• Fast bias (4× rate) halved the span rise and cut excess early overload from 2.84 to 0.54
• Balance loss reduced the score offset and early overload but did not eliminate the span rise
• Partial backbone partitions (experts only, attention+dense FFN) raised span more than full Muon backbone in some cases without reproducing the shared direction
• The large model's first MoE layer remained overloaded at end of training (2.2-2.8× fair share vs. 1.07 under AdamW)
• Loss spikes occurred only with router on AdamW in large-model runs, correlated with span near its ceiling

Theoretical and Practical Implications

Mechanistic Understanding: The paper provides a clear mechanism for optimizer-induced routing difficulty:

1. Muon on the backbone causes first-router inputs to develop anisotropy (shared direction)
2. The shared direction creates a token-independent expert preference via the batch-mean offset term
3. Biases lag behind (growing by at most $u$ per step), causing transient overload
4. The effect is specific to the first MoE layer, consistent with reports from production models

Practical Recommendations for Small Models:

• Keep router on Muon when backbone is on Muon (avoids warmup churn and slight loss penalty)
• Halving Muon's learning rate substantially mitigates the span rise while improving loss
• Input centering is effective but tested only in short probes
• The fast bias provides a practical mitigation for overload

Connection to Production Models: The results contextualize reports from Nemotron 3 Ultra (busiest first-MoE expert at 4.8× fair share early, ~12× by 12T tokens), DeepSeekMoE's dense first layer, and IQuest-Q1's similar choice. The first MoE layer appears structurally challenging across optimizers and scales.

Limitations: Models are 50-2,000× smaller than production models; single architecture and dataset; three seeds per comparison; probes cover only 3,000 steps; the sigmoid router's bounded scores differ from production models using unbounded scoring.

Conclusion

Muon training raises the span of the first MoE layer to roughly double that of AdamW at tuned rates, caused by a shared direction in router inputs induced by Muon on the backbone. This creates transient overloads of 4-7× fair share that biases take hundreds of steps to correct. Key mitigations include halving Muon's learning rate (removes ~75% of the rise with better loss), input centering (removes overload with negligible loss change), and faster bias updates (halves the rise). At this scale, keeping the router on Muon is favored. The paper identifies three open questions: what spreads router rows along the shared direction, why partial backbone partitions raise span without the shared direction, and why only the full backbone on Muon produces the anisotropic inputs. Future work should determine whether these effects grow or settle at production scale and whether a running mean can replace the batch mean in practical centering.