粘性路由：训练MoE模型以实现内存高效推理

# 粘性路由：训练MoE模型以实现高效内存推理
来源：https://arxiv.org/html/2607.08780

###### 摘要

混合专家（MoE）模型每个token仅激活稀疏的专家子集，然而连续token频繁激活*不同的*专家——导致边缘设备上慢速存储与快速内存之间持续进行权重交换。现有的补救措施要么是系统级的（缓存启发式策略），要么是事后性的（路由器微调），在预训练期间未改变根本原因。我们提出**StickyMoE**，一种可微的*路由一致性损失*，该损失惩罚相邻token之间的突然专家切换，鼓励路由器在语义连贯的跨度内维持相同的专家分配。**StickyMoE**无需架构更改，仅添加一个超参数λ，并且与事后方法不同，它允许专家表示和路由决策从第一步训练开始就共同适应。在小型和中等规模MoE语言模型上的实验表明，**StickyMoE**将专家切换率降低了高达59%，同时在中等规模模型上*提高了*困惑度，并将缓存未命中率降低了高达**3.92×**，在质量-局部性前沿上帕累托优于事后微调。路由时间局部性最有效的方式是在训练时注入。代码地址：https://github.com/alikayyam/sticky_moe.git

## 1 引言

混合专家（MoE）架构已成为扩展大型语言模型（LLMs）同时不按比例增加每个token计算量的最重要实用技术之一(Shazeer 等人., 2017 (https://arxiv.org/html/2607.08780#bib.bib1); Fedus 等人., 2022 (https://arxiv.org/html/2607.08780#bib.bib2); Lepikhin 等人., 2021 (https://arxiv.org/html/2607.08780#bib.bib3))。通过将密集的前馈网络（FFN）子层替换为N个专家FFN的集合和一个学习到的稀疏路由器，MoE模型可以拥有数十亿参数，同时仅为每个输入token激活一小部分——通常是8-128个专家中的1或2个。这一特性使得MoE架构在理论上对边缘部署具有吸引力：如果任何时候快速内存中只需要驻留64个专家中的2个，那么有效工作集大小仅是总参数数量的一小部分。

#### 专家交换问题。

在实践中，这一承诺因标准MoE路由器的*时间不一致性*而受到损害。因为路由器仅基于每个token的隐藏表示为该token做出独立的路由决策，序列中的连续token通常会激活完全不同的专家(Zhou 等人., 2025 (https://arxiv.org/html/2607.08780#bib.bib7))。在GPU VRAM或设备上SRAM有限的系统上，每次这样的切换都需要驱逐当前缓存的专家权重，并从CPU RAM或闪存加载新的专家权重。PCIe和NVMe带宽比GPU内存带宽低一到两个数量级；因此，单次专家权重传输可能会主导前向传播本身的延迟(Xue 等人., 2024 (https://arxiv.org/html/2607.08780#bib.bib10); Zhu 等人., 2025 (https://arxiv.org/html/2607.08780#bib.bib13))。结果是，尽管具有理论上的稀疏性优势，MoE模型在内存受限硬件上可能比相同规模的密集模型*更慢*。

#### 为什么现有解决方案不足。

先前的工作从两个方面解决这个问题。*系统级*方法在快速内存中维护一个专家缓存，并使用驱逐策略（LRU、LFU）或学习到的预取器来降低缓存未命中率(Xue 等人., 2024 (https://arxiv.org/html/2607.08780#bib.bib10); Kamahori 等人., 2024 (https://arxiv.org/html/2607.08780#bib.bib11); Zhang 等人., 2025 (https://arxiv.org/html/2607.08780#bib.bib12))。虽然实际有用，但这些方法本质上是反应性的：它们试图利用模型路由器恰好产生的任何局部性，但无法改善本质上不一致的路由器。*事后*方法仅微调已训练模型的路由器以增加专家复用(Zhu 等人., 2026 (https://arxiv.org/html/2607.08780#bib.bib8))，这种方法更有针对性，但仍然在治标而非治本：专家表示已经由从未激励局部性的训练机制塑造，而短暂的微调过程只能部分重新对齐路由表面。第三项工作，最接近的是Oracle-MoE(Zhou 等人., 2025 (https://arxiv.org/html/2607.08780#bib.bib7))，它重新设计了路由架构，使得路由器在注意力推导的“神谕空间”中运行，该空间在token间固有地更具语义稳定性。Oracle-MoE从头开始训练，这是朝着正确方向迈出的一步，但局部性属性是结构性的——嵌入在路由输入的选择中——而不是显式优化的。我们提出一种不同的方法：**StickyMoE**，它直接将路由局部性作为训练目标注入，无需架构更改和事后纠正。

#### 我们的提案。

我们认为最简洁、最通用的解决方案是将路由时间局部性作为显式的训练目标。具体来说，我们引入了一个*路由一致性损失*——对连续门控分布施加一个可微的l2惩罚——以直接优化MoE路由中的时间局部性：

L_cons = 1/(T-1) ∑_{t=2}^T ||g_t - g_{t-1}||_2^2, (1) 其中 g_t ∈ R^N 是路由器为token t 生成的softmax门控概率向量，T 是序列长度。总训练损失为

L = L_CE + λ L_cons + μ L_bal, (2) 其中 L_CE 是标准交叉熵语言建模损失，L_bal 是来自 Fedus 等人 (2022 (https://arxiv.org/html/2607.08780#bib.bib2)) 的负载均衡辅助损失，λ, μ ≥ 0 是标量超参数。这种形式：

- • 需要*无需架构修改*——它适用于任何标准的 top-k MoE，无需更改路由器结构；
- • 是*处处可微的*，并且增加的计算开销可忽略；
- • 是*架构无关的*，同样适用于小型实验模型和大规模MoE；
- • 与语言建模目标*共同训练*，使得专家表示和路由决策从第一次更新开始就朝着局部一致性共同进化。

#### 贡献。

我们做出以下具体贡献：

1. 1. 我们提出了路由一致性损失，一个简单的可微训练目标，直接惩罚相邻token之间的专家切换，无需架构更改，仅添加一个超参数λ（第4节 (https://arxiv.org/html/2607.08780#S4)）。
2. 2. 我们通过实验表明 L_cons 和 L_bal 是互补的：在所有设置中，专家利用率熵保持在1.92比特以上（最高为 log_2 4 = 2.0 比特），确认一致性损失不会导致专家坍缩（第5节 (https://arxiv.org/html/2607.08780#S5)）。
3. 3. 我们提出一个*软-硬*变体，通过将逐步骤软惩罚与段级锚点约束相结合来解决长程路由漂移，在无额外质量成本的情况下实现更强的局部性保证（第4节 (https://arxiv.org/html/2607.08780#S4)）。
4. 4. 我们在WikiText-2上训练的小型和中等规模MoE语言模型上呈现了控制实验，与标准MoE基线、硬窗口消融、模拟ReMoE事后基线以及简化的Oracle-MoE重新实现进行比较，测量困惑度、切换率、缓存命中率和利用率熵（第5节 (https://arxiv.org/html/2607.08780#S5)）。

#### 论文组织。

第2节 (https://arxiv.org/html/2607.08780#S2) 回顾了MoE架构和专家交换瓶颈。第3节 (https://arxiv.org/html/2607.08780#S3) 调查了相关工作，并精确地定位了**StickyMoE**。第4节 (https://arxiv.org/html/2607.08780#S4) 详细介绍了该方法。第5节 (https://arxiv.org/html/2607.08780#S5) 描述了实验设置并呈现了结果。第6节 (https://arxiv.org/html/2607.08780#S6) 讨论了发现、局限性和更广泛的影响。第7节 (https://arxiv.org/html/2607.08780#S7) 得出结论并概述了未来方向。

## 2 背景

### 2.1 混合专家语言模型

一个标准MoE变换器将每个FFN子层替换为N个专家网络 {E_i}_{i=1}^N 和一个路由器 R。给定token t 的隐藏表示 h_t ∈ R^d，路由器计算门控logits l_t = W_r h_t ∈ R^N，并根据概率选择 top-k 个专家：

g_t = softmax(l_t), (3)
S_t = top-k(g_t), (4)
y_t = ∑_{i∈S_t} g_{t,i} E_i(h_t). (5)

路由器权重矩阵 W_r ∈ R^{N×d} 是相对于密集FFN层的唯一额外参数。专家网络 E_i 通常是两层FFN，结构与它们替换的密集FFN相同。

### 2.2 负载均衡辅助损失

没有明确的正则化，MoE路由器倾向于坍缩：一小部分专家接收了不成比例的大量token，而其余专家很少被激活且未能发展出有用的表示(Fedus 等人., 2022 (https://arxiv.org/html/2607.08780#bib.bib2))。为了防止这种情况，我们包含了标准的负载均衡辅助损失：

L_bal = N ∑_{i=1}^N f_i · p_i, (6)

其中求和跨越所有N个专家，f_i 是分配给专家i的top-k分配比例（计数所有k个槽位），p_i = 1/T ∑_{t=1}^T g_{t,i} 是分配给专家i的平均门控概率，两者都在当前批次上计算。当路由在专家间均匀分布时，乘积 f_i · p_i 最小化；乘以N确保损失幅度与专家数量无关。

### 2.3 内存受限硬件上的专家交换

设 M_fast 为快速内存容量（VRAM或设备上SRAM），w_e 为单个专家权重张量的大小。一个设备最多可以同时持有 ⌊M_fast / w_e⌋ 个专家权重矩阵，同时还需容纳注意力和嵌入参数。当解码步骤t的路由器选择一个当前不在快速内存中的专家时，系统必须驱逐一个缓存的专家并从慢速内存加载新的专家——即发生了*缓存未命中*。每个步骤的延迟为：

τ_t = τ_compute + 1[miss_t] · τ_load, (7)

在带宽受限设备上，τ_load >> τ_compute。经历缓存未命中的步骤的预期比例由*专家切换率*决定：

SR = 1/(T-1) ∑_{t=2}^T 1[S_t ≠ S_{t-1}]. (8)

这个单序列定义说明了概念；跨越层的完整经验估计量在第5节 (https://arxiv.org/html/2607.08780#S5) 中正式定义。因此，降低SR直接等同于降低内存约束下的预期推理延迟。

## 3 相关工作

### 3.1 系统级专家卸载

大量工作将专家权重视为两级内存层次结构，并专注于在推理时最小化从慢速内存加载专家到快速内存的延迟成本。MoE-Infinity (Xue 等人., 2024 (https://arxiv.org/html/2607.08780#bib.bib10)) 在个人机器上引入了稀疏性感知的专家缓存，将注意力权重保留在GPU内存中，并从主机RAM流式传输专家权重。Fiddler (Kamahori 等人., 2024 (https://arxiv.org/html/2607.08780#bib.bib11)) 使用CPU本身进行专家计算以避免PCIe传输。DuoServe (Zhang 等人., 2025 (https://arxiv.org/html/2607.08780#bib.bib12)) 使用学习到的专家激活路径预测模型分别优化预填充和解码阶段。EdgeMoE (Zhu 等人., 2025 (https://arxiv.org/html/2607.08780#bib.bib13)) 针对边缘场景提出了重要性驱动的专家调度。文献中的一个相关线索研究了已部署模型中的局部性：Liang 等人 (2025 (https://arxiv.org/html/2607.08780#bib.bib9)) 进行了一项实证研究，表明局部路由一致性在不同模型族之间差异巨大，并且某些模型从根本上不适合基于卸载的部署。

所有上述方法都是*推理时*干预。它们利用训练模型的路由器提供的任何局部性，但无法改变路由器的归纳偏置。**StickyMoE**是正交且互补的：它可以与这些系统中的任何一个结合使用，提供已经倾向于一致性的路由器。

### 3.2 事后路由器自适应

ReMoE (Zhu 等人., 2026 (https://arxiv.org/html/2607.08780#bib.bib8)) 是最直接相关的事后方法。它仅微调预训练MoE检查点的路由器参数，以增加短视野专家复用，结合使用语言建模损失和门控分布上的局部性感知正则化器。ReMoE报告了缓存未命中率的有意义降低和基于vLLM的服务的吞吐量提升。其关键优势在于实用性：它可以应用于任何预训练模型而无需从头重新训练。其关键局限性在于专家表示本身是固定的——由从未激励局部性的训练过程塑造——因此路由器被要求在一个并非为局部性设计的表示空间中改变其决策。**StickyMoE**通过从一开始就注入局部性偏置来避免这种不匹配，允许专家表示和路由决策共同适应。在我们的控制实验中，我们无法从事后路由器微调中重现有意义的SR降低；我们将此归因于第6节 (https://arxiv.org/html/2607.08780#S6) 中讨论的表示不匹配。

### 3.3 训练时架构重新设计

Oracle-MoE (Zhou 等人., 2025 (https://arxiv.org/html/2607.08780#bib.bib7)) 在动机和训练机制方面与我们的工作最为接近。它确定了相同的根本原因——时间路由不一致性——并提出了一种结构化解决方案：不是在标准的隐藏状态空间中对token进行路由，而是在一个从注意力分数推导出的“神谕空间”中进行路由，该空间在经验上对连续token更具语义稳定性。Oracle-MoE从头开始训练，并在不牺牲任务性能的情况下在边缘设备上实现了最先进的推理速度。然而，它需要非平凡的架构修改（用注意力推导的特征替换路由器输入）。