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.
@article{arxiv.2601.02144,
title = {Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts},
author = {Boxuan Lyu and Soichiro Murakami and Hidetaka Kamigaito and Peinan Zhang},
journal= {arXiv preprint arXiv:2601.02144},
year = {2026}
}