Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts
Fuente:
arXiv
Gespeichert in:
| Hauptverfasser: | , , , |
|---|---|
| Format: | Preprint |
| Veröffentlicht: |
2026
|
| Schlagworte: | |
| Online-Zugang: | |
| Tags: |
Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
|
| _version_ | 1866910252905005056 |
|---|---|
| author | Lyu, Boxuan Murakami, Soichiro Kamigaito, Hidetaka Zhang, Peinan |
| author_facet | Lyu, Boxuan Murakami, Soichiro Kamigaito, Hidetaka Zhang, Peinan |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_02144 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts Lyu, Boxuan Murakami, Soichiro Kamigaito, Hidetaka Zhang, Peinan Computation and Language Artificial Intelligence 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. |
| title | Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts |
| topic | Computation and Language Artificial Intelligence |
| url | https://arxiv.org/abs/2601.02144 |