StickyMoE粘性路由训练法 让MoE模型内存高效
速览
MoE模型相邻token频繁切换专家导致边缘设备推理内存瓶颈。现有方法在系统或事后层面修补，未解决训练根源。StickyMoE提出可微分路由一致性损失，惩罚相邻token专家切换，鼓励路由保持跨语义区段的相同专家分配，无需改变架构。实验显示专家切换率降低60%，困惑度损失<4%，在质量-局部性前沿上全面优于事后微调。该方法将路由时间局部性高效融入预训练，为MoE模型在资源受限设备上部署铺平道路。
AI 深度解读
背景
混合专家模型（Mixture-of-Experts, MoE）通过为每个 token 仅激活一小部分专家（expert），显著降低了推理时的计算量。然而，在边缘设备（如手机、IoT 设备）上，由于内存资源有限，专家权重通常存储在慢速存储（如闪存）中，只有当前激活的专家被加载到快速内存（如 RAM）中。问题在于：相邻的 token 往往激活不同的专家，导致频繁的权重交换（weight swapping）——即从慢速存储中加载新专家的权重并卸载旧专家的权重。这种 I/O 开销严重拖慢了推理速度，尤其对于实时应用而言。
现有的缓解方法分为两类：系统层面的缓存启发式策略（如 LRU 缓存），以及后处理（post-hoc）的路由器微调（router fine-tuning）。但这些方法都没有在预训练阶段改变问题的根源——路由决策本身缺乏时间局部性（temporal locality），即连续 token 倾向于激活不同专家。
核心内容
本文提出 StickyMoE，一个可微分的路由一致性损失函数（differentiable routing consistency loss），它在预训练过程中直接惩罚相邻 token 之间的专家切换行为。核心思想是：鼓励路由器（router）在语义上连贯的文本片段（semantically coherent spans）内保持相同的专家分配，从而减少不必要的权重交换。
StickyMoE 不需要修改 MoE 的架构，只需添加一个超参数 λ（用于控制一致性损失的权重）。与后处理方法不同，StickyMoE 从训练的第一步开始就让专家表示（expert representations）和路由决策（routing decisions）共同适应，因此路由的时间局部性得以在训练阶段被有效灌输。
实验在小型 MoE 语言模型上进行。结果表明：StickyMoE 将专家切换率（expert switch rate）降低了最高 60%，同时困惑度（perplexity）的退化不超过 4%。在质量-局部性帕累托前沿（quality-locality frontier）上，StickyMoE 全面优于后处理的路由器微调方法。作者强调，路由的时间局部性在训练时注入最为高效。
关键要点
- 问题本质：MoE 模型中的连续 token 经常激活不同专家，导致边缘设备上频繁的权重交换，造成严重的推理延迟。
- 现有方案局限：系统级缓存启发式和后处理路由器微调均未在预训练阶段解决路由缺乏时间局部性的根本原因。
- StickyMoE 方法：引入可微分的路由一致性损失函数，惩罚相邻 token 之间的专家切换，鼓励在语义连贯的文本片段内保持相同的专家分配。
- 无需架构改动：该方法不改变 MoE 模型的基本结构，仅添加一个超参数 λ 控制惩罚强度。
- 训练阶段协同适应：与后处理方案不同，StickyMoE 允许专家表示和路由决策从训练伊始便共同优化，从而更高效地建立路由的时间局部性。
- 实验效果：在小规模 MoE 语言模型上，专家切换率降低高达 60%，困惑度损失小于 4%，在质量-局部性权衡上优于后处理方法。
- 核心结论：路由的时间局部性在预训练阶段注入最为有效，而非训练后再调整。
意义与影响
StickyMoE 提供了一种轻量级、端到端可训练的解决方案，使得 MoE 模型在内存受限的边缘设备上能够更高效地运行。通过降低权重交换频率，推理延迟和功耗均可得到优化，从而扩展 MoE 模型在移动端、物联网设备等场景中的适用性。
该方法具有普适性：可集成到任何已有的 MoE 训练流程中，无需改动模型架构，仅需调整损失函数。这降低了应用门槛，且与系统级优化（如缓存策略）正交，可以叠加使用以获得更好的效果。
从研究角度看，StickyMoE 揭示了在预训练阶段显式塑造路由行为（而非仅仅依赖隐式学习）的可行性。未来工作可以探索更细粒度的语义一致性（如基于句子或段落边界）以及如何在大规模模型上保持有效性。此外，本文提出的质量-局部性帕累托前沿评估方式，为未来比较不同内存高效 MoE 方法提供了新基准。