Routers Learn the Geometry of Their Experts: Geometric Coupling in Sparse Mixture-of-Experts

Forum topic · 小凯 · 2026-05-14

Summary

This paper investigates how routing decisions form mechanistically in Sparse Mixture-of-Experts (SMoE) language models. The authors reveal a geometric coupling between routers and their corresponding experts: for a given token, the router weights of the selected expert and the expert weights processing it receive gradients along the same input direction, differing only by scalar coefficients, so matched router-expert directions accumulate the same routed token history. Experiments on a 1B-parameter SMoE trained from scratch confirm that higher router scores predict stronger expert neuron activation, showing routing decisions are mirrored inside selected experts. The study further shows that auxiliary load-balancing losses disrupt this coupled geometry by spreading input-directional gradients across router weights, making router directions nearly three times more similar to each other. Finally, a parameter-free online K-Means router—where each expert keeps a running mean of hidden states routed to it and tokens are assigned by cosine similarity—achieves the lowest load imbalance with only moderate perplexity increase, indicating geometric coupling captures most of what routers learn. arXiv: 2605.12476.

Paper Overview

Field: NLP Authors: Sagi Ahrac, Noya Hochwald, Mor Geva Published: 2026-05-12 arXiv: 2605.12476

Abstract

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, this work studies how routing decisions in SMoEs are formed mechanistically.

Key Findings

Geometric coupling: 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. Matched router-expert directions therefore accumulate the same routed token history.
Empirical confirmation: In a 1B SMoE trained from scratch, higher router scores predict stronger expert neuron activation, showing that routing decisions are mirrored inside the selected experts.
Impact of load-balancing losses: Auxiliary load-balancing losses disrupt this coupled structure by spreading input-directional gradients across router weights, making different router directions nearly three times more similar to each other.
Parameter-free routing: An online K-Means router, where each expert maintains a running mean of hidden states routed to it and tokens are assigned by cosine similarity, achieves the lowest load imbalance compared to auxiliary-loss and loss-free balancing, with only a moderate perplexity increase.

Conclusion

Geometric coupling appears central to effective routing and captures most of what routers learn, offering a mechanistic lens on SMoE routing dynamics and a direction for better router design.

---

*Auto-collected on 2026-05-14*