Muon Makes the First Mixture-of-Experts Layer Harder to Balance
Abstract
Using Muon to train mixture-of-experts language models increases routing imbalance in early layers, which can be mitigated by centering router inputs or reducing the optimizer learning rate.
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.
Community
Open MoE recipes disagree on whether Muon should train the router. DeepSeek trains it with Muon, Qwen keeps it on AdamW.
We found the router's optimizer isn't the problem. With Muon on the rest of the network, the inputs to the router of the first MoE layer share a direction, and right after warmup its busiest expert gets 4-7x its fair share, vs 1.1x under AdamW. Moving only the router to AdamW makes it worse. Subtracting the batch mean from the router's inputs removes the overload and barely changes the loss.
Code and all 145 runs are public.
This is an automated message from the Librarian Bot. I found the following papers similar to this paper.
The following papers were recommended by the Semantic Scholar API
- ExpertMuon-Compass: Alignment-Guided Step Sizes for Mixture-of-Experts Training (2026)
- How to Loop MoE: Flatten the Experts, Untie the Attention (2026)
- HobbyLM: A Sparse Mixture-of-Experts Model Family on a Hobby Budget (2026)
- EnergyCal: Taming Loss Spikes in Muon-Family Optimizers for Sparse MoE Training (2026)
- How Sparse Probability Maps Shape Mixture-of-Experts Routing (2026)
- The Router Already Knows (2026)
- Universal Test-Time Training (2026)
Please give a thumbs up to this comment if you found it helpful!
If you want recommendations for any Paper on Hugging Face checkout this Space
You can directly ask Librarian Bot for paper recommendations by tagging it in a comment: @librarian-bot recommend
Get this paper in your agent:
hf papers read 2610.37914 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 2
marcoshernanz/llm-lab-fineweb-edu-sample10bt-bpe-16384-full
Spaces citing this paper 0
No Space linking this paper