Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

TL;DR

MoE models rely on a static, parametric router that can fail under distribution shift. The authors propose kNN-MoE, a retrieval-augmented routing framework that builds a memory of router inputs paired with offline-optimized expert assignments from a reference set, and blends this memory with the frozen router using a confidence-based interpolation. The memory is constructed by optimizing per-token routing logits to maximize the likelihood of ground-truth tokens, enabling the memory to provide high-quality routing guidance at inference time through nearest-neighbor retrieval. Experiments across three MoE models and multiple challenging benchmarks show that kNN-MoE improves zero-shot performance and is competitive with supervised fine-tuning while avoiding online parameter updates, offering a scalable path to robust MoE routing under distribution shifts.

Abstract

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 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 aggregate 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 kNN-MoE outperforms zero-shot baselines and rivals computationally expensive supervised fine-tuning.

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

TL;DR

MoE models rely on a static, parametric router that can fail under distribution shift. The authors propose kNN-MoE, a retrieval-augmented routing framework that builds a memory of router inputs paired with offline-optimized expert assignments from a reference set, and blends this memory with the frozen router using a confidence-based interpolation. The memory is constructed by optimizing per-token routing logits to maximize the likelihood of ground-truth tokens, enabling the memory to provide high-quality routing guidance at inference time through nearest-neighbor retrieval. Experiments across three MoE models and multiple challenging benchmarks show that kNN-MoE improves zero-shot performance and is competitive with supervised fine-tuning while avoiding online parameter updates, offering a scalable path to robust MoE routing under distribution shifts.

Abstract

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 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 aggregate 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 kNN-MoE outperforms zero-shot baselines and rivals computationally expensive supervised fine-tuning.
Paper Structure (36 sections, 9 equations, 3 figures, 12 tables)

This paper contains 36 sections, 9 equations, 3 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Expert Assignment Optimization for MoE
  4. Non-Parametric and Retrieval-Augmented LMs
  5. Mixture-of-Experts
  6. Proposed Method: kNN-MoE
  7. Memory Construction
  8. Data Collection
  9. Deriving Optimal Expert Assignments
  10. Confidence-Aware Adaptive Mixing
  11. Experiments
  12. Experimental Settings
  13. Datasets
  14. Models
  15. Baselines
  16. ...and 21 more sections

Figures (3)

  • Figure 1: Schematic illustration of the kNN-MoE inference mechanism. Given a router input, the model simultaneously obtains a prediction from the frozen "Parametric Router" and retrieves similar past cases from the "Memory". The "Adaptive Mixing" module fuses these two signals based on a confidence score (derived from the aggregate similarity of retrieved neighbors) to produce the final expert assignment.
  • Figure 2: Prompt template for the zero-shot setting.
  • Figure 3: Prompt template for the 5-shot setting.