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

> Router and expert weights in sparse mixture-of-experts models naturally align through shared gradient updates, but auxiliary load-balancing losses disrupt this coupling, while a parameter-free K-Means router preserves it with minimal perplexity cost.

- **Source:** [arXiv](https://arxiv.org/abs/2605.12476)
- **Published:** 2026-08-29
- **Permalink:** https://picx.dev/p/Hdy570
- **Whiteboard:** https://picx.dev/p/Hdy570/image

## Summary

# Sparse Mixture-of-Experts: Geometric Coupling Between Routers and Experts

## Summary (Overview)
- **Core Finding**: Router weight vectors and expert input-side weights in Sparse Mixture-of-Experts (SMoE) models develop coupled geometry because both receive gradient updates proportional to the same routed hidden states during training.
- **Empirical Evidence**: In a 1B SMoE trained from scratch, router scores strongly predict expert gate-neuron activations (correlation ρ = 0.43), confirming that routing decisions are mirrored inside expert computation.
- **Auxiliary Loss Effect**: Standard auxiliary load-balancing losses break this coupling by injecting global interference gradients into all router weight vectors, making distinct router directions nearly three times more similar (cosine similarity 0.63 vs. 0.18).
- **Constructive Test**: A parameter-free K-Means router using EMA centroids of routed hidden states achieves the lowest load imbalance (MaxVio = 0.037) with only modest perplexity cost (2.6% increase), demonstrating that centroid tracking captures a substantial component of learned routing.

## Introduction and Theoretical Foundation

### Background
Sparse Mixture-of-Experts (SMoE) models scale language models efficiently by activating only a fraction of parameters per token. A router network selects K experts per input from N total experts, enabling a division of labor where different experts specialize in different input patterns.

### The Routing Mechanism
For a hidden state **x**, the router computes scores:

$$p = \sigma(\mathbf{z} + \mathbf{m}), \quad \mathbf{z} = W_r \mathbf{x}$$

where $W_r \in \mathbb{R}^{N \times d}$ contains router weight vectors $\mathbf{r}_i$ for each expert $i$, and the SMoE layer combines selected expert outputs:

$$\mathbf{y} = \sum_{i \in \mathcal{T}_K} p_i E_i(\mathbf{x})$$

Experts use gated SwiGLU activations:

$$E_i(\mathbf{x}) = W_i^{\text{down}}\left(\sigma(W_i^{\text{gate}}\mathbf{x}) \odot W_i^{\text{up}}\mathbf{x}\right)$$

### Training Challenges
- **Representation collapse**: Without intervention, routing concentrates on few experts
- **Auxiliary load-balancing losses** [1, 3, 2]: Add penalties to balance expert utilization but may reduce specialization
- **Loss-free methods** [7, 4]: Use adaptive per-expert routing biases instead

## Methodology

### Theoretical Derivation of Geometric Coupling
The paper derives that both router weights and expert input-side weights receive gradient updates proportional to the input direction **x**:

**Expert side** (for row $\mathbf{w}_{i,k}$ of $W_i^{\text{gate}}$):
$$\nabla_{\mathbf{w}_{i,k}} \mathcal{L} = \delta_{i,k} \mathbf{x}^{\top} \propto \mathbf{x}^{\top}$$

**Router side** (for router weight vector $\mathbf{r}_i$):
$$\nabla_{\mathbf{r}_i} \mathcal{L} = \gamma_i \mathbf{x} \propto \mathbf{x}$$

This shared gradient structure means matched router–expert pairs accumulate the same routed token history, creating geometric alignment.

### Experimental Setup
- **Model**: 1B SMoE with 9 layers, hidden size 1024, 64 routed experts, top-K = 6 routing
- **Training**: OLMoE-mix-0924 dataset, ~50B tokens, AdamW optimizer
- **Balancing variants tested**: Aux-Loss (Switch-style), Loss-Free (bias-only), Loss-Free + Seq-Aux (DeepSeek-V3 style), K-Means (centroid-based)

### K-Means Router Design
Replaces learned router weights with non-learnable EMA centroids:

$$s_i(\mathbf{x}) = \frac{\mathbf{c}_i^{\top}\mathbf{x}}{\|\mathbf{c}_i\|\|\mathbf{x}\|} + b_i$$

$$\mathbf{c}_i \leftarrow \alpha \mathbf{c}_i + (1-\alpha)\bar{\mathbf{x}}_i$$

where $\bar{\mathbf{x}}_i$ is the mean of tokens routed to expert $i$ in the current micro-batch, and $b_i$ uses the loss-free bias update rule.

## Empirical Validation / Results

### 1. Router Scores Predict Expert Activations
- **Finding**: Monotone relationship between router score and expert gate-neuron activations
- **Correlation**: ρ = 0.43 (p-value 1.2 × 10⁻⁸¹)
- **Significance**: These are computed independently in the same forward pass, providing direct functional evidence of coupling

### 2. Auxiliary Loss Collapses Router Geometry
Pairwise cosine similarities between router weight vectors:

| Layer | Aux-Loss (μ) | Loss-Free (μ) |
|-------|-------------|---------------|
| 0     | 0.63        | 0.32          |
| 4     | 0.63        | 0.18          |
| 8     | 0.57        | 0.13          |

The auxiliary loss makes router directions nearly three times more similar, eroding expert-specific signal.

### 3. K-Means Router Performance
| Method | Router Params | Train PPL | C4-en PPL | Pile PPL | MaxVio |
|--------|--------------|-----------|-----------|----------|--------|
| Aux-Loss | 0.59M | 15.09 | 20.54 | 11.82 | 0.526 |
| Loss-Free + Seq-Aux | 0.59M | 15.03 | 20.44 | 11.77 | 0.102 |
| Loss-Free | 0.59M | 15.01 | 20.40 | 11.76 | 0.084 |
| **K-Means (ours)** | **0** | **15.40** | **21.01** | **12.09** | **0.037** |

Key results:
- K-Means achieves the lowest load imbalance (0.037) with zero trainable router parameters
- Only 2.6% perplexity increase over Loss-Free (15.40 vs. 15.01)
- No auxiliary losses or balancing gradients needed

## Theoretical and Practical Implications

### Theoretical Contributions
1. **Gradient-level explanation**: Provides the first formal account of why router–expert correlation emerges naturally from backpropagation through SMoE layers
2. **Mechanism for load-balancing degradation**: Explains why auxiliary losses reduce specialization - they inject global interference gradients that collapse router weight vector diversity
3. **Centroid tracking as implicit mechanism**: Shows learned routers essentially perform online clustering of routed hidden states

### Practical Implications
1. **Router design guidance**: Auxiliary balancing losses may be counterproductive; bias-only or centroid-based approaches preserve geometric coupling
2. **Evaluation framework**: Router quality should be assessed not just by load balance but by preservation of router–expert geometry
3. **Architecture insights**: The success of parameter-free K-Means routing suggests learned routers have significant redundancy, motivating hybrid designs

## Conclusion

The paper establishes that router–expert coupling is a fundamental emergent property of SMoE training, arising from shared input-directed gradients. This coupling:
- Is empirically visible in token-level expert activations
- Is fragile under auxiliary load-balancing interventions
- Can be largely recovered through explicit centroid tracking

**Future directions**: 
- Validating coupling at larger scales and different architectures
- Extending activation analysis to other expert weight matrices ($W_i^{\text{up}}$)
- Developing hybrid routers that preserve centroid geometry while closing the perplexity gap
- Using the geometric lens to understand other SMoE pathologies (representation collapse, expert dominance)

---

_Markdown view of https://picx.dev/p/Hdy570, served by PicX — AI-generated visual whiteboard summaries of research papers._