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

TL;DR

This paper introduces kNN-MoE, a retrieval-augmented routing method for mixture-of-experts models that improves routing decisions by leveraging a memory of similar past cases, enhancing robustness and performance.

Contribution

The paper proposes a novel kNN-based routing framework for MoE models that reuses expert assignments from a memory, addressing brittleness under distribution shifts.

Findings

kNN-MoE outperforms zero-shot baselines.

It rivals expensive supervised fine-tuning.

The method improves routing robustness.

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.