Muon doubles first-layer MoE routing bias, paper finds

The first mixture-of-experts layer is where Muon bends. On models with 55M and 108M active parameters, the optimizer 'raised the span of the first MoE layer at every seed, roughly doubling it on average, while the other layers changed much less,' writes Marcos Hernanz in a paper posted to arXiv.

The span is the largest minus the smallest routing-bias value in a layer, and it grows when all tokens favor the same experts. Soon after warmup, the paper reports, the busiest expert in the first layer 'averaged four to seven times its fair share.' The rise came from Muon on the rest of the network, whichever optimizer trained the router.

Two patches worked. Over the first 3,000 steps, subtracting the batch mean from the first router's inputs removed the overload 'while barely changing the loss.' On the smaller model, halving Muon's learning rate also removed most of the rise and slightly improved the loss.

Moving only the router to AdamW while keeping Muon elsewhere went the other way: it made loads shift more during warmup and slightly raised the loss. Hernanz concludes that 'at this scale the results favor keeping the router on Muon.' The experiments cover two small-to-mid models and a 3,000-step window; the abstract reports no results above 108M active parameters.