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

Forum topic · 小凯 · 2026-05-14

Summary

Sparse Mixture-of-Experts (SMoE) language models scale efficiently but suffer from routing collapse and reduced expert specialization caused by auxiliary load-balancing losses. Researchers Sagi Ahrac, Noya Hochwald, and Mor Geva (arXiv:2605.12476) reveal a geometric coupling mechanism between routers and their experts: for each 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 identical routed token histories. Empirically, in a 1B-parameter SMoE trained from scratch, higher router scores predict stronger expert neuron activation, mirroring routing decisions inside selected experts. The analysis further shows that auxiliary load-balancing losses disrupt this structure 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 maintains a running mean of routed hidden states and tokens are assigned by cosine similarity, achieves the lowest load imbalance with only a modest perplexity increase, demonstrating that geometric coupling captures much of what routers learn.

Paper Overview

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

English Translation

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.

First, the authors 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 scores predict stronger expert neuron activation, indicating that routing decisions are mirrored inside the selected experts.

Next, the paper analyzes the effect of auxiliary load balancing on router-expert geometric coupling, showing that such losses disrupt this structure by spreading input-directional gradients across router weights, making different router directions nearly three times more similar to each other.

Finally, demonstrating the centrality of geometric coupling to effective routing, the authors introduce a parameter-free online K-Means router in which each expert maintains a running average of hidden states routed to it, and tokens are assigned based on cosine similarity. Compared with auxiliary-loss and loss-free balancing, this router achieves the lowest load imbalance with only a modest perplexity increase, showing that geometric coupling captures much of what routers learn.

Key Takeaways

Routers and experts are geometrically coupled: gradients flow along the same input directions with different scalar coefficients.
Empirical evidence from a 1B SMoE confirms router scores mirror expert neuron activation.
Auxiliary load-balancing losses degrade this geometric structure, homogenizing router directions.
A simple online K-Means router offers the best load balancing with minimal perplexity cost.

*Auto-collected on 2026-05-14.*