Routers Learn the Geometry of Their Experts: Geometric Coupling in Sparse Mixture-of-Experts
Sagi Ahrac*
Noya Hochwald*
Mor Geva
arXiv:2605.12476v1 [cs.LG] 12 May 2026
Blavatnik School of Computer Science and AI, Tel Aviv University {sagiahrac@mail, noyahochwald@mail, morgeva@tauex}.tau.ac.il
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 1B 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.
1
Introduction
Sparse Mixture-of-Experts (SMoE) has emerged as a leading architecture for scaling language model parameters without a proportional increase in inference latency [1–3]. Recent implementations, such as DeepSeek-V3 [4] and OLMoE [5], match or outperform dense models while activating only a fraction of their total parameters. These efficiency gains stem from a division of labor in which a router, typically implemented as a small gating network, directs each input to a subset of independent expert networks [1, 2]. Still, despite their wide adoption, training SMoEs with effective routing remains a challenge. Without intervention, routing concentrates on a shrinking subset of experts, leading to representation collapse [6]. Approaches to mitigate this apply auxiliary load-balancing losses [1, 3, 2], which encourage balanced routing but often reduce expert specialization [7, 4]. Motivated by these pathologies, we seek to understand the inner dynamics of routing in SMoEs, focusing on how routing decisions are derived mechanistically. We tackle this question from a geometric view. First, we analyze the gradients that shape both sides of the routing decision. Although routers and experts are often treated as separate modules [1, 8], SMoE gradients reveal a shared input-directed structure. For each assigned token, the router weights Preprint.
Figure 1: Router–expert geometric coupling in SMoEs. The router scores a hidden state x and selects a sparse set of top-K experts. For each selected expert, the matched router direction and expert input-side weights receive backpropagation updates proportional to the same hidden-state direction x. Repeated updates make matched router–expert pairs accumulate a common routed-token history, which is read out at inference as higher router scores predicting stronger gate-neuron activations.
associated with the selected expert and the expert weights that process the token receive updates along the same input direction, differing only in scalar coefficients. Notably, this alignment is not a generic consequence of joint training: a shared computation graph derives both modules to receive gradients, but does not imply that their updates accumulate along shared input directions. In SMoE layers, however, the chain rule enforces this proportional form, inducing a geometric coupling where matched router–expert directions evolve as coupled accumulators of their shared routed token history. Next, we investigate if this theoretical coupling translates into empirical routing dynamics. To this end, we compare the router’s score for an expert and that expert’s neuron activations in response to the same token. In a 1B SMoE trained from scratch for approximately 50B tokens, we find that experts ranked higher by the router consistently exhibit stronger activations than experts not selected by the router. This shows that routing decisions are not merely an external assignment, but are mirrored in the computation of the selected expert, thus providing empirical evidence of the geometric coupling between the router and experts. Having established the router–expert geometric coupling, we revisit common routing instabilities and load-balancing side effects. Specifically, we study how the router’s geometry is influenced by the common auxiliary load-balancing loss employed in existing models [1, 3, 2]. This loss encourages balanced routing by penalizing uneven expert load. From a theoretical point of view, this optimization sends input-directed gradients to every router weight vector on every token, regardless of which experts were chosen. This is expected to unify different router directions over training, weakening the geometric coupling between the router and the experts. We evaluate this prediction by training two 1B SMoEs that differ only in the balancing rule, finding that distinct router weight vectors become nearly three times more similar with the auxiliary loss than without it. These results show that auxiliary balancing regulates expert usage, yet breaks the geometric coupling, unifying expert-specific directions and eroding the specialization that coupling produces. Last, we demonstrate the centrality of geometric coupling in routing by evaluating a parameter-free centroid router that follows this natural coupling. Concretely, if router weights learn to summarize 2
the hidden states assigned to each expert, then explicit summaries may recover much of their role. Following this idea, we implement the centroid router using online K-Means [9]. Each expert maintains a centroid of its assigned hidden states, and tokens are routed by cosine similarity to these centroids. On our 1B SMoE setup, this router achieves the lowest load imbalance among all variants while maintaining comparable perplexity to the learned loss-free router. This indicates that geometric coupling is not only a byproduct of training, but a substantial component of what routers learn. In conclusion, our work makes the following contributions: 1. We reveal a geometric coupling between routers and experts in SMoEs, where matched router–expert pairs receive gradient updates along the same routed-token directions and evolve as coupled accumulators of their shared token history. 2. We show empirically that router scores are reflected in the selected expert’s internal computation, with higher-ranked experts exhibiting stronger activations for the same tokens. 3. We show that common auxiliary load-balancing losses act directly on router geometry, injecting input-directed gradients into every router weight vector regardless of selection and making distinct router directions nearly three times more similar. 4. Motivated by coupling, we instantiate a simple online K-Means router in which non-learnable EMA centroids replace learned router weights, and tokens are routed by cosine similarity with sign-updated biases. This router achieves the lowest load imbalance among our variants with only a modest perplexity increase. Taken together, our results answer how routers form assignment geometry that supports an effective division of labor. More broadly, they suggest that future routing methods may benefit from preserving the natural router–expert geometry that emerges during training. We release our code at https: //github.com/sagearc/router-expert-geometry.
2
SMoE architecture and training
The SMoE architecture SMoEs are typically built on top of the Transformer architecture [10], replacing feed-forward blocks with a set of N distinct expert networks {E1 , . . . , EN } and a router R that selects a small subset of experts for each input [1, 3, 2, 11]. For a hidden state x ∈ Rd , the router computes one score per expert: p = σ(z + m),
z = Wr x.
(1)
where Wr ∈ RN ×d are the router weights, with the i-th row ri corresponding to the expert i, m ∈ RN is a mask for selecting the top-K entries in z, and σ is a nonlinear activation function. Throughout the paper, we refer to zi and pi as the router’s score and the routing weight of expert i, respectively, and denote the set of top-K selected experts by TK . The SMoE layer combines only the selected expert outputs: X y= pi Ei (x). (2) i∈TK
Each expert Ei maps the hidden state through an intermediate expert dimension. In the gated SwiGLU experts used by recent SMoEs [12, 5, 4], this computation can be written as Ei (x) = Widown σ(Wigate x) ⊙ Wiup x , (3) where σ is the SiLU activation and Wigate , Wiup ∈ Rdi ×d are input-side expert matrices with an intermediate dimension di . Throughout the paper, we refer to the coordinates of σ(Wigate x) as the expert’s gate-neuron activations; these are the activations measured in our empirical analysis. SMoE training and load balancing Training SMoEs requires both meaningful token-to-expert assignments and balanced expert utilization. For a batch of training examples B, let fi denote the fraction of top-K assignments received by expert i: 1 X 1 fi = 1{i ∈ TK }, τ= . (4) K|B| N x∈B
3
Here TK is the set of top-K expert indices for x, and τ is the target uniform utilization. Standard SMoE training often adds an auxiliary load-balancing loss [1, 3, 2], L = LLM + λaux Lbalance ,
(5)
where Lbalance penalizes the unequal distribution of tokens across experts, and λaux dictates the strength of this penalty. In contrast, recent loss-free methods instead add adaptive per-expert routing biases [7, 4], z̃i = zi + bi , bi ← bi + γ sign(τ − fi ), (6) where γ > 0 is the bias update rate. In our experiments, this bias affects top-K selection but is not added to the expert weights used in Eq. (2).
3
Router–expert geometric coupling
We formalize the gradient structure underlying the learned routing, showing that for a token representation x routed to expert i, both the corresponding router weight vector ri and the input-side weights of expert i receive updates proportional to x. Over training, matched router–expert pairs accumulate the same routed token history with different weights, a structure that follows from the chain rule rather than from an added constraint. To derive the coupling between router weight vectors and expert input matrices, we consider a standard backpropagation step through the SMoE layer in Eq. (2). On the expert side, each row of Wigate takes an inner product with x to produce one coordinate of the gate activations. Let wi,k be k-th row of Wigate , its respective gradient can be written as ∇wi,k L = δi,k x⊤ ∝ x⊤ ,
(7)
where δi,k collects all scalar factors other than the input direction. Therefore, the rows of Wigate evolve into weighted sums of the hidden states processed by expert i. Notably, the same input-directed form also applies to Wiup , with a different scalar coefficient. For a detailed mathematical derivation, see Appendix A. Simultaneously, the router weight vector ri follows the same input-directed update structure. Given the logit zi = x⊤ ri , the gradient with respect to ri is ∂L ∂pi ∇r i L = x = γi x ∝ x, (8) ∂pi ∂zi where γi is the scalar coefficient for the router update (see Appendix A.3 for the full router-gradient derivation). For both the router and the expert input-side matrices, weights are updated by adding the hidden-state direction x when it reduces the loss and subtracting x when it increases it. Thus, over training, matched router–expert pairs accumulate the same routed token history with different weights, making this coupling a direct consequence of backpropagation through the SMoE layer rather than an added constraint. This shared gradient structure suggests that the paired router and expert weights should develop aligned geometry in hidden space (Figure 1). This accumulation is expert-specific, as experts outside the top-K set do not contribute to the layer’s output. For such an expert j ∈ / TK , the routing weight pj is zero, and its router weight vector rj receives no gradient from this token. Each router weight vector therefore accumulates only the tokens routed to its expert, yielding the expert-specific accumulation that drives geometric coupling. This view implies that router preference should be reflected inside the selected expert. Specifically, if ri and the rows of Wigate are shaped by the same routed hidden states, then a token that receives a high router score for expert i should also activate that expert’s gate neurons more strongly. Section 4 tests this relationship directly.
4
Empirical evidence of geometric coupling
The derivation in Section 3 shows that router scores and expert activations share a common geometry because both are shaped by the same routed hidden states. Here, we examine this relationship empirically, asking whether router preference is reflected inside the selected expert’s computation. 4
For a token x, the router score zi = r⊤ i x measures the router’s preference for expert i. For the same token, the expert’s gate-neuron activations are the coordinates of σ(Wigate x). Importantly, these quantities are computed independently in the same forward pass, i.e., the router score is not used to compute the expert activations. If ri and the expert input-side vectors are shaped by the same routed inputs during training, then higher router scores should correspond to stronger expert activations. We measure expert activations by averaging over the expert’s gate neurons. For our experiment, we use the 1B SMoE configuration from Wang et al. [7], and train it from scratch on the OLMoE-mix-0924 [5] dataset for ∼50B tokens. Each model has L=9 SMoE layers, hidden size d=1024, N =64 routed experts, topK=6 routing, Ns =2 shared experts, and expert hidden width 512. Specifically, we use the model trained with the loss-free balancing rule of Wang et al. [7], which raises the routing bias of underused experts and lowers it for overused experts, so no balancing gradient reaches the router and perexpert input accumulation is preserved. This is the regime in which the geometric coupling derived in Section 3 should be most directly observable. We evaluate the model on token sequences from English Wikipedia. For each routed token–expert pair and each SMoE layer ℓ, we record the raw (ℓ) router score zi (x(ℓ) ) = r⊤ and, independently i x inside the selected expert, the average activation of its gate neurons for the same token. We normalize router scores and expert activations separately within each layer and expert, then pool all routed pairs.
Figure 2: Expert activations increase with router score. For each routed token–expert pair, we compare the router score with the average activation of that expert’s gate neurons. Scores and activations are normalized separately for each layer and expert before pooling. We observe that router scores and expert activations are correlated (ρ = 0.43, p-value 1.2×10−81 ).
Router scores predict expert activations Figure 2 shows a monotone relationship between router score and expert activation. Tokens with higher router scores produce stronger expert activations, as predicted by the router–expert coupling. During the same forward pass, the router score comes from ri , while expert activations are measured from σ(Wigate x). This monotone trend therefore gives direct functional evidence of geometric coupling. Inputs that shaped ri during training also shaped the input vectors of Wigate , so a new token aligned with that history receives a higher router score and produces stronger expert activations. This complements the static observation of Lo et al. [13] by showing that router–expert correlation is not only visible in trained weights, but also appears in token-level expert activations during routing.
5
Auxiliary load-balancing breaks geometric coupling
The previous section showed that router preferences are reflected inside experts when coupling is preserved. We now examine what happens to the geometric coupling under the auxiliary loadbalancing loss, a widely used practice for avoiding load-imbalance across experts. The auxiliary loss breaks expert-specific accumulation In Section 3 we showed that, under the language-modeling loss, only selected experts contribute to the SMoE layer output. For an unselected expert j ∈ / TK , the routing weight pj is zero and rj receives no gradients from this token. This preserves expert-specific accumulation. Conversely, we observe that the standard auxiliary load-balancing loss [1, 2] breaks this expert-specific structure: Lbalance = N
N X
fi Pi ,
i=1
Pi =
1 X pi , |B|
(9)
x∈B
where Pi is computed from an unmasked softmax over all N experts, i.e., softmax(z), and therefore depends on every router weight vector, not only on the selected experts. The chain rule yields, for 5
Figure 3: Auxiliary loss collapses router geometry. Each panel shows pairwise cosine similarities between router weight vectors within one model layer. The top row uses auxiliary load-balancing loss. The bottom row uses bias-only balancing with the same 1B architecture and training setup. Off-diagonal means (µ) appear below each panel. every token x in the batch and every expert j, ∇rj Lbalance = βj x,
βj ̸= 0
(10)
where βj is a scalar coefficient that depends on the auxiliary loss and routing probabilities. Every router weight vector therefore absorbs an input-directed contribution from every token, regardless of which experts were chosen. Instead of accumulating only the tokens routed to expert j, rj also accumulates signal from the global token stream. Notably, these are the “interference gradients” of Wang et al. [7]. Since hidden states at a given layer are not centered across the token stream, their average over many tokens can be nonzero. Because the auxiliary-loss gradient on each router weight vector is proportional to x on every token, all router weight vectors can accumulate a shared mean component, and across training this shared component can pull the vectors together, giving a coupling-failure signature which we now measure directly. The auxiliary loss collapses router weight vectors We train two 1B SMoEs from scratch using the same setup as in Section 4, starting from the same initialization and training for the same number of tokens. The models use the same SMoE architecture and differ only in the balancing rule. The run without auxiliary loss uses the bias-only balancing rule [7], where no balancing gradients reach the router. The run with auxiliary loss replaces the bias rule with the auxiliary load-balancing loss used by Switch Transformers [2] and the router z-loss from ST-MoE [14], a recipe used in recent SMoEs such as OLMoE and Mixtral [5, 12]. After training, we extract the router weight matrix Wr ∈ RN ×d at three stratified layers and report pairwise cosine similarities between its router weight vectors. Figure 3 presents the results, showing a collapse in router geometry. With the auxiliary loss, the off-diagonal mean cosine similarity is µ = 0.63, 0.63, and 0.57 at layers 0, 4, and 8, respectively. By contrast, under loss-free bias balancing the same architecture remains more diverse, with substantially lower similarity scores of µ = 0.32, 0.18, and 0.13. Overall, this showcases that auxiliary loadbalancing aligns otherwise distinct router weight vectors, making them nearly three times more similar. Section 4 showed that router scores carry information when coupling is intact. Here we identify a mechanism that erodes this information at the router itself. When router weight vectors become strongly aligned, the expert-specific differences between router scores shrink relative to their shared 6
component. The softmax therefore has less expert-specific signal to use for distinguishing between experts. Prior work has shown that interference gradients from auxiliary loss can impair language modeling [7] and degrade expert specialization [15]. Our measurement provides a geometric explanation for this degradation, localizing the effect directly in the router’s own weight vectors. This motivates the design we develop next, a router whose directions come directly from the inputs each expert has processed, without balancing gradients.
6
Centroid tracking captures a substantial component of learned router–expert coupling
The geometric coupling suggests that router weight vectors accumulate toward summaries of the inputs routed to each expert. This turns the gradient dynamics into a form of online centroid tracking, where each router direction is repeatedly pulled toward the hidden states assigned to its expert and gradually becomes a running summary of that expert’s routed-token cluster. If routing is largely centroid tracking, then much of the router’s role should be recoverable without trainable routing weights or router-side gradients. We test this by replacing the standard gating weights with non-learnable centroids updated as exponential moving averages of the hidden states assigned to each expert. The resulting router has no trainable routing parameters and it absorbs no balancing gradients. In what follows, we describe our router in detail, and evaluate its performance compared to existing routing approaches. Centroid update rule Our router adapts online clustering to SMoE routing. At every layer, we maintain one centroid per routed expert, corresponding to the router weight vectors in standard SMoEs. Here, as before, the total number of centroids and experts is N , and the number of experts selected per token is k. For a token with hidden state x, expert i receives the score si (x) =
c⊤ i x + bi , ∥ci ∥ ∥x∥
(11)
where ci ∈ Rd is the router’s centroid corresponding to expert i and bi is a scalar bias. The token is routed to the k experts with the largest scores. Centroids are updated at every step by an exponential moving average over the inputs assigned to each expert, P ci ← α ci + (1 − α) x̄i , x̄i = |T1i | x∈Ti x, (12) where Ti is the set of tokens routed to expert i in the current micro-batch. We update the biases using the loss-free rule of Wang et al. [7], bi ← bi + γ sign(τ − fi ), where fi is the realized load fraction, τ = k/N is the target load, and γ > 0 is the update rate. Both centroids and biases are gradient-free running statistics, rather than trainable parameters. Thus, the router makes the predicted centroid-tracking dynamics explicit, with no trainable router weights, no auxiliary loss, and no gradient flow into the routing decision. Experimental setup To evaluate our new router, we use the same 1B SMoE configuration as in Section 4, and train it from scratch on OLMoE-mix-0924 [5] multiple times with different routing rules. Each run trains for 21k steps (50B tokens), using AdamW with β1 =0.9, β2 =0.95, a peak learning rate of 10−3 decayed to 10−4 , a 1k-step warmup, and a batch size of 2.36M tokens per step. We compare four routing variants: • Aux-Loss: Uses the Switch-style auxiliary load-balancing loss together with the router z-loss regularization [2, 14]. This is the same model analyzed in Section 5. • Loss-Free: Uses the bias-only balancing rule of Wang et al. [7], adjusting per-expert routing biases while leaving router weights free of balancing gradients. • Loss-Free + Seq-Aux: In addition to the bias-based balancing of Loss-Free, it employs the sequence-wise auxiliary loss of DeepSeek-V3 [4], which encourages each sequence to distribute its routed tokens more evenly across experts. This discourages within-sequence routing bottlenecks. • K-Means: Our router, which follows the natural geometric coupling, using the centroid rule above with α=0.99 and γ=10−3 . These hyperparameters were chosen so the centroid and bias updates are similar in magnitude to the other three variants. 7
Table 1: K-Means achieves comparable perplexity to the loss-free family and yields the lowest steady-state load imbalance. All four runs use the 1B SMoE configuration from Wang et al. [7] and are trained for 50B tokens; they differ only in the routing rule. “Router params” counts the learnable parameters of the router, summed over the nine layers; the centroids and biases used by K-Means are non-learnable running statistics. MaxVio is averaged across the nine layers at step 21k. Method
Router parameters
Auxiliary losses
0.59M 0.59M 0.59M 0
load balance + z seq-aux – –
Aux-Loss Loss-Free + Seq-Aux Loss-Free K-Means (ours)
Train PPL C4-en PPL Pile PPL MaxVio 15.09 15.03 15.01 15.40
Full range
MaxVio = LI 1
0 6 5 4 3 2 1 0
10
Training tokens (B)
20
30
20.54 20.44 20.40 21.01
11.82 11.77 11.76 12.09
0.526 0.102 0.084 0.037
Loss-free family (zoom) 40
0.30
Aux-Loss Loss-Free + Seq-Aux Loss-Free K-Means (ours)
0
10
Training tokens (B)
20
30
40
0.25 0.20 0.15 0.10 0.05
0
2500
5000
0.00
7500 10000 12500 15000 17500 20000
Training step
0
2500
5000
7500 10000 12500 15000 17500 20000
Training step
Figure 4: Load imbalance during training (log scale). Layer-averaged MaxVio versus training step for the four routing variants on the 1B SMoE of Wang et al. [7]. Curves are 200-step rolling means; the logarithmic y-axis separates the Aux-Loss collapse from the loss-free family while keeping both visible. K-Means settles to the lowest plateau (MaxVio ≈ 0.037 at step 21k) without any learned routing parameters or balancing gradient.
We evaluate these variants in terms of language modeling performance and load imbalance. For language modeling, we report perplexity (PPL) over the C4-en and the Pile held-out validation slices from OLMo, derived from the C4 [16] and Pile corpora [17]. For load-balancing, we use the MaxVio metric by Wang et al. [7], defined as MaxVio = maxi fi /f¯ − 1 and averaged across the nine layers. MaxVio measures the worst relative overload: a value of 0 means perfectly balanced routing, while larger values indicate that at least one expert receives more tokens than the average expert. Results Table 1 reports perplexity and steady-state load imbalance of the trained models (using their last checkpoints). K-Means achieves the lowest sustained load imbalance across variants, with MaxVio = 0.037 compared to 0.084 for Loss-Free and 0.526 for Aux-Loss, despite using no trainable routing parameters or balancing gradients. This improvement comes with a modest perplexity cost relative to the learned loss-free router, with training PPL increasing by 2.6% (15.40 → 15.01) and similar gaps on C4-en and Pile. Figure 4 further shows the MaxVio score over training. The loss-free variants rapidly reduce imbalance and remain stable, while K-Means settles to the lowest and smoothest plateau. Since its routing state consists only of exponential moving average centroids and load-balancing biases, the bias rule regulates expert usage without altering centroid directions through gradients. Discussion Sections 3–5 suggest that router weight vectors accumulate toward summaries of the inputs processed by their corresponding experts, while auxiliary balancing losses perturb this expert-specific accumulation. Under this view, a router that explicitly tracks these centroids, without trainable routing weights or balancing gradients, should train stably and balance at least as well as the bias-only baseline. Our results support this prediction, as K-Means achieves tighter balance than Loss-Free with a 2.6% training perplexity cost. We interpret the perplexity gap as the cost of removing routing degrees of freedom beyond centroid tracking, and its modest size as evidence that centroid tracking captures a substantial component of what gradient-trained routers learn. 8
7
Related work
Geometric analyses of SMoE routing and experts Recent work analyzes SMoEs along several dimensions, including expert specialization, knowledge attribution, and routing behavior [15, 18–20]. Closest to ours are geometric analyses of expert and routing structure. Liu et al. [21] promote expert diversity through an orthogonal optimizer, and Lo et al. [13] observe correlated router and expert gate-projection weights in pretrained models. Concurrent with our work, Huang et al. [22] document aligned expert subspaces from shared low-rank input structure. We complement these empirical analyses with a gradient-level account. Dikkala et al. [23] study router–expert correspondence under clustered-data assumptions, and Lv et al. [24] enforce it through an explicit alignment loss; we show that this correspondence instead emerges under standard SMoE training from shared input-directed gradients between matched router rows and expert input weights, and is visible in expert activations. Decoupled and constrained routing Many routing methods intervene on the router separately from expert weight evolution, either through auxiliary load-balancing losses [1, 3, 2] or bias-based balancing [7, 4]. Other works decouple routing more directly: Dai et al. [8] freeze a distilled router during later training, Sukhbaatar et al. [25] train experts independently before merging them into an MoE, Pan et al. [26] defer sparse routing to inference, and Xu et al. [27] distill fixed routing structures from pretrained models. Our analysis suggests that such interventions should be evaluated not only by load balance, but also by whether they preserve router–expert geometry: auxiliary balancing can inject interference gradients into router directions, while stronger decoupling may discard the shared input-directed updates through which router rows and expert weights co-evolve. Geometric and centroid-based routing Recent works replace learned MoE routing with explicit geometric structure, using shared eigenbases, per-expert subspaces, Grassmannian structure, or hidden-state subspaces to route tokens [28–31]. Most similar to our K-Means router, Yang [32] route by cosine similarity to EMA-updated prototypes, though theirs are learned in a regularized latent projection space rather than computed as non-learnable averages of routed hidden states. Centroidbased routing has also appeared in different settings, including online spherical k-means for sparse attention [33] and EMA centroid routing for frozen-backbone adapters [34], neither targeting SMoE feed-forward routing during language-model pretraining. These works impose geometric structure on routing through architectural design. We show, in contrast, that such structure already emerges in standard learned SMoE routers, since matched router vectors and expert input weights co-accumulate the hidden-state directions assigned to their experts. To test this mechanism directly, our K-Means router replaces learned router weights with non-learnable EMA centroids of routed hidden states.
8
Conclusion and discussion
We show that learned routing in SMoEs is not independent of expert evolution but emerges from gradient dynamics shared with the experts it selects. Routers and the experts they select co-evolve as coupled accumulators of the tokens routed to each expert. This theory-driven geometric coupling also appears empirically in the model’s computation, is fragile under interventions that bypass it, and largely recoverable from the routed token stream alone. Together, our findings suggest that router–expert coupling is a substantial part of what gradient-trained routers learn, and that SMoE training should preserve this geometry rather than perturb it. Limitations and future work Our empirical results are based on a single 1B SMoE configuration, and validating router–expert coupling across larger scales and architectures remains an important next step. At the activation level, our evidence focuses on the gate branch, SiLU(Wigate x); extending this readout analysis to Wiup and other expert weights is left for future work. Beyond scale, our analysis of how auxiliary load-balancing gradients reduce separation between router directions captures only one consequence of router–expert geometry; SMoE training involves other pathologies and routing interventions, including representation collapse, expert dominance, and alternative balancing mechanisms. A broader direction is to test whether the same geometric lens can help explain these phenomena by measuring how they preserve, distort, or bypass router–expert coupling. Last, the K-Means router is a constructive test of the coupling view rather than a practical router; its perplexity gap suggests that learned routers use degrees of freedom beyond centroid position, motivating hybrid designs that preserve the centroid geometry while closing this gap. 9
Acknowledgments We thank Elad Shikley for his help with figure design, and Or Shafran for feedback on the manuscript.
References [1] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017. URL https://arxiv.org/abs/1701.06538. [2] William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022. URL https://arxiv.org/abs/2101.03961. [3] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding, 2020. URL https://arxiv.org/abs/2006.16668. [4] DeepSeek-AI, Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Daya Guo, Dejian Yang, Deli Chen, Dongjie Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, Guowei Li, H. Zhang, Han Bao, Hanwei Xu, Haocheng Wang, Haowei Zhang, Honghui Ding, Huajian Xin, Huazuo Gao, Hui Li, Hui Qu, J. L. Cai, Jian Liang, Jianzhong Guo, Jiaqi Ni, Jiashi Li, Jiawei Wang, Jin Chen, Jingchang Chen, Jingyang Yuan, Junjie Qiu, Junlong Li, Junxiao Song, Kai Dong, Kai Hu, Kaige Gao, Kang Guan, Kexin Huang, Kuai Yu, Lean Wang, Lecong Zhang, Lei Xu, Leyi Xia, Liang Zhao, Litong Wang, Liyue Zhang, Meng Li, Miaojun Wang, Mingchuan Zhang, Minghua Zhang, Minghui Tang, Mingming Li, Ning Tian, Panpan Huang, Peiyi Wang, Peng Zhang, Qiancheng Wang, Qihao Zhu, Qinyu Chen, Qiushi Du, R. J. Chen, R. L. Jin, Ruiqi Ge, Ruisong Zhang, Ruizhe Pan, Runji Wang, Runxin Xu, Ruoyu Zhang, Ruyi Chen, S. S. Li, Shanghao Lu, Shangyan Zhou, Shanhuang Chen, Shaoqing Wu, Shengfeng Ye, Shengfeng Ye, Shirong Ma, Shiyu Wang, Shuang Zhou, Shuiping Yu, Shunfeng Zhou, Shuting Pan, T. Wang, Tao Yun, Tian Pei, Tianyu Sun, W. L. Xiao, Wangding Zeng, Wanjia Zhao, Wei An, Wen Liu, Wenfeng Liang, Wenjun Gao, Wenqin Yu, Wentao Zhang, X. Q. Li, Xiangyue Jin, Xianzu Wang, Xiao Bi, Xiaodong Liu, Xiaohan Wang, Xiaojin Shen, Xiaokang Chen, Xiaokang Zhang, Xiaosha Chen, Xiaotao Nie, Xiaowen Sun, Xiaoxiang Wang, Xin Cheng, Xin Liu, Xin Xie, Xingchao Liu, Xingkai Yu, Xinnan Song, Xinxia Shan, Xinyi Zhou, Xinyu Yang, Xinyuan Li, Xuecheng Su, Xuheng Lin, Y. K. Li, Y. Q. Wang, Y. X. Wei, Y. X. Zhu, Yang Zhang, Yanhong Xu, Yanhong Xu, Yanping Huang, Yao Li, Yao Zhao, Yaofeng Sun, Yaohui Li, Yaohui Wang, Yi Yu, Yi Zheng, Yichao Zhang, Yifan Shi, Yiliang Xiong, Ying He, Ying Tang, Yishi Piao, Yisong Wang, Yixuan Tan, Yiyang Ma, Yiyuan Liu, Yongqiang Guo, Yu Wu, Yuan Ou, Yuchen Zhu, Yuduan Wang, Yue Gong, Yuheng Zou, Yujia He, Yukun Zha, Yunfan Xiong, Yunxian Ma, Yuting Yan, Yuxiang Luo, Yuxiang You, Yuxuan Liu, Yuyang Zhou, Z. F. Wu, Z. Z. Ren, Zehui Ren, Zhangli Sha, Zhe Fu, Zhean Xu, Zhen Huang, Zhen Zhang, Zhenda Xie, Zhengyan Zhang, Zhewen Hao, Zhibin Gou, Zhicheng Ma, Zhigang Yan, Zhihong Shao, Zhipeng Xu, Zhiyu Wu, Zhongyu Zhang, Zhuoshu Li, Zihui Gu, Zijia Zhu, Zijun Liu, Zilin Li, Ziwei Xie, Ziyang Song, Ziyi Gao, and Zizheng Pan. Deepseek-v3 technical report, 2025. URL https://arxiv.org/abs/2412.19437. [5] Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, and Hannaneh Hajishirzi. Olmoe: Open mixture-of-experts language models, 2025. URL https://arxiv.org/abs/2409.02060. [6] Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, Heyan Huang, and Furu Wei. On the representation collapse of sparse mixture of experts, 2022. URL https://arxiv.org/abs/2204.09179. [7] Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Auxiliary-loss-free load balancing strategy for mixture-of-experts. arXiv preprint arXiv:2408.15664, 2024. [8] Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. Stablemoe: Stable routing strategy for mixture of experts, 2022. URL https://arxiv.org/abs/2204.08396. [9] James B. MacQueen. Some methods for classification and analysis of multivariate observations. In Proceedings of the Fifth Berkeley Symposium on Mathematical Statistics and Probability, volume 1, pages 281–297. University of California Press, 1967. [10] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017. URL https://arxiv.org/abs/1706. 03762.
10
[11] Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024. URL https://arxiv.org/abs/2401.06066. [12] Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie-Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts, 2024. URL https://arxiv.org/abs/2401.04088. [13] Ka Man Lo, Zeyu Huang, Zihan Qiu, Zili Wang, and Jie Fu. A closer look into mixture-of-experts in large language models, 2025. URL https://aclanthology.org/2025.findings-naacl.251/. [14] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models, 2022. URL https: //arxiv.org/abs/2202.08906. [15] Hongcan Guo, Haolang Lu, Guoshun Nan, Bolun Chu, Jialin Zhuang, Yuan Yang, Wenhao Che, Xinye Cao, Sicong Leng, Qimei Cui, and Xudong Jiang. Advancing expert specialization for better moe. arXiv preprint arXiv:2505.22323, 2025. [16] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1–67, 2020. URL https://jmlr.org/papers/v21/ 20-074.html. [17] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020. URL https: //arxiv.org/abs/2101.00027. [18] Junzhuo Li, Bo Wang, Xiuze Zhou, Peijie Jiang, Jia Liu, and Xuming Hu. Decoding knowledge attribution in mixture-of-experts: A framework of basic-refinement collaboration and efficiency analysis. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 22431–22446, 2025. [19] Xingyi Yang, Constantin Venhoff, Ashkan Khakzar, Christian Schroeder de Witt, Puneet K. Dokania, Adel Bibi, and Philip Torr. Mixture of experts made intrinsically interpretable. In Proceedings of the 42nd International Conference on Machine Learning, 2025. [20] Jeremy Herbst, Jae Hee Lee, and Stefan Wermter. The expert strikes back: Interpreting mixture-of-experts language models at expert level. In Proceedings of the 43rd International Conference on Machine Learning, 2026. To appear. [21] Boan Liu, Liang Ding, Li Shen, Keqin Peng, Yu Cao, Dazhao Cheng, and Dacheng Tao. Diversifying the mixture-of-experts representation for language models with orthogonal optimizer, 2024. URL https: //arxiv.org/abs/2310.09762. [22] Ruijun Huang, Fang Dong, Xin Zhang, Hengjie Cao, Zhendong Huang, Anrui Chen, Jixian Zhou, Mengyi Chen, Yifeng Yang, Mingzhi Dong, Yujiang Wang, Jinlong Hou, Qin Lv, Robert P. Dick, Yuan Cheng, Fan Yang, Tun Lu, Chun Zhang, and Li Shang. Sd-moe: Spectral decomposition for effective expert specialization, 2026. URL https://arxiv.org/abs/2602.12556. [23] Nishanth Dikkala et al. On the benefits of learning to route in mixture-of-experts models, 2023. URL https://aclanthology.org/2023.emnlp-main.583/. [24] Ang Lv, Jin Ma, Yiyuan Ma, and Siyuan Qiao. Coupling experts and routers in mixture-of-experts via an auxiliary loss. arXiv preprint arXiv:2512.23447, 2025. [25] Sainbayar Sukhbaatar, Olga Golovneva, Vasu Sharma, Hu Xu, Xi Victoria Lin, Baptiste Rozière, Jacob Kahn, Daniel Li, Wen tau Yih, Jason Weston, and Xian Li. Branch-train-mix: Mixing expert llms into a mixture-of-experts llm, 2024. URL https://arxiv.org/abs/2403.07816. [26] Bowen Pan, Yikang Shen, Haokun Liu, Mayank Mishra, Gaoyuan Zhang, Aude Oliva, Colin Raffel, and Rameswar Panda. Dense training, sparse inference: Rethinking training of mixture-of-experts language models, 2024. URL https://arxiv.org/abs/2404.05567.
11
[27] Yuqi Xu, Rizhen Hu, Zihan Liu, Mou Sun, and Kun Yuan. Grouter: Decoupling routing from representation for accelerated moe training, 2026. URL https://arxiv.org/abs/2603.06626. [28] Anzhe Cheng, Shukai Duan, Shixuan Li, Chenzhong Yin, Mingxi Cheng, Shahin Nazarian, Paul Thompson, and Paul Bogdan. Emoe: Eigenbasis-guided routing for mixture-of-experts, 2026. URL https://arxiv. org/abs/2601.12137. [29] Anzhe Cheng, Shukai Duan, Shixuan Li, Chenzhong Yin, Mingxi Cheng, Heng Ping, Tamoghna Chattopadhyay, Sophia I Thomopoulos, Shahin Nazarian, Paul Thompson, and Paul Bogdan. Ermoe: Eigenreparameterized mixture-of-experts for stable routing and interpretable specialization, 2025. URL https://arxiv.org/abs/2511.10971. [30] Ibne Farabi Shihab, Sanjeda Akter, and Anuj Sharma. Grassmannian mixture-of-experts: Concentrationcontrolled routing on subspace manifolds, 2026. URL https://arxiv.org/abs/2602.17798. [31] Jama Hussein Mohamud, Drew Wagner, and Mirco Ravanelli. Self-routing: Parameter-free expert routing from hidden states, 2026. URL https://arxiv.org/abs/2604.00421. [32] Jiajie Yang. Latent prototype routing: Achieving near-perfect load balancing in mixture-of-experts, 2025. URL https://arxiv.org/abs/2506.21328. [33] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9: 53–68, 2021. [34] Nusrat Jahan Prottasha, Md Kowsher, Chun-Nam Yu, Chen Chen, and Ozlem Garibay. Monkey jump: Moe-style peft for efficient multi-task learning, 2026. URL https://arxiv.org/abs/2601.06356.
A
Theoretical foundation: the geometric alignment principle
To better understand the meaning of the router and expert weights, we analyzed how these weights are updated during training, focusing on the gradients involved. A.1
Setup
A.2
Expert gradient derivation
We derive the gradient for Wiup . In SwiGLU, the expert computes Ei (x) = Widown hi ,
hi = gi ⊙ Wiup x,
gi = SiLU(Wigate x).
(13)
Crucially, Wiup x enters the computation linearly: the gate branch gi depends on Wigate , not on Wiup . This is precisely why we analyze Wiup —the gradient path through the linear branch avoids activation-function derivatives entirely, yielding a clean input-directed update. P ∂y The SMoE output is y = j pj Ej (x), so ∂E = pi . Applying the chain rule: i 1. Through the SMoE weighted sum and down-projection: ∂L = (Widown )T (pi Ly ) ∈ Rdf f ∂hi
(14)
2. Through the Hadamard gate (linear in Wiup x): Since hi = gi ⊙ Wiup x and gi does not depend on Wiup , ∂L ∂L = ⊙ gi ∈ Rdf f (15) ∂(Wiup x) ∂hi 3. Final matrix gradient (outer product with input): ∂L ∂L = xT ∈ Rdf f ×d ∂Wiup ∂(Wiup x) 12
(16)
Table 2: Mathematical notations and setup parameters. Category
Notation
Input
x
The input token representation (or hidden state).
Rd
Router
Wr
The router weight matrix.
RN ×d
ri
The i-th row of Wr ; the router weight vector for expert i.
Rd
zi
The routing score (logit) for the i-th expert: zi = x⊤ ri
R
pi
The gating probability assigned to expert i, obtained via softmax Pover the logits: pi = ezi / j ezj
R
Expert
Description
Dimension
Wigate
The gate projection matrix for expert i.
Rdf f ×d
Wiup
The up-projection matrix for expert i.
Rdf f ×d
Widown
The down-projection matrix for expert i.
Rd×df f
Ei (x)
The output of the i-th expert (SwiGLU): Ei (x) = Widown (SiLU(Wigate x) ⊙ Wiup x)
Rd
Output
y
The final SMoE layer output, calculated as the weighted P sum of the expert outputs: y = j pj Ej (x)
Rd
Loss
Ly := ∂L ∂y
The upstream gradient, representing the derivative of the total loss L with respect to the output y.
Rd
The k-th row of Wiup (hidden neuron uk ) therefore receives the gradient ∂L = δi,k xT ∈ R1×d , ∂uk
δi,k = (Widown )T (pi Ly ) k · gi,k {z } |{z} | error signal
(17)
gate value
where δi,k is the scalar error term for neuron k, which absorbs the upstream loss gradient, the routing probability, the down-projection, and the gate activation—all factors that are independent of Wiup . The directional component of every row update is xT : each row of Wiup evolves as a weighted sum of the input representations it processes. A.3
Router gradient derivation
The routing logit for expert i is zi = xT ri , where ri is the i-th row of Wr . From the SMoE output P y = j pj Ej (x), the loss gradient with respect to pi is: ∂L = LTy Ei (x) ∈ R ∂pi
(18)
∂pi ∂L = LTy Ei (x) x = γi x ∈ Rd ∂ri ∂zi | {z }
(19)
i Since ∂z ∂ri = x, the chain rule gives:
γi
where γi is a scalar that absorbs the upstream loss gradient, the expert output, and the softmax Jacobian—all independent of ri itself. 13
A.4
The shared structure
Both gradients have the same form: a scalar error signal multiplied by the input x. ∂L = δi,k xT ∂uk
(expert row)
(20)
∂L = γi x ∂ri
(router weight vector)
(21)
The scalars δi,k and γi differ—they encode different error signals—but the directional component is the same input x in both cases. Over training, each row of Wiup and each router weight vector ri accumulates a weighted sum of the hidden states routed to expert i. Because matched router–expert pairs process the same token stream, this shared input-directed update structure predicts that they will develop statistically aligned geometry in hidden space.
B
Resources
All experiments were run on an Nvidia H100 node, or an AMD MI325X node.
14