Research map / Architectures and efficiency / Mixture of experts
Routers Learn the Geometry of Their Experts: Geometric Coupling in Sparse Mixture-of-Experts
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that SMoE routers and their selected experts receive gradients along the same input directions, that auxiliary balancing losses erode this coupling, and that a parameter-free online K-Means (EMA centroid) router achieves the lowest load imbalance with a modest perplexity cost.
Problem
Routing in sparse mixture-of-experts (SMoE) language models is hard to train: routing collapses onto few experts, and auxiliary load-balancing losses can reduce specialization. The mechanisms by which routing decisions are formed inside SMoEs are not well understood.
Approach
Backpropagation gives matched router rows and expert input-side weights gradients along the same hidden-state direction, so they accumulate a shared routed-token history; auxiliary balancing losses inject input-directed gradients into all router rows, eroding this coupling. A parameter-free router using EMA centroids of routed hidden states with cosine-similarity routing and sign-updated biases te
Key results
- Router scores correlate with expert gate-neuron activations (rho=0.43, p=1.2e-81) in a 1B SMoE trained on ~50B tokens.
- Auxiliary loss raises mean off-diagonal router cosine similarity to 0.63/0.63/0.57 at layers 0/4/8, versus 0.32/0.18/0.13 under loss-free balancing (nearly three times more similar).
- K-Means router reaches the lowest steady-state MaxVio (0.037 vs 0.084 for Loss-Free and 0.526 for Aux-Loss) with training PPL 15.40 vs 15.01 for Loss-Free (2.6% increase).
Surprising findings
- A parameter-free router with no learned routing weights and no balancing gradients achieves the lowest load imbalance among the variants tested.
Limitations
- Results are based on a single 1B SMoE configuration; scale and architecture generality untested.
- Activation-level evidence covers only the gate branch, not the up projection or other expert weights.
- K-Means router has a 2.6% training perplexity gap relative to the learned loss-free router, so it is not a practical replacement.
Experimental setup
Datasets: OLMoE-mix-0924 English Wikipedia C4-en The Pile
Benchmarks: C4-en validation perplexity Pile validation perplexity
Metrics: perplexity MaxVio load imbalance Spearman correlation between router score and expert activat mean pairwise cosine similarity of router weight vectors
Baselines: Aux-Loss (Switch-style load balance + z-loss) Loss-Free (bias-only balancing) Loss-Free + Seq-Aux (DeepSeek-V3 sequence-wise aux loss)
Builds on: OLMoE Switch Transformer load-balancing loss loss-free balancing ST-MoE router z-loss online k-means
Closest papers by meaning
- Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary Loss (ICLR 2026 · similarity 0.92)
- Layerwise Recurrent Router for Mixture-of-Experts (ICLR 2025 · similarity 0.90)
- Understanding the Mixture-of-Experts with Nadaraya-Watson Kernel (ICLR 2026 · similarity 0.90)
- Tight Clusters Make Specialized Experts (ICLR 2025 · similarity 0.90)
- Routing Experts: Learning to Route Dynamic Experts in Existing Multi-modal Large Language Models (ICLR 2025 · similarity 0.89)
- Statistical Advantages of Perturbing Cosine Router in Mixture of Experts (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.
