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

TL;DR

This paper uncovers a geometric coupling mechanism in sparse mixture-of-experts models, explaining how routing decisions align with expert activations and how auxiliary losses affect this structure, leading to improved routing strategies.

Contribution

It reveals the geometric coupling between routers and experts, analyzes the impact of load balancing losses, and introduces a parameter-free online K-Means router for better routing.

Findings

Router and expert weights share gradient directions, forming a geometric coupling.

Auxiliary load balancing disrupts this coupling, increasing similarity among router directions.

A parameter-free K-Means router achieves low load imbalance with minimal perplexity increase.

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, 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 $1$B SMoE trained from scratch, higher router scores predict stronger expert neuron activations, showing that routing decisions are mirrored inside the selected expert. Next, we analyze the effects of auxiliary load balancing on the router--expert geometric coupling, showing that such losses break this structure by spreading input-directed gradients across router weights, making distinct router directions nearly three times more similar to each other. Last, we demonstrate the centrality of geometric coupling for effective routing with a parameter-free online K-Means router, in which each expert maintains a running average of the 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, indicating that geometric coupling captures a substantial part of what the router learns. Overall, our results explain how routers form assignment geometry that supports an effective division of labor.