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

Muon Makes the First Mixture-of-Experts Layer Harder to Balance. By Marcos Hernanz, from the Autonomous University of Madrid. Published on the twenty-seventh of October.

Many open mixture-of-experts language models train with the Muon optimizer instead of AdamW. Most of them balance the load across experts with a bias on each expert's routing score, adjusted after every step. Their recipes disagree on one point: whether Muon should also train the router. The claim here sits to one side of that disagreement. Muon makes the first MoE layer harder to balance, and the rise comes from Muon on the rest of the network, not from whichever optimizer trains the router.

The measure is the span of the routing biases in a layer: the largest bias minus the smallest. The span grows when all tokens favour the same experts, so it reads as how hard a layer is to balance. The models have fifty-five million and one hundred and eight million active parameters, trained at learning rates tuned on shorter runs. Under those conditions, Muon raised the span of the first MoE layer at every seed, roughly doubling it on average. The other layers changed much less. And the rise came from Muon on the rest of the network, whichever optimizer trained the router.

Looking closer at those runs, the inputs to the first router shared a direction. Soon after the warmup, the busiest expert in that layer averaged four to seven times its fair share. Two interventions changed the picture. Over the first three thousand steps, subtracting the batch mean from those router inputs removed the overload while barely changing the loss. And on the smaller model, halving Muon's learning rate removed most of the rise and slightly improved the loss.

The scale is small: two models, the larger of them with one hundred and eight million active parameters. The effect is reported at learning rates tuned on shorter runs. The router question, too, is answered only at this scale. Moving just the router to AdamW made the loads shift more during the warmup and slightly raised the loss, so here the results favour keeping the router on Muon. Whether that holds for much larger models is outside what these experiments test.

The disagreement between recipes was framed around the router, but the router is not where the imbalance starts. The first layer receives inputs that share a direction, and the overload follows from that. Centring those inputs on their batch mean was enough to remove it, at least over the early steps examined.