🤖 machine learning

Sticky Routing: Training MoE Models for Memory-Efficient Inference

该论文提出了 StickyMoE，一种可微的路由一致性损失函数，旨在训练混合专家模型在语义连贯的标记跨度内保持专家分配，从而在显著减少边缘设备上内存密集型权重交换的同时，实现极小的困惑度下降。

原作者： Ali Kayyam

发布于 2026-07-13

📖 1 分钟阅读☕ 轻松阅读

原作者： Ali Kayyam

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你拥有一个由 64 位不同领域的专家图书管理员组成的巨大图书馆，但你的微型边缘设备（比如手机或智能音箱）在桌面上只能腾出足够的空间来放置 2 位图书管理员。剩下的都被存放在一个缓慢且尘土飞扬的地下室里。

在标准的“混合专家”（Mixture-of-Experts, MoE）模型中，每当 AI 阅读一个新词时，它都会抛一次硬币来决定询问哪位图书管理员。问题在于？它为每一个词都抛一次硬币。所以，它在读到“苹果”时询问图书管理员 A，读到“派”时询问 B，读到“饼皮”时询问 C。每当它切换时，它都必须跑下楼，把图书管理员 B 踢出去，再把图书管理员 C 拉上来。这种“来回奔波”非常缓慢，以至于 AI 花在取回图书管理员上的时间比阅读故事本身还要多。

这篇论文介绍了一种名为 StickyMoE 的新方法，旨在训练这些模型，让它们不再如此疯狂地变换决策。通过这种方法，模型学会了变得“粘性”（sticky）。如果它在读到“苹果”时询问了图书管理员 A，那么只要主题合理，它就会被鼓励在读到“派”和“饼皮”时也继续询问图书管理员 A。

核心发现：要在训练时解决，而不是事后修补

作者尝试了几种方法来解决这个“来回奔波”的问题，以下是他们的发现：

“事后”修补失败了： 他们尝试拿一个已经训练好的模型，仅仅通过“微调”那个抛硬币的机制（路由），使其不再那么随机。他们称之为 ReMoE。结果是：没用。 切换率几乎没有变化（绝对变化量小于 0.5%）。作者认为这是因为图书管理员（专家）已经被训练成了针对特定、孤立词汇的专家。如果你想让一个“硬币投掷器”变得更具“粘性”，但专家本身还没准备好处理长期的连续话题，那是不行的。这就像试图通过只改变短跑运动员的起跑姿势，来教他成为一名马拉松选手；肌肉（专家权重）是为短跑爆发而生的，而非长跑。
“硬性”修复过于僵化： 他们还尝试了一种方法，即使用严格的规则强制模型在一定数量的词语内固定使用一位图书管理员（“硬窗口”法）。虽然这减少了切换，但也让 AI 的回答质量大幅下降（使“困惑度”上升了高达 6.8%，针对中型模型）。这就像是强迫一位图书管理员去读完一整本小说，即使主题已经完全转向了另一种类型；故事的质量因此受损。
“粘性”方案奏效了： 获胜的方法 StickyMoE 在训练的第一步加入了一个温和的“一致性惩罚”。它告诉模型：“嘿，如果你在两个相邻的词之间更换了专家，那会受到一点惩罚。”这鼓励模型在学习语言的同时，也学会了保持连贯性。

结果：具有“粘性”更好

当他们在 WikiText-2 数据集（维基百科文章集合）上对小型和中型模型进行测试时，结果非常明确：

更少的切换： 通过调节一个单一的旋钮（称为 λ 的数值），他们可以将模型更换专家的次数减少高达 59%。例如，在中型模型上，切换率从 0.71（几乎每次都在切换）降到了 0.29。
更好的速度（模拟）： 由于模型切换减少了，他们模拟了一个可以容纳 2 位专家的缓存（桌面架子）。“缓存命中率”（即正确的图书管理员已经在桌上的频率）从 0.54 跳升至 0.88。这意味着模型需要跑下楼去地下室的次数比之前减少了 3.92 倍。
更好的质量： 令人惊讶的是，在中型模型上，使用适度的“粘性”（λ = 0.05）实际上让 AI 变得更聪明了，它将“困惑度”（衡量混乱程度的指标）降低了 4.1%。模型意识到，保持对主题的专注有助于它更好地理解上下文。
没有崩溃： 一个巨大的担忧是模型会变得“懒惰”，从而只使用一位图书管理员。但“利用熵”（Utilisation Entropy）保持在高位（在最大值 2.0 bits 中的 1.92 bits 以上），这意味着所有 4 位专家仍在使用中。模型并没有变懒，它只是变得更有组织了。

什么行不通（以及为什么）

论文明确排除了在模型构建完成后再解决此问题的想法。ReMoE 方法（对旧模型进行微调）被证明是无效的，因为“专家表示”（图书管理员的大脑）已经由一个并不关心是否能持续关注主题的训练过程所塑造。你无法后期补救“局部性”，你必须从一开始就将其植入其中。

他们还指出，模型的第一个层（第 0 层）是最难实现“粘性”的。即使使用最好的设置，该层的切换率仍然高于 0.49。这是因为第一层看到的是原始词汇，在它们与上下文融合之前，它们还没有被混合。因此，“苹果”和“派”在这一阶段看起来可能确实截然不同。

总结

论文表明，如果你希望你的 AI 在内存有限的设备上运行得更快，你不应该仅仅尝试用聪明的缓存规则，或者试图在模型建成后再去修复它。相反，你应该从第一天起就训练模型的“粘性”。通过对相邻词语之间频繁更换专家施加简单的惩罚，你会得到一个天然能保持主题一致、需要更少往返于缓慢地下室、并且能更好地理解故事的模型。

作者承认，这是基于模拟以及在特定数据集（WikiText-2）和中小规模模型（最高达 22M 参数）上的训练。他们暗示，随着模型规模的增大，这种“粘性”可能会表现得更好，但尚未在数十亿参数的大型模型上进行测试。他们还注意到，目前该方法并不知道句子何时结束，因此它可能会在段落分隔处表现得“过于粘性”，这是未来研究需要解决的问题。

技术摘要：用于内存高效 MoE 推理的粘性路由（Sticky Routing）

1. 问题陈述：专家切换瓶颈

混合专家（Mixture-of-Experts, MoE）模型旨在通过在每个 token 时仅激活一小部分专家子集来扩展大语言模型（LLM）的规模。理论上，这使得在拥有海量参数的同时保持较低的单 token 计算量，使其在边缘部署场景中极具吸引力——因为在任何给定时刻，只需要将少数专家驻留在快速内存（VRAM/SRAM）中。

然而，本文指出在实践中存在一个关键的失效模式：时间路由不一致性（temporal routing inconsistency）。标准的 MoE 路由根据每个 token 的隐藏表示独立做出决策。因此，连续的 token 经常会激活完全不同的专家。在内存受限的硬件上，每一次切换都需要从慢速存储（CPU RAM/Flash）中驱逐缓存的专家权重并加载新权重。由于 PCIe/NVMe 带宽比 GPU 显存带宽低几个数量级，这些频繁的“专家切换”主导了推理延迟，往往导致 MoE 模型比同等规模的稠密模型运行得更慢。

现有的解决方案被认为不足以解决该问题：

系统级方法（例如缓存启发式算法、预取器）：是响应式的；它们利用已有的局部性，但无法修复本质上不一致的路由器。
事后处理方法（例如仅微调路由器）：之所以失败，是因为专家的表示是由从未激励过局部性的训练机制塑造的。路由器无法在并非为此设计的表示空间内轻松重新对齐决策。
架构重设计（例如 Oracle-MoE）：通常需要对模型的结构或训练流程进行重大更改。

2. 方法论：StickyMoE

作者提出了 StickyMoE，这是一种将路由时间局部性视为显式目标的训练时方法。其核心创新是一个可微路由一致性损失（differentiable routing consistency loss, $L_{cons}$ ），该损失会对相邻 token 之间专家分配的突变进行惩罚。

路由一致性损失

该损失函数最小化连续 token 的 Softmax 门控概率向量（ $g_t$ ）之间的 $\ell_2$ 距离：
$L_{cons} = \frac{1}{T-1} \sum_{t=2}^{T} \| g_t - g_{t-1} \|_2^2$
其中 $g_t \in \mathbb{R}^N$ 是第 $t$ 个 token 的路由输出。

该损失与标准的交叉熵损失（ $L_{CE}$ ）和负载均衡损失（ $L_{bal}$ ）一起集成到总训练目标中：
$L = L_{CE} + \lambda L_{cons} + \mu L_{bal}$

关键机制

协同进化（Co-adaptation）：与事后处理方法不同，StickyMoE 同时训练路由器和专家表示。最小化 $L_{cons}$ 会更新路由器权重（ $W_r$ ）并通过隐藏状态（ $h_t$ ）进行反向传播，从而鼓励模型学习在语义连贯的区间内具有自然一致性的表示。
软-硬变体（Soft-Hard Variant）：为了解决“长程漂移”（即微小的步进式变化累积成远距离的切换），作者引入了一个锚点约束。序列被划分为若干窗口；窗口的第一个 token 设置为一个锚点，后续 token 被惩罚偏离该锚点，且惩罚权重随距离窗口起始点的距离线性增加。
兼容性：该方法无需架构更改，且增加的计算开销微乎其微。它兼容任何标准的 top- $k$ MoE 以及现有的负载均衡技术。

3. 核心贡献

路由一致性损失：一种简单的、可微的训练目标，直接惩罚相邻 token 之间的专家切换，仅需一个超参数（ $\lambda$ ）。
与负载均衡的互补性：实证证据表明，一致性损失与负载均衡损失是互补的。一致性损失在不引起“专家崩溃”（利用熵保持在高位，接近理论最大值）的情况下，塑造了专家使用的时机。
软-硬变体：结合了逐步软惩罚和段级锚点约束，在不增加质量成本的前提下保证了极强的局部性。
全面评估：在小型和中型 MoE 模型上，通过受控实验将 StickyMoE 与原始 MoE、事后路由器微调（ReMoE）以及架构基准（Oracle-MoE）进行了对比。

4. 实验结果

实验在 WikiText-2 数据集上的小型（约 8.8M 参数）和中型（约 22M 参数）MoE 模型上进行。

专家切换率 (SR)：StickyMoE 将专家切换率降低了高达 59%（例如，在中型模型上从 0.71 降至 0.29）。
缓存命中率 (CHR)：切换率的降低转化为缓存缺失的显著减少。在中型模型使用 2 个专家槽位的缓存配置下，缓存缺失减少了 3.92 倍（CHR 从 ~0.54 提高到 ~0.88）。
困惑度（质量）：
- 在适度的 $\lambda$ 值（如 0.05）下，StickyMoE 比基准模型改善了困惑度（例如，在中型模型上提升了 -4.1%），这表明一致性损失起到了有益的正则化作用。
- 过大的 $\lambda$ 值会引入权衡，但 StickyMoE 的帕累托前沿（Pareto frontier）严格优于其他方法。
事后微调的失败：ReMoE 基准（仅在预训练模型上微调路由器）未能显著降低切换率（变化 <0.5%），这证实了如果表示已经固定，则无法追溯性地实现局部性。
层级分析：一致性损失在深层（L1–L3）最为有效。输入层（L0）对局部性改进表现出抵抗性，这可能是因为 token 嵌入缺乏上下文混合。

5. 重要性与主张

论文认为，路由时间局部性是一个合法的训练目标，应当在预训练阶段就植入其中，而不是将其视为系统层面的事后补救。

范式转变：正如量化感知训练（quantization-aware training）已成为标准一样，作者断言，“局部性感知训练”应当成为面向内存受限部署的 MoE 模型的默认选择。
根源解决：StickyMoE 解决了根本原因（不一致的路由决策），而非症状（缓存缺失）。通过协同训练表示和路由决策，它避免了使事后方法陷入困境的表示失配问题。
实用性：该方法与架构无关，且不需要更改推理流水线。训练后的路由器自然产生一致的门控分布，使得标准的缓存策略（如 LRU, LFU）能够在无需额外工程投入的情况下显著提高效率。

作者总结道，虽然系统级优化和事后调整仍然是有效的补充，但它们无法取代一个从未被训练为“缓存友好”的模型。StickyMoE 为在边缘设备上实现 MoE 架构的理论内存效率提供了一条路径。