通过可微分损失让MoE模型减少专家切换，提升边缘设备推理效率。
Sticky Routing: Training MoE Models for Memory-Efficient Inference
- 引入可微分的路由一致性损失，抑制相邻词元间频繁切换专家。
- 专家切换率降低60%，困惑度仅下降4%以内，效果优于后期优化方法。
- 适合部署在内存受限的边缘设备上的MoE模型训练。
Mixture-of-Experts（MoE）模型每令牌仅激活少量专家，但连续令牌常触发不同专家，导致边缘设备中缓慢存储与快速内存间频繁权重交换。现有方案或为系统级（缓存启发式），或为后期修复（路由器微调），未改变预训练阶段的根本问题。本文提出StickyMoE，一种可微分的路由一致性损失，惩罚相邻令牌间的突变专家切换，促使路由器在语义连贯片段内保持相同专家分配。该方法无需架构改动，仅增加一个超参数λ，且相比后期方法，使专家表示与路由决策从训练初始即可协同优化。小规模MoE语言模型实验表明，StickyMoE将专家切换率最高降低60%，困惑度降幅不足4%，在质量-局部性权衡前沿上优于后处理微调，证明路由时间局部性应在训练阶段高效植入。
原文摘要 · Abstract (English)
Mixture-of-Experts (MoE) models activate only a sparse subset of experts per token, yet consecutive tokens frequently activate different experts -- causing constant weight swapping between slow storage and fast memory on edge devices. Existing remedies are either system-level (caching heuristics) or post-hoc (router fine-tuning), leaving the root cause unchanged during pretraining. We propose StickyMoE, a differentiable routing consistency loss that penalises abrupt expert switches between adjacent tokens, encouraging the router to maintain the same expert assignment across semantically coherent spans. StickyMoE requires no architectural changes, adds a single hyperparameter lambda, and unlike post-hoc methods, allows expert representations and routing decisions to co-adapt from the first training step. Experiments on small-scale MoE language models show that StickyMoE reduces the expert switch rate by up to 60% with less than 4% perplexity degradation, Pareto-dominating post-hoc fine-tuning on the quality-locality frontier. Routing temporal locality is most efficiently instilled at training time.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。