💬 NLP

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

This paper introduces kNN-MoE, a retrieval-augmented routing framework that dynamically improves Mixture-of-Experts models by leveraging a memory of optimized past routing decisions and confidence-driven fallback mechanisms to handle distribution shifts more effectively than frozen routers.

Original authors: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

Published 2026-05-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive, highly intelligent library of experts (a "Mixture-of-Experts" model). When you ask this library a question, a Router acts like a librarian. Its job is to look at your question and quickly decide which specific expert in the library is best suited to answer it.

Usually, this librarian is trained once and then frozen in place. They rely on their memory of what they learned during training. But here's the problem: if you ask a weird, difficult, or brand-new type of question that the librarian hasn't seen before, they might guess wrong and send you to the wrong expert. This is like a librarian sending a medical question to a history expert because the words look vaguely similar.

The paper introduces kNN-MoE, a clever upgrade that gives this librarian a "cheat sheet" based on past successes. Here is how it works, broken down into simple concepts:

1. The "Cheat Sheet" (Memory Construction)

Before the system ever answers a real question, the researchers take a set of practice questions (a reference dataset) and run them through the library.

The Experiment: For every single word in these practice questions, they ask: "If we could magically change the librarian's decision right now, which expert would have given the absolute best answer?"
The Result: They find the "perfect" expert for each specific moment and write it down. They store these pairs: "This specific question input" + "The perfect expert choice."
The Analogy: Imagine the librarian is studying for a final exam. Instead of just memorizing the rules, they create a massive index card deck. On one side is a tricky question; on the other is the exact expert who solved it perfectly in the past.

2. The "Smart Lookup" (Inference)

Now, when a real user asks a question, the system does two things simultaneously:

The Frozen Librarian: The original, frozen router makes its best guess based on its training.
The Cheat Sheet: The system looks up the user's question in the index card deck to find the most similar past questions.

3. The "Confidence Vote" (Adaptive Mixing)

This is the most important part. The system doesn't blindly trust the cheat sheet. It checks how similar the past cases are to the current question.

High Confidence: If the past cases are almost identical to the current question, the system says, "The librarian is probably unsure about this tricky one, but our cheat sheet says the perfect expert is X." It then mixes the librarian's guess with the cheat sheet's advice, leaning heavily on the cheat sheet.
Low Confidence: If the current question is totally unique and nothing in the cheat sheet matches well, the system says, "The cheat sheet is useless here; we might just add noise." It ignores the lookup and trusts the original librarian's frozen decision.

Why This Matters

The paper claims this method is a "sweet spot" between two extremes:

Doing Nothing (Zero-Shot): Just using the frozen librarian. This is fast but can fail on hard or new questions.
Retraining (Supervised Fine-Tuning): Teaching the librarian new rules from scratch. This works well but is incredibly slow, expensive, and requires re-doing the whole process for every new task.

kNN-MoE gets the performance boost of retraining without the heavy cost. It's like giving the librarian a dynamic, searchable memory of their best past moves rather than forcing them to go back to school.

Key Findings from the Paper

It works on hard questions: The system helps the most when the original librarian is confused (high "perplexity"). When the librarian is already confident, the system steps back to avoid messing things up.
Less is more: Surprisingly, looking at just one past example (the single closest match) worked better than averaging many examples. The paper suggests that for expert routing, having too many "opinions" from the past actually dilutes the signal.
Speed vs. Accuracy: It is much faster to build this "cheat sheet" (offline) than to retrain the whole model. During the actual answering phase, it adds a tiny bit of delay (about 3-4% slower) but significantly improves accuracy on difficult tasks like medical exams and coding.

The Catch (Limitations)

The paper notes that this system needs a "reference set" of labeled data to build the cheat sheet. If you are in a situation where you have absolutely no similar past examples to learn from, this method can't help. It relies on the assumption that "what worked before" is a good guide for "what will work now."

In short, kNN-MoE is a way to make AI experts smarter by letting them peek at their own history of perfect decisions, but only when it's safe to do so.

Technical Summary: kNN-MoE

Problem Statement

Mixture-of-Experts (MoE) architectures scale Large Language Models (LLMs) efficiently by using a parametric "router" to dispatch tokens to a sparse subset of experts. However, standard routers are lightweight classifiers trained once and then frozen. While effective on in-distribution data, they lack flexibility during inference when facing distribution shifts (e.g., specialized domains or out-of-distribution tasks). In such scenarios, the router must extrapolate from fixed priors, often leading to suboptimal expert selection and high perplexity. Existing test-time adaptation methods, such as C3PO, attempt to fine-tune routing logits online but incur significant computational overhead (approximately 5x the baseline). Other approaches require permanent parameter updates via supervised fine-tuning (SFT), which is computationally intensive and less efficient for inference.

Methodology: kNN-MoE

The authors propose kNN-MoE, a retrieval-augmented routing framework that refines expert assignment decisions without modifying the model's parameters. The method operates in two distinct phases:

1. Memory Construction (Offline)

The system builds a non-parametric memory store $\mathcal{M}$ for each MoE layer, consisting of key-value pairs $(k_t, v_t)$ :

Key ( $k_t$ ): The router input (hidden state) $x_t$ from a reference dataset $D_{ref}$ .
Value ( $v_t$ ): A "locally optimal" expert assignment derived by optimizing the routing logits to maximize the likelihood of the ground-truth token $y_t$ $y_{t}$ .
- Instead of storing the frozen router's output, the method solves an optimization problem for each token in $D_{ref}$ : finding a logit vector $r^*_t$ that minimizes the negative log-likelihood of the target token, given the frozen network parameters.
- This optimization is performed via a few steps of gradient descent (typically $S=1$ ) starting from the original parametric logits.
- The resulting optimized logits $r^{(S)}_t$ serve as the value $v_t$ in the memory.

2. Confidence-Aware Adaptive Mixing (Inference)

During inference, for a given router input $x$ :

Retrieval: The system retrieves the $K$ nearest neighbors $\mathcal{N}(x)$ from the memory $\mathcal{M}$ based on the similarity between $x$ and the stored keys.
Non-Parametric Proposal: A memory-based logit proposal $r_{mem}(x)$ is computed by aggregating the retrieved optimal logits, weighted by their similarity to the input.
Confidence-Driven Mixing: The final routing logits are a linear interpolation between the original parametric router logits $r(x)$ $r (x)$ and the memory-based proposal $r_{mem}(x)$ $r_{m e m} (x)$ :
$r_{final}(x) = (1 - \lambda(x))r(x) + \lambda(x)r_{mem}(x)$
The mixing coefficient $\lambda(x)$ $λ (x)$ is defined as the average similarity of the retrieved neighbors. This acts as a confidence score:
- High Similarity ( $\lambda \approx 1$ ): The model trusts the retrieved, locally optimal assignment.
- Low Similarity ( $\lambda \approx 0$ ): The model falls back to the original frozen router, avoiding noise from irrelevant cases.
Output: The final logits are passed through a Top-K Softmax to generate the sparse expert assignment, maintaining the model's original sparsity constraints.

Key Contributions

Retrieval-Augmented Routing: Introduces a framework that reuses offline-computed, locally optimal expert assignments from a memory of similar past cases to improve routing decisions.
Non-Parametric Adaptation: Unlike SFT or online adaptation methods, kNN-MoE does not update model parameters. It shifts the computational burden to an offline memory construction phase, enabling efficient inference.
Confidence-Driven Fallback: The method utilizes the average similarity of retrieved neighbors as a dynamic mixing coefficient, ensuring robustness by reverting to the frozen router when no relevant historical cases are found.
Efficiency: The approach requires only a single gradient descent step per token during memory construction and avoids the heavy computational costs of online optimization or full model fine-tuning.

Experimental Results

The authors evaluated kNN-MoE on three MoE models (OLMoE, GPT-OSS, and Qwen3) across six benchmarks (GPQA, SuperGPQA, MMLU, USMLE, MedMCQA, and MBPP).

Performance vs. Baselines: kNN-MoE consistently outperformed the Zero-shot baseline and the "SFT (Router Only)" baseline across all models and datasets.
Competitiveness with SFT: The method achieved performance competitive with full Supervised Fine-Tuning (SFT) and, in several data-scarce settings (e.g., GPQA, USMLE), outperformed SFT, which occasionally suffered from overfitting on small reference sets.
Stability: Unlike 5-shot prompting, which degraded performance on certain models (GPT-OSS, OLMoE) due to input length issues, kNN-MoE exhibited stable improvements.
Ablation Studies:
- Neighbors ( $K$ ): Contrary to typical kNN-LM findings, performance peaked at $K=1$ . Aggregating more neighbors ( $K>1$ ) introduced noise, suggesting that truly beneficial routing signals are rare and specific.
- Optimization Steps ( $S$ ): Increasing gradient descent steps beyond $S=1$ yielded negligible accuracy gains while increasing construction costs.
- Reference Set Size: Performance improved monotonically with the size of the reference set $|D_{ref}|$ , with diminishing returns observed beyond 0.5k–1k examples.
Efficiency: Memory construction took 0.46 hours for 1k examples (vs. ~1.34h for SFT). The IVFPQ (Inverted File + Product Quantization) variant reduced inference latency overhead to ~3.67% and VRAM usage growth to ~9.43% compared to the baseline.

Significance and Claims

The paper claims that kNN-MoE bridges the gap between static parametric routing and dynamic inference needs. Its primary significance lies in offering a scalable pathway to enhance MoE adaptability without the computational burden of fine-tuning.

The authors emphasize that the method is particularly effective when the parametric router is uncertain (correlated with high perplexity), as the retrieval mechanism provides a corrective signal from similar, high-quality historical cases. While the method relies on the availability of a labeled reference set with a distribution similar to the test data, it provides a practical, non-parametric alternative to expensive adaptation techniques, achieving SFT-competitive results with significantly lower preparation time and no parameter updates.