Muon optimizer doubles first MoE layer bias span
Original title: Muon Makes the First Mixture-of-Experts Layer Harder to Balance
Muon raises router bias span in early MoE layers, complicating load balancing at tuned learning rates.
Mixture-of-experts models train with Muon instead of AdamW. Bias span on router scores measures balancing difficulty. Muon on the network doubles span of the first layer at every seed, while later layers change little. Inputs to first router share direction, leading to one expert averaging four to seven times fair share. Subtracting batch mean removes overload without loss change. Results favor keeping router on Muon for smaller models.
- span grows when tokens favor same experts
- doubles first layer span at every seed
- largest rise from Muon on rest of network
- inputs share direction post-warmup
- halving Muon LR removes rise
- Organizations
- Hugging Face
- AI models
- mixture-of-experts language models
From the original article
1 min read at huggingface.co
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.
By Marcos Hernanz · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.