# Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Expert routing strategies, one of 12 directions in Mixture-of-experts.
- Builds on 19 of its 29 references that are in the map.
- First posted on arXiv 5 Jan 2026.

## Placed in

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

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Measuring Massive Multitask Language Understanding](https://mlsanity.com/paper/2009.03300/) (2020)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [GPQA: A Graduate-Level Google-Proof Q&A Benchmark](https://mlsanity.com/paper/2311.12022/) (2023)
- [Billion-scale similarity search with GPUs](https://mlsanity.com/paper/1702.08734/) (2017)
- [gpt-oss-120b & gpt-oss-20b Model Card](https://mlsanity.com/paper/2508.10925/) (2025)
- [What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Exams](https://mlsanity.com/paper/2009.13081/) (2020)
- [Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models](https://mlsanity.com/paper/2506.05176/) (2025)
- [Generalization through Memorization: Nearest Neighbor Language Models](https://mlsanity.com/paper/1911.00172/) (2019)
- [Why do Nearest Neighbor Language Models Work?](https://mlsanity.com/paper/2301.02828/) (2023)
- [Coupling Experts and Routers in Mixture-of-Experts via an Auxiliary Loss](https://mlsanity.com/paper/2512.23447/) (2025)
- [C3PO: Critical-Layer, Core-Expert, Collaborative Pathway Optimization for Test-Time Expert Re-Mixing](https://mlsanity.com/paper/2504.07964/) (2025)
- [Rewiring Experts on the Fly:Continuous Rerouting for Better Online Adaptation in Mixture-of-Expert models](https://mlsanity.com/paper/2510.14853/) (2025)
- [Case-Based Decision-Theoretic Decoding with Quality Memories](https://mlsanity.com/paper/2509.12677/) (2025)

## Abstract (arXiv)

Mixture-of-Experts (MoE) architectures scale large language models efficiently by employing a parametric ``router'' to dispatch tokens to a sparse subset of experts. Typically, this router is trained once and then frozen, rendering routing decisions brittle under distribution shifts. We address this limitation by introducing kNN-MoE, a retrieval-augmented routing framework that reuses locally optimal expert assignments from a memory of similar past cases. This memory is constructed offline by directly optimizing token-wise routing logits to maximize the likelihood on a reference set. Crucially, we use the average similarity of retrieved neighbors as a confidence-driven mixing coefficient, thus allowing the method to fall back to the frozen router when no relevant cases are found. Experiments show that kNN-MoE outperforms the zero-shot baseline and is competitive with computationally intensive supervised fine-tuning.