💬 NLP

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

本文介绍了 kNN-MoE，这是一种检索增强型路由框架，它通过利用优化后的历史路由决策记忆和基于置信度的回退机制来动态改进混合专家模型，从而比固定路由更有效地处理分布偏移。

原作者： Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

发布于 2026-05-26

📖 1 分钟阅读☕ 轻松阅读

原作者： Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你拥有一个庞大且高度智能的专家图书馆（即“混合专家”模型）。当你向这座图书馆提问时，路由器就像一位图书管理员。它的任务是审视你的问题，并迅速决定图书馆中哪位特定专家最适合回答它。

通常，这位图书管理员只经过一次训练，随后便被“冻结”在原地。它们依赖的是训练期间所学的记忆。但问题在于：如果你问了一个奇怪、困难或全新类型的问题，而图书管理员从未见过，它们可能会猜错，并将你引向错误的专家。这就像图书管理员因为某些词汇看起来 vaguely 相似，就将一个医学问题派给了历史专家。

本文介绍了kNN-MoE，这是一种巧妙的升级，它为这位图书管理员提供了一份基于过往成功的“作弊小抄”。其工作原理可拆解为以下简单概念：

1. “作弊小抄”（记忆构建）

在系统回答任何实际问题之前，研究人员会选取一组练习题（参考数据集），并将它们输入图书馆。

实验过程：对于练习题中的每一个词，他们都会问：“如果我们能在此刻神奇地改变图书管理员的决策，哪位专家会给出绝对最佳答案？”
结果：他们为每个特定时刻找出“完美”的专家并记录下来。他们存储了这些配对：“这个特定的问题输入” + “完美的专家选择”。
类比：想象图书管理员正在为期末考试做准备。他们不是仅仅死记硬背规则，而是制作了一副巨大的索引卡片。卡片的一面是棘手的问题，另一面则是过去完美解决该问题的确切专家。

2. “智能查找”（推理）

现在，当真实用户提问时，系统会同时做两件事：

冻结的图书管理员：原始的、被冻结的路由器基于其训练做出最佳猜测。
作弊小抄：系统在索引卡片中查找用户的问题，以找到最相似的过往问题。

3. “置信度投票”（自适应混合）

这是最关键的部分。系统不会盲目信任作弊小抄。它会检查过往案例与当前问题有多相似。

高置信度：如果过往案例与当前问题几乎完全相同，系统会说：“图书管理员可能对这个棘手问题不太确定，但我们的作弊小抄指出完美专家是 X。”随后，它将图书管理员的猜测与小抄的建议进行混合，并 heavily 依赖小抄。
低置信度：如果当前问题独一无二，且作弊小抄中没有任何匹配良好的案例，系统会说：“小抄在这里没用；我们可能只会增加噪音。”此时它会忽略查找结果，转而信任图书管理员原本冻结的决策。

为何这很重要

本文声称，这种方法处于两个极端之间的“甜蜜点”：

无所作为（零样本）：仅使用冻结的图书管理员。这很快，但在面对困难或新问题时会失败。
重新训练（监督微调）：从头开始教导图书管理员新规则。这效果很好，但极其缓慢、昂贵，且每次面对新任务都需要重做整个过程。

kNN-MoE 在不付出高昂代价的情况下，获得了重新训练带来的性能提升。这就像给图书管理员提供了一份动态的、可搜索的“最佳过往操作”记忆，而不是强迫他们重返校园。

论文的关键发现

它在难题上有效：当原始图书管理员感到困惑（高“困惑度”）时，该系统提供的帮助最大。当图书管理员已经很有信心时，系统则会退后，以免弄巧成拙。
少即是多：令人惊讶的是，仅查看一个过往示例（即最接近的单个匹配项）的效果，优于对许多示例取平均值。论文指出，对于专家路由而言，拥有太多来自过去的“意见”实际上会稀释信号。
速度与准确性：构建这份“作弊小抄”（离线进行）比重新训练整个模型要快得多。在实际回答阶段，它仅增加极小的延迟（约慢 3-4%），但显著提高了在医学考试和编程等困难任务上的准确性。

局限性（注意事项）

论文指出，该系统需要一组带标签的“参考集”数据来构建作弊小抄。如果你处于一种完全没有相似过往案例可供学习的境地，这种方法就无法提供帮助。它依赖于这样一个假设：“过去行之有效的方法”是“现在行之有效的方法”的良好指南。

简而言之，kNN-MoE 是一种让 AI 专家变得更聪明的方法，它允许它们在安全的情况下，窥探自己完美决策的历史记录。

技术摘要：kNN-MoE

问题陈述

混合专家（MoE）架构通过利用参数化“路由器”将令牌分发给稀疏的专家子集，从而高效地扩展大语言模型（LLM）。然而，标准路由器是轻量级分类器，仅训练一次后即被冻结。虽然它们在分布内数据上表现有效，但在面对分布偏移（例如特定领域或分布外任务）时，标准路由器在推理过程中缺乏灵活性。在此类场景中，路由器必须基于固定的先验进行外推，往往导致专家选择次优且困惑度（perplexity）较高。现有的推理时适应方法（如 C3PO）试图在线微调路由 logits，但会带来显著的计算开销（约为基线的 5 倍）。其他方法则需通过监督微调（SFT）进行永久性参数更新，这不仅计算密集，而且在推理效率上较低。

方法论：kNN-MoE

作者提出了 kNN-MoE，这是一种检索增强的路由框架，旨在在不修改模型参数的前提下优化专家分配决策。该方法分为两个 distinct 阶段：

1. 记忆构建（离线）

系统为每个 MoE 层构建一个非参数化记忆存储 $\mathcal{M}$ ，由键值对 $(k_t, v_t)$ 组成：

键（ $k_t$ ）： 来自参考数据集 $D_{ref}$ 的路由器输入（隐藏状态） $x_t$ 。
值（ $v_t$ ）： 通过优化路由 logits 以最大化真实令牌 $y_t$ $y_{t}$ 似然性而得出的“局部最优”专家分配。
- 该方法不存储冻结路由器的输出，而是针对 $D_{ref}$ 中的每个令牌求解一个优化问题：在给定冻结网络参数的情况下，寻找一个能最小化目标令牌负对数似然的 logits 向量 $r^*_t$ 。
- 该优化通过从原始参数化 logits 开始的少量梯度下降步骤（通常 $S=1$ ）执行。
- 由此产生的优化 logits $r^{(S)}_t$ 作为记忆中的值 $v_t$ 。

2. 置信度感知自适应混合（推理）

在推理过程中，对于给定的路由器输入 $x$ ：

检索： 系统基于 $x$ 与存储键之间的相似度，从记忆 $\mathcal{M}$ 中检索 $K$ 个最近邻 $\mathcal{N}(x)$ 。
非参数化提议： 通过聚合检索到的最优 logits 并以其与输入的相似度为权重，计算基于记忆的 logits 提议 $r_{mem}(x)$ 。
置信度驱动的混合： 最终的路由 logits 是原始参数化路由器 logits $r(x)$ $r (x)$ 与基于记忆的提议 $r_{mem}(x)$ $r_{m e m} (x)$ 之间的线性插值：
$r_{final}(x) = (1 - \lambda(x))r(x) + \lambda(x)r_{mem}(x)$
混合系数 $\lambda(x)$ $λ (x)$ 定义为检索到的邻居的平均相似度。这充当了一个置信度分数：
- 高相似度（ $\lambda \approx 1$ ）： 模型信任检索到的局部最优分配。
- 低相似度（ $\lambda \approx 0$ ）： 模型回退到原始冻结路由器，避免无关案例带来的噪声。
输出： 最终 logits 通过 Top-K Softmax 生成稀疏专家分配，保持模型原有的稀疏性约束。

主要贡献

检索增强路由： 引入了一种框架，利用离线计算的、来自相似过往案例记忆中的局部最优专家分配，以改进路由决策。
非参数化适应： 与 SFT 或在线适应方法不同，kNN-MoE 不更新模型参数。它将计算负担转移到离线记忆构建阶段，从而实现高效的推理。
置信度驱动的回退： 该方法利用检索邻居的平均相似度作为动态混合系数，确保在找不到相关历史案例时回退到冻结路由器，从而增强鲁棒性。
效率： 该方法在记忆构建期间每个令牌仅需单次梯度下降步骤，并避免了在线优化或全模型微调的高昂计算成本。

实验结果

作者在三个 MoE 模型（OLMoE、GPT-OSS 和 Qwen3）上评估了 kNN-MoE，涵盖了六个基准测试（GPQA、SuperGPQA、MMLU、USMLE、MedMCQA 和 MBPP）。

性能与基线对比： kNN-MoE 在所有模型和数据集上始终优于 Zero-shot 基线和“SFT（仅路由器）”基线。
与 SFT 的竞争力： 该方法取得了与全监督微调（SFT）相媲美的性能，并在多个数据稀缺场景（如 GPQA、USMLE）中超越了 SFT，后者在小型参考集上偶尔会出现过拟合问题。
稳定性： 与因输入长度问题导致某些模型（GPT-OSS、OLMoE）性能下降的 5-shot 提示不同，kNN-MoE 表现出稳定的提升。
消融研究：
- 邻居数量（ $K$ ）： 与典型的 kNN-LM 发现相反，性能在 $K=1$ 时达到峰值。聚合更多邻居（ $K>1$ ）会引入噪声，表明真正有益的路由信号是罕见且具体的。
- 优化步数（ $S$ ）： 将梯度下降步数增加至 $S=1$ 以上带来的精度提升微乎其微，却增加了构建成本。
- 参考集大小： 性能随参考集 $|D_{ref}|$ 的大小单调提升，但在 0.5k–1k 个示例后观察到收益递减。
效率： 构建 1k 个示例的记忆耗时 0.46 小时（而 SFT 约为 1.34 小时）。IVFPQ（倒排文件 + 乘积量化）变体将推理延迟开销降低至约 3.67%，并将 VRAM 使用量增长控制在约 9.43%（相比基线）。

意义与主张

该论文声称，kNN-MoE 弥合了静态参数化路由与动态推理需求之间的差距。其主要意义在于提供了一条可扩展的路径以增强 MoE 的适应性，而无需承担微调的计算负担。

作者强调，当参数化路由器不确定时（与高困惑度相关），该方法尤为有效，因为检索机制能从相似的高质量历史案例中提供校正信号。虽然该方法依赖于分布与测试数据相似的标注参考集，但它提供了一种实用的非参数化替代方案，以替代昂贵的适应技术，在显著降低准备时间且无需参数更新的情况下，实现了与 SFT 相媲美的结果。