通过类比路由:kNN增强的专家分配用于Mixture-of-Experts
计算与语言
2026-05-26 v2 人工智能
摘要
Mixture-of-Experts (MoE) 架构通过采用参数化\'router\'将 token 分配到稀疏子集中的专家,从而高效扩展大型语言模型。通常情况下,该 router 在训练后被冻结,导致在分布迁移下路由决策变得脆弱。我们通过引入 kNN-MoE,一种基于检索增强的路由框架,来解决这一限制。该框架复用记忆库中类似过去案例的局部最优专家分配。记忆库通过直接优化 token 级别的路由 logits 以最大化参考集合上的似然概率来构建。关键在于,我们使用检索到的邻居平均相似度作为置信驱动的混合系数,从而允许方法在未找到相关案例时回退到冻结的 router。实验表明,kNN-MoE 在零样本基线上优于后者,而且在计算密集的监督微调方面表现出竞争力。
引用
@article{arxiv.2601.02144,
title = {Routing by Analogy: kNN-Augmented Expert Assignment for Mixture-of-Experts},
author = {Boxuan Lyu and Soichiro Murakami and Hidetaka Kamigaito and Peinan Zhang},
journal= {arXiv preprint arXiv:2601.02144},
year = {2026}
}