混合专家（MoE）模型通过稀疏激活来提升效率，但在边缘设备上面临一个关键挑战：相邻的token往往会激活不同的专家，导致模型权重需要在慢速存储（如闪存）和快速内存之间频繁交换。这种权重交换不仅增加了推理延迟，还造成了严重的内存瓶颈，限制了MoE模型在手机、物联网设备等资源受限平台上的部署。
现有的解决方案主要分为两类：系统级的缓存启发式策略，以及事后的路由器微调。然而，这些方法都未能从根源上解决问题——训练过程中的路由不稳定性。缓存策略只能缓解症状，而事后微调虽然可以改善局部性，但往往会损害模型质量，并且需要额外的训练步骤。
来自arXiv的新研究提出了StickyMoE，这是一种可微分的路由一致性损失函数。其核心思想是在训练过程中直接惩罚相邻token之间的专家切换，鼓励路由器在语义连贯的文本跨度内保持相同的专家分配。这种方法不需要改变模型架构，仅仅引入了一个超参数λ来控制惩罚的强度。更重要的是，与事后微调不同，StickyMoE允许专家表示和路由决策从训练的第一步起就协同适应，从而在保持模型质量的同时，自然地学习到具有时间局部性的路由策略。
研究者在小型MoE语言模型上进行了实验，结果显示StickyMoE能够将专家切换率降低最多60%，同时仅带来不到4%的困惑度下降。在质量-局部性权衡曲线上，StickyMoE全面优于事后微调方法，实现了帕累托改进。这一工作强调了在训练阶段注入路由时间局部性是最有效的方式，为未来在边缘设备上高效部署MoE模型提供了新的思路。随着边缘AI的快速发展，StickyMoE有望成为推动MoE模型落地的重要技术之一。