Standard load-balancing losses degrade SMoE expert specialization by 3x

Researchers at Tel Aviv University identified a geometric coupling between routers and experts in Sparse Mixture-of-Experts models that explains why standard load-balancing losses degrade specialization. In experiments on an 11B-parameter SMoE trained on 5,050 billion tokens, auxiliary losses present in Mixtral, Switch Transformers, and DeepSeek-V3 made router weight vectors nearly three times more similar to each other than training without the loss.

The mechanism operates at the gradient level. When a token is routed to an expert, both the router's weight vector for that expert and the expert's input-side weights receive updates along the same input direction, differing only in scalar coefficients. The chain rule in an SMoE layer enforces this proportional form: matched router–expert direction pairs co-evolve as coupled accumulators of the token histories routed through them.

Across the 11B model, experts ranked higher by the router produced stronger neuron activations for the same tokens than experts the router did not select. Routing decisions embed themselves in the expert's internal computation — a geometric signature visible at inference time.

Load-balancing losses work by sending input-directed gradient signals to every router weight vector on every token, regardless of which expert was chosen. That broadcast collapses the directional fingerprints that coupling builds up. The researchers compared two 11B SMoEs trained with and without the auxiliary loss: loss-free training preserved expert differentiation at the cost of worse load balance.

For teams training or fine-tuning SMoE models—a cohort that expanded sharply after DeepSeek-V3 and OLMoE showed the architecture matches dense models at inference cost savings—this reveals the gradient-level cost of balancing. The researchers introduce a parameter-free K-Means router: each expert maintains a running average of routed hidden states; new tokens are assigned by cosine similarity to those centroids. The K-Means variant achieves the lowest load imbalance of three configurations, with a modest perplexity increase.

The experiments are bounded to the 11B range; whether coupling dynamics persist at 100B-plus scales remains untested. The K-Means router's perplexity tradeoff is not quantified against downstream task performance. Yet the diagnostic value is concrete: teams observing routing collapse now have a specific mechanism to interrogate and a clear prediction—check gradient similarity of router directions before tuning the balancing coefficient.