# Improving Routing in Sparse Mixture of Experts with Graph of Tokens

arXiv: https://arxiv.org/abs/2505.00792 · HTML: https://mlsanity.com/paper/2505.00792/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts.
- Builds on 34 of its 50 references that are in the map.
- First posted on arXiv 1 May 2025.

## Placed in

- Mixture-of-experts / [Expert routing strategies](https://mlsanity.com/moe/expert-routing-strategies/)

## Builds on (most cited first, 30 of 34)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models](https://mlsanity.com/paper/2301.12597/) (2023)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [The Many Faces of Robustness: A Critical Analysis of Out-of-Distribution Generalization](https://mlsanity.com/paper/2006.16241/) (2020)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [Natural Adversarial Examples](https://mlsanity.com/paper/1907.07174/) (2019)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Scaling Vision with Sparse Mixture of Experts](https://mlsanity.com/paper/2106.05974/) (2021)
- [Efficient Intent Detection with Dual Sentence Encoders](https://mlsanity.com/paper/2003.04807/) (2020)
- [VLMo: Unified Vision-Language Pre-Training with Mixture-of-Modality-Experts](https://mlsanity.com/paper/2111.02358/) (2021)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [BASE Layers: Simplifying Training of Large, Sparse Models](https://mlsanity.com/paper/2103.16716/) (2021)
- [Hash Layers For Large Sparse Models](https://mlsanity.com/paper/2106.04426/) (2021)
- [Sequential Modeling Enables Scalable Learning for Large Vision Models](https://mlsanity.com/paper/2312.00785/) (2023)
- [Efficient Large Scale Language Modeling with Mixtures of Experts](https://mlsanity.com/paper/2112.10684/) (2021)
- [Conditional Computation in Neural Networks for faster models](https://mlsanity.com/paper/1511.06297/) (2015)
- [On the Representation Collapse of Sparse Mixture of Experts](https://mlsanity.com/paper/2204.09179/) (2022)
- [Towards Foundation Models for Scientific Machine Learning: Characterizing Scaling and Transfer Behavior](https://mlsanity.com/paper/2306.00258/) (2023)
- [StableMoE: Stable Routing Strategy for Mixture of Experts](https://mlsanity.com/paper/2204.08396/) (2022)
- [Sparse MoE as the New Dropout: Scaling Dense and Self-Slimmable Transformers](https://mlsanity.com/paper/2303.01610/) (2023)
- [Improving Transformer Models by Reordering their Sublayers](https://mlsanity.com/paper/1911.03864/) (2019)

## Built on by (most cited first, 4 of 4)

- [Pre-Attention Expert Prediction and Prefetching for Mixture-of-Experts Large Language Models](https://mlsanity.com/paper/2511.10676/) (2025)
- [Modeling Expert Interactions in Sparse Mixture of Experts via Graph Structures](https://mlsanity.com/paper/2510.16411/) (2025)
- [Region-Graph Optimal Transport Routing for Mixture-of-Experts Whole-Slide Image Classification](https://mlsanity.com/paper/2604.07298/) (2026)
- [Hierarchical Copula-Gumbel-Top-K Routing: Two-Sided Dependence Control for Frozen Mixture-of-Experts at Fixed Per-Token Routing Laws](https://mlsanity.com/paper/2607.28670/) (2026)

## Abstract (arXiv)

Sparse Mixture of Experts (SMoE) has emerged as a key to achieving unprecedented scalability in deep learning. By activating only a small subset of parameters per sample, SMoE achieves an exponential increase in parameter counts while maintaining a constant computational overhead. However, SMoE models are susceptible to routing fluctuations--changes in the routing of a given input to its target expert--at the late stage of model training, leading to model non-robustness. In this work, we unveil the limitation of SMoE through the perspective of the probabilistic graphical model (PGM). Through this PGM framework, we highlight the independence in the expert-selection of tokens, which exposes the model to routing fluctuation and non-robustness. Alleviating this independence, we propose the novel Similarity-Aware (S)MoE, which considers interactions between tokens during expert selection. We then derive a new PGM underlying an (S)MoE-Attention block, going beyond just a single (S)MoE layer. Leveraging the token similarities captured by the attention matrix, we propose the innovative Attention-Aware (S)MoE, which employs the attention matrix to guide the routing of tokens to appropriate experts in (S)MoE. We theoretically prove that Similarity/Attention-Aware routing help reduce the entropy of expert selection, resulting in more stable token routing mechanisms. We empirically validate our models on various tasks and domains, showing significant improvements in reducing routing fluctuations, enhancing accuracy, and increasing model robustness over the baseline MoE-Transformer with token routing via softmax gating.