Marcos Hernanz
7 min
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.
Alex: Let’s pin down the experiment’s scale before interpreting the result. How large were the models?
Sam: The smaller model had about fifty-five million active parameters, meaning weights used in processing a token. The larger had about one hundred eight million. Within each experiment, comparisons used matching initial weights and training batches, with three random seeds.
Alex: What happened to the balancing corrections when Muon replaced AdamW?
Sam: At learning rates selected in shorter tuning runs, the first expert layer’s average bias span roughly doubled. The increase appeared at every seed in the main comparison and recurred in later experiments. Other layers changed much less, while Muon achieved lower validation loss.
Alex: The suspected feedback weakened, but balancing still got harder. How did the separate optimizer choices locate the source?
Sam: The span mainly followed the backbone, meaning the tested network components other than the routers. A Muon-trained backbone produced the rise whether its routers used Muon or AdamW. Changing only the router had a much smaller effect on that first layer’s span.
Alex: That points upstream of the selector. What changed in the information arriving at it?
Sam: The first router’s input vectors increasingly pointed in a shared direction under Muon. Think of each token’s input as a common component plus its individual deviation. The common component can favor the same experts for every token, depending on how the router responds to that direction.
Alex: Then the balancer has to cancel a preference affecting the whole batch, not just occasional tokens. Can it respond quickly enough?
Sam: Not immediately, because its corrections move by a fixed amount each step. After the initial learning-rate warmup, the busiest expert averaged four to seven times its fair share in these runs. The biases needed hundreds of steps to catch up with the changing preference.
Alex: Shared directions and overload could still be fellow symptoms. Did the author intervene on the shared component itself?
Sam: He repeated the first three thousand training steps, subtracting the batch’s mean input before the first router saw it. The experts still received their original inputs. This centering removed the early overload on both model sizes and barely changed the loss.
Alex: That supports the proposed route from shared input to early overload. Does it explain every large bias span?
Sam: No. Some experiments put only selected backbone components on Muon. They produced large spans without reproducing the strong shared direction. The author also doesn’t know why the whole Muon-trained backbone creates that direction, or fully understand what spreads the router’s responses along it.
Alex: Before adding a new routing operation, I’d want to know whether ordinary tuning solves most of the problem. Did lowering the learning rate help?
Sam: On the smaller model, halving Muon’s selected learning rate removed about three quarters of the excess span. It also slightly improved the final loss. The shared direction remained, but the router’s responses along it were less spread out.
Alex: So this is partly a warning about tuning on shorter runs. The selected rate wasn’t the best of the tested alternatives over the longer training period.
Sam: That qualification matters. Faster bias updates also reduced overload, while an added balance loss did not eliminate the span difference. And on the larger model, the first expert layer remained overloaded at training’s end.
Alex: What about the original recipe dispute? If the backbone uses Muon, does moving the router to AdamW help?
Sam: At these scales, the comparisons favored keeping the router on Muon. Moving it to AdamW made workloads shift more during warmup and slightly worsened validation loss. In the larger-model experiment, loss spikes appeared only with that mixed optimizer setup.
Alex: How much should researchers trust the practical implications? These aren’t production-sized systems.
Sam: The within-setup evidence is persuasive, but generalization is limited. Each comparison has few seeds, and the report makes no correction for multiple comparisons. It tests one architecture and dataset, with models far smaller and training much shorter than the production recipes motivating the question.
Alex: And did the overloaded experts actually make these experiments slower or less accurate?
Sam: The report found no measurable cost from overload, apart from the possible connection to loss spikes. These models had no capacity limit. With experts on separate devices, overload can slow training; with capacity limits, excess tokens can be dropped. Centering also remains a short-run test, not a validated training recipe.
Alex: Given those boundaries, who should read the full report, and where should they start?
Sam: Researchers studying expert routing or optimizer choices should read The Router and the Backbone, then The Shared Direction. Those sections separate the source from the symptom. Follow with The Learning Rate of Muon and the limitations before changing a recipe. For everyone else: balanced workloads can hide a router that strongly prefers the same experts.
Alex: Watch the corrections, not just the balance.