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

AI summary

Overview

Research area: Machine learning — training dynamics of mixture-of-experts (MoE) language models, specifically how the choice of optimizer (AdamW versus Muon) affects expert-load balancing.

Technical level: Intermediate to Advanced. The paper is readable with a working knowledge of Transformer training, but its core measurements (bias span, input cosine similarity, common share of a router update) and the Muon update rule require familiarity with optimizer internals.

Scope: A controlled, small-scale empirical study (three seeds per comparison, models of 55M and 108M active parameters trained on FineWeb-Edu) showing that Muon, applied to the network outside the router, makes the first MoE layer substantially harder for the standard expert bias to balance.

What This Paper Is About

Most recent open MoE language models balance expert load with a per-expert bias that is nudged up or down after every step, and many of them train their weight matrices with Muon rather than AdamW. Their published recipes disagree on whether the router itself should also be trained with Muon, and no report gives data comparing the two choices. This paper measures how hard the bias has to work — its span, the largest minus the smallest bias in a layer — under each optimizer, and finds that Muon roughly doubles the span in the first MoE layer while leaving the other layers largely unchanged.

Key Contributions

1. A new way to measure balancing difficulty. The authors propose the bias span (largest minus smallest bias in a layer) as a proxy for how hard a layer is to balance, and give it a theoretical ceiling of 1 + 2u, where u is the bias update rate. They complement it with a set of diagnostics: input cosine similarity, score offset, overload, load churn, common share, update–load correlation, usage perplexity, and specialization.

2. A partition-based experimental design. Rather than switching the optimizer for the whole model, they split parameters into six groups (routers, routed experts, shared experts, attention, dense FFN of block 0, and embedding/head/norm gains) so that the router, the experts, and the rest of the network can be moved to Muon separately. This isolates which part of the network causes the effect.

3. The empirical finding that Muon raises the first layer's span. At learning rates tuned per optimizer on shorter runs, Muon raised the first-MoE-layer span from 0.084 to 0.194 on average at every seed, and again at all nine seeds of three later experiments (probability ≈ 2⁻⁹ ≈ 0.002 that two identical arms would produce this).

4. A mechanistic explanation and a cheap fix. The rise traces to the backbone, not the router. Under Muon the first router's inputs share a direction (mean cosine similarity about 0.4 versus 0.07 under AdamW), and the mean input gives every token the same expert preference. Subtracting the batch mean from the first router's input removed the overload over the first 3,000 steps while changing the loss by less than 0.003.

Main Findings

• Muon roughly doubles the first layer's span. In the main comparison (small model, 1B tokens, seeds 0 to 2), the first MoE layer's span went from 0.084 under AdamW to 0.194 under Muon, a paired difference of +0.110 (range +0.085 to +0.151). Averaged over all layers the span rose from 0.105 to 0.137.

• The other layers barely move. Every MoE layer except the first kept its integrated overload below 1.32 under both optimizers, and the remaining six layers differed in span by 0.02 on average.

• The effect is largest early and persists. Over steps 501 to 1,000 the first-layer span was 0.74 under Muon against 0.19 under AdamW; over steps 8,001 to 15,259 it was still 0.139 against 0.062.

• The prediction was the opposite. The authors expected Muon to weaken the rich-get-richer feedback, because under Muon an expert's update has approximately the same spectral norm regardless of how many tokens it receives. The update–load correlation did fall as predicted, from 0.379 to 0.024, yet the span rose anyway.

• The backbone, not the router, drives the effect. In the ablation, the first-layer span was 0.198 with backbone and router both on Muon, 0.175 with the backbone on Muon and the router on AdamW, and 0.099 and 0.095 with the backbone on AdamW. Changing only the router's optimizer moved the span by 0.023 or less on average, in a seed-dependent direction.

• The router inputs become anisotropic. The input cosine similarity of the first router was about 0.4 with the backbone on Muon against 0.07 on AdamW. The common share of the router update was 0.160 under Muon and 0.129 under AdamW, against a bound of about 0.18 for E = 32.

• Concrete overload. With the bias alone at its usual rate, the busiest expert of the first MoE layer received on average four to seven times its fair share soon after warmup. The early overload was 4.06 with the full Muon partition and 6.63 with the router on AdamW, against 1.14 for AdamW. The ceiling for these models (E = 32, k = 4) is 8.

• Centering the first router's input removes the overload. In probes repeating the first 3,000 steps of both models, subtracting the batch mean from this router's input removed the overload while changing the loss by less than 0.003.

• Halving Muon's learning rate helps on the small model. It removed about three quarters of the rise (span 0.119 versus 0.198), gave a slightly better loss (3.166 versus 3.175), and left the shared direction in place (input cosine 0.401).

• A larger span can occur without the shared direction. Two partitions that put only part of the backbone on Muon did not reproduce the shared direction yet raised the first-layer span more than the full backbone did: experts on Muon gave a span of 0.409 and attention plus dense FFN on Muon gave 0.229.

• A faster bias partially compensates. A bias moving at 4u per step halved the rise and removed most of the overload after warmup. The rise also persisted with an added sequence-wise balance loss (weight α = 10⁻⁴) and on the large model, where the first MoE layer was still overloaded at the end of training.

• Keeping the router on Muon looks better at this scale. Moving the router to AdamW made loads shift more during warmup on both model sizes and at every learning rate of the stress test, and at the tuned rates it raised the loss slightly at every seed. In the large-model experiment, loss spikes occurred only in the Muon runs with the router on AdamW.

• Muon still wins on loss. In the main comparison Muon reached a validation loss of 3.177 against 3.255 for AdamW, a difference of −0.078, and a higher specialization (0.118 versus 0.105).

Methodology in Plain English

The authors train two small decoder-only Transformer MoE models from scratch on a 10-billion-token sample of FineWeb-Edu. The small model has 8 blocks of width 512 and 55M active parameters out of 303M; the large one has 12 blocks of width 768 and 108M active out of 496M. Both use 32 routed experts per MoE layer, select 4 per token, keep a dense FFN in block 0, and balance load only with the DeepSeek-V3 style bias at rate u = 0.001.

Instead of flipping one global switch, they divide parameters into groups and assign AdamW or Muon to each group independently. Runs sharing a seed start from identical weights and see identical batches, so every comparison is paired by seed and reported as the mean of three paired differences with their range.

Learning rates are tuned separately for each optimizer on 0.5B-token runs (selecting 2×10⁻³ for AdamW, 4×10⁻³ for Muon, and 5×10⁻⁴ for a non-orthogonalized control), then the models train on 1B tokens (small) or 3B tokens (large). A "stress test" sweeps learning rates at 1, 2, 4, 8, and 16 times the tuned values with no decay. The key metric is the bias span, supplemented by cosine similarity between router inputs, the load of the busiest expert as a multiple of its fair share, and how much of each router update acts along the shared direction of the inputs. Comparisons were pre-registered as "planned" before training, and everything else is labeled exploratory.

Why This Matters

Impact on research. The paper supplies the first direct measurements comparing router-on-Muon against router-on-AdamW in an MoE balanced by the bias, at a time when leading open models (DeepSeek-V4, Kimi K3, GLM-5.3 on one side; Qwen3.8-Flash-Next, MiMo-V2.6, IQuest-Q1 on the other) have split on exactly this choice without published evidence. It also reframes the problem: since the bias forces loads toward fair shares regardless of what the router prefers, the loads themselves say little about the router, and the span is a better probe. The finding that the first MoE layer is the fragile one connects to a body of reports (OLMoE, Motif 3, Trinity Large, DeepSeekMoE's dense first layer, DeepSeek-V4's hashed first three MoE layers, Nemotron 3 Ultra's 4.8× to 12× overload) that singled out early layers without knowing why.

Real-world applications:

• Training-cost reduction: a first layer that stays overloaded wastes capacity and risks divergence; catching it with a cheap diagnostic (the span) or a cheap fix (centering the first router's input) avoids expensive restarts.
• Optimizer selection for new MoE training runs: teams choosing between Muon and AdamW now have evidence about a specific side effect rather than only a loss comparison.
• Router design and normalization: the centering result is a practical, low-risk architectural tweak for the router input.
• Diagnostics for large runs: the span and overload metrics can be logged in production pretraining to flag imbalance before it becomes a loss spike.

Industry relevance. Muon is now the default for weight matrices in several frontier open MoE recipes, so any systematic interaction between Muon and the balancing mechanism matters directly to anyone reproducing or extending those recipes. The paper's practical recommendation — keep the router on Muon when the backbone is on Muon, and watch the first MoE layer — is actionable, and its warning that halving Muon's rate removed most of the rise suggests the effect is partly a learning-rate artifact.

Future Directions

• Scale up the comparison. The authors state plainly that their models have about 50 to 2,000 times fewer active parameters than the models whose recipes motivate the question, so whether the first-layer rise and the router-on-AdamW penalty persist at frontier scale is open.
• Explain the partial-backbone anomaly. Two partitions that put only part of the backbone on Muon raised the span more than the full backbone without producing the shared direction, so a large span can arise by a second, unexplained route.
• Test other balancers and routers. The main experiments use the fixed-rate sign update at u = 0.001; the paper's own related-work section describes alternatives (Kimi K3's quantile rule, Trinity Large's damped update, φ-balancing, proportional-integral readings of the bias) that were not compared here.
• Turn centering into an architectural change. Subtracting the batch mean from the first router's input was tested only in 3,000-step probes; whether it helps or hurts over full training runs, and whether it should be applied to all routers, is untested.

Target Audience

Researchers and engineers training or reproducing mixture-of-experts language models, particularly those adopting Muon or deciding which parameter groups should receive it. It is also useful for anyone working on load balancing, routing collapse, or training stability in sparse models, and for students looking for a compact example of a carefully controlled, seed-paired ablation study with pre-registered predictions. Readers without background in MoE routing or the Muon update rule will need to consult the background section and its references first.

Authors’ abstract

Many open mixture-of-experts (MoE) language models train with Muon instead of AdamW. Most balance the load with a bias on each expert's routing score, adjusted after every step. Their recipes disagree on whether Muon should also train the router. In models with 55M and 108M active parameters, we measure how the optimizer changes the span of these biases, the largest minus the smallest in a layer. The span grows when all tokens favor the same experts, so it shows how hard a layer is to balance. At learning rates tuned on shorter runs, Muon raised the span of the first MoE layer at every seed, roughly doubling it on average, while the other layers changed much less. The rise came from Muon on the rest of the network, whichever optimizer trained the router. In those runs, the inputs to the first router shared a direction, and soon after the warmup its busiest expert averaged four to seven times its fair share. Over the first 3,000 steps, subtracting the batch mean from these inputs removed this overload while barely changing the loss. On the smaller model, halving Muon's learning rate removed most of the rise and slightly improved the loss. Moving only the router to AdamW made the loads shift more during the warmup and slightly raised the loss, so at this scale the results favor keeping the router on Muon.