Muon Makes the First Mixture-of-Experts Layer Harder to Balance
The first layer’s busiest expert received four to seven times its fair share of tokens shortly after warmup.
In models with 55 million and 108 million active parameters, Muon roughly doubled the first MoE layer’s routing-bias span on average, while other layers changed much less. The increase came from Muon training the rest of the network, regardless of which optimizer trained the router. Subtracting the batch mean from the first router’s inputs removed the overload over the first 3,000 steps with little change in loss. On the smaller model, halving Muon’s learning rate removed most of the span increase and slightly improved loss; switching only the router to AdamW slightly worsened loss.
HF Daily Papers' note
In models with 55 million and 108 million active parameters, Muon roughly doubled the first MoE layer’s routing-bias span on average, while other layers changed much less. The increase came from Muon training the rest of the network, regardless of which optimizer trained the router. Subtracting the batch mean from the first router’s inputs removed the overload over the first 3,000 steps with little change in loss. On the smaller model, halving Muon’s learning rate removed most of the span increase and slightly improved loss; switching only the router to AdamW slightly worsened loss.
HF Daily Papers' note
score 4