arXivDaily arXiv每日学术速递 周一至周五更新
​通知:联邦节假日调休:哥伦布日 / 原住民日(Columbus Day / Indigenous Peoples' Day),今日(13日)无更新。
arXiv 2601.02144cs.CLcs.AI

类比路由:用于混合专家模型的kNN增强专家分配

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

Institute of Science Tokyo(东京科学研究院);CyberAgent;Nara Institute of Science and Technology(奈良科学技術大學)

Boxuan Lyu, Soichiro Murakami, Hidetaka Kamigaito, Peinan Zhang

AI总结:

提出kNN-MoE框架,通过检索历史相似案例的局部最优专家分配来增强MoE路由,使用检索邻居的平均相似度作为置信度混合系数,在分布偏移下提升鲁棒性。

摘要:

混合专家(MoE)架构通过使用参数化“路由器”将令牌分派给稀疏的专家子集,高效地扩展大型语言模型。通常,该路由器被训练一次然后冻结,导致路由决策在分布偏移下变得脆弱。我们通过引入kNN-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 locally 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 average 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 that kNN-MoE outperforms the zero-shot baseline and is competitive with computationally intensive supervised fine-tuning.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑