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

📄 arXiv: 2601.02144v1 📥 PDF

作者: Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

分类: cs.CL, cs.AI

发布日期: 2026-01-05

💡 一句话要点

提出kNN-MoE，通过检索增强路由解决MoE模型在分布偏移下的脆弱性问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 路由算法 k近邻检索 分布偏移 检索增强学习

📋 核心要点

  1. 现有MoE模型中的路由通常是静态的，无法适应分布偏移，导致性能下降。
  2. kNN-MoE通过检索与当前输入相似的历史样本，并重用其专家分配来增强路由决策。
  3. 实验表明，kNN-MoE在零样本设置下表现优异，并能与监督微调方法竞争。

📝 摘要（中文）

混合专家模型(MoE)通过使用参数化的“路由”将token分派到稀疏的专家子集，从而有效地扩展大型语言模型。通常，此路由在训练一次后被冻结，导致路由决策在分布偏移下变得脆弱。为了解决这个限制，我们引入了kNN-MoE，这是一个检索增强的路由框架，它重用来自相似过去案例记忆中的最佳专家分配。该记忆通过直接优化token级别的路由logits以最大化参考集上的似然度来离线构建。至关重要的是，我们使用检索到的邻居的聚合相似性作为置信度驱动的混合系数，从而允许该方法在未找到相关案例时回退到冻结的路由。实验表明，kNN-MoE优于零样本基线，并且可以与计算成本高昂的监督微调相媲美。

🔬 方法详解

问题定义：MoE模型中的路由层通常在训练后被冻结，这使得模型在面对数据分布偏移时难以做出准确的路由决策。现有的静态路由方法无法有效利用历史信息，导致模型泛化能力受限。

核心思路：kNN-MoE的核心思想是利用k近邻(kNN)检索来增强路由决策。通过维护一个包含历史token及其最佳专家分配的记忆库，当新的token到来时，模型可以检索到与其最相似的token，并利用这些相似token的专家分配信息来指导当前token的路由。这种方法允许模型根据历史经验动态调整路由策略，从而更好地适应分布偏移。

技术框架：kNN-MoE的整体框架包括以下几个主要模块：1) 离线记忆库构建：使用参考集优化token级别的路由logits，并将token及其对应的最佳专家分配存储到记忆库中。2) 在线kNN检索：对于每个新的token，使用kNN算法在记忆库中检索与其最相似的k个token。3) 相似度加权混合：使用检索到的邻居的聚合相似度作为置信度驱动的混合系数，将kNN检索结果与原始的静态路由器的输出进行加权混合，得到最终的路由logits。

关键创新：kNN-MoE的关键创新在于将kNN检索与MoE路由相结合，从而实现了动态的、自适应的路由决策。与传统的静态路由方法相比，kNN-MoE能够利用历史信息来指导路由，从而更好地适应分布偏移。此外，使用相似度加权混合的方式，使得模型能够在找到相关案例时利用kNN检索结果，而在没有相关案例时回退到原始的静态路由器，从而保证了模型的鲁棒性。

关键设计：关键设计包括：1) 记忆库的构建方式：通过直接优化token级别的路由logits来构建记忆库，保证了记忆库中存储的是最佳的专家分配信息。2) 相似度度量方式：选择合适的相似度度量方式（例如余弦相似度）来衡量token之间的相似性。3) 相似度加权混合系数的计算方式：使用检索到的邻居的聚合相似度作为混合系数，使得模型能够根据检索结果的置信度来调整kNN检索结果的权重。

🖼️ 关键图片

fig_0

📊 实验亮点

实验结果表明，kNN-MoE在零样本设置下优于基线方法，并且能够与计算成本高昂的监督微调方法相媲美。具体来说，kNN-MoE在多个数据集上取得了显著的性能提升，证明了其在处理分布偏移方面的有效性。此外，实验还验证了相似度加权混合策略的有效性，表明该策略能够有效地平衡kNN检索结果和原始静态路由器的输出。

🎯 应用场景

kNN-MoE可应用于各种需要处理分布偏移的MoE模型，例如自然语言处理、语音识别和计算机视觉等领域。该方法能够提高模型在实际应用中的鲁棒性和泛化能力，尤其是在数据分布不稳定或存在对抗性攻击的情况下。此外，kNN-MoE还可以用于在线学习场景，通过不断更新记忆库来适应新的数据分布。

📄 摘要（原文）

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.