Paper Overview
Research area: NLP Authors: Sagi Ahrac, Noya Hochwald, Mor Geva Published: 2026-05-12 arXiv: 2605.12476
Summary
Sparse Mixture-of-Experts (SMoE) models enable efficient scaling of language models, but training them remains challenging: routing can collapse onto a few experts, and auxiliary load-balancing losses can reduce specialization. Motivated by these hurdles, the authors study how routing decisions in SMoEs are formed mechanistically.
Key Contributions
- Geometric coupling between routers and experts: For a given token, the router weights for the selected expert and the expert weights processing it receive gradients along the same input direction, differing only in scalar coefficients. Consequently, matched router–expert directions accumulate the same routed token history.
- Empirical evidence: In a 1B-parameter SMoE trained from scratch, higher router scores predict stronger activation of expert neurons, showing that routing decisions are mirrored inside the selected experts.
- Effect of auxiliary load-balancing losses: These losses disrupt the router–expert geometric coupling by diffusing input-directional gradients across router weights, making different router directions nearly three times more similar to one another.
- Parameter-free online K-Means router: Each expert maintains a running mean of hidden states routed to it; tokens are assigned by cosine similarity. Compared to auxiliary-loss and loss-free balancing, this router achieves the lowest load imbalance with only a modest perplexity increase.
Conclusion
Geometric coupling is central to effective routing and captures much of what routers learn, suggesting balance mechanisms should be designed to preserve rather than destroy this structure.
Original Abstract (excerpt)
Sparse Mixture-of-Experts (SMoE) models enable scaling language models efficiently, but training them remains challenging, as routing can collapse onto few experts and auxiliary load-balancing losses can reduce specialization. Motivated by these hurdles, we study how routing decisions in SMoEs are formed mechanistically. First, we reveal a geometric coupling between routers and their corresponding experts. For a given token, the router weights for the selected expert and the expert weights processing it receive gradients along the same input direction, differing only in scalar coefficients. Thus, matched router–expert directions accumulate the same routed token history. This theoretical coupling also appears empirically in routing dynamics. In a 1B SMoE trained from scratch, higher router...
---
*Auto-collected on 2026-05-14*