Blog Tags GitHub
Back to blog
mixture-of-experts diffusion-model inference-serving

Elastic MoE: Unlocking the Inference-Time Scalability of Mixture-of-Experts

通过随机共激活采样和层次化路由器损失，实现MoE模型推理时弹性扩展，有效扩展范围达训练时k的2-3倍

Elastic MoE: Unlocking the Inference-Time Scalability of Mixture-of-Experts

一、论文概述

项目内容
标题Elastic MoE: Unlocking the Inference-Time Scalability of Mixture-of-Experts
作者Naibin Gu, Zhenyu Zhang, Yuchen Feng, Yilong Chen, Peng Fu, Zheng Lin, Shuohuan Wang, Yu Sun, Hua Wu, Weiping Wang, Haifeng Wang
机构未明确标注（推测为百度/哈工大）
论文arXiv:2509.21892
发布2025-09-26 (v1), 2026-05-11 (v2)
领域计算与语言 (cs.CL)
投稿NeurIPS

二、核心思想

问题定义

MoE模型通常在训练和推理时固定激活专家数量 kk。然而，实际部署面临异构硬件、波动工作负载和多样化质量-延迟需求，为每种场景训练单独模型成本高昂。考虑到MoE模型 already operates with sparse activation，调整激活专家数量提供了一条自然路径来服务不同预算。

核心问题：能否通过在推理时激活更多专家来释放MoE模型的潜力？

核心发现

推理时扩展墙（Inference-Time Scaling Wall）：

标准MoE模型无法支持弹性推理。当模型用 kk 个专家训练时，推理时的有效扩展范围极其狭窄——将 k′k' 增加到略大于 kk 时，性能会迅速下降。

根本原因：专家之间缺乏 learned collaboration。通过专家共现矩阵分析发现，训练时和推理时的专家激活模式存在显著差异。

解决方案概述

Elastic MoE (EMoE) 是一个新颖的训练框架，使MoE模型能够在推理时弹性变化激活专家数量：

  1. 随机共激活采样（Stochastic Co-activation Sampling）：将训练预算 ktraink_{\text{train}} 与组合覆盖 kidealk_{\text{ideal}} 解耦，通过从更大候选池中采样来暴露模型于多样化的专家组合
  2. 层次化路由器损失（Hierarchical Router Loss）：鼓励路由器产生稳定、果断的专家排名，确保在不同推理预算下的可靠性能

核心结果：

  • 有效扩展范围达训练时 kk 的 2-3×
  • 性能随推理预算单调递增
  • 无需额外训练开销
  • 适用于4种MoE架构（7B-21B）和9个基准测试

三、技术架构

推理时扩展墙现象

推理时扩展墙

关键观察：

  • 当推理时激活专家数 k′k' 超过训练预算 kk 时，性能短暂保持后迅速下降
  • 更大的训练预算（如 k=6k=6）虽然移动了峰值，但计算成本 prohibitive
  • 回到原始推理预算时性能崩溃

专家共现分析

专家共现矩阵

专家共现矩阵 M(k)∈RN×NM^{(k)} \in \mathbb{R}^{N \times N}：

  • 训练时 k=2k=2 的专家共现模式与推理时 k′=6k'=6 存在巨大差异
  • 共现距离与模型性能强相关

EMoE架构

EMoE架构

EMoE vs 标准Top-k MoE：

组件标准Top-k MoEEMoE
专家选择固定top-k候选池 + 随机采样
训练预算ktraink_{\text{train}}ktraink_{\text{train}}（相同）
组合覆盖仅top-k组合多样化组合（kidealk_{\text{ideal}}）
路由器损失标准 + 负载均衡+ 层次化损失 LHR\mathcal{L}_{\text{HR}}
推理弹性无2-3× 扩展范围

随机共激活采样

核心思想：将训练预算与组合覆盖解耦。

算法流程：

  1. 确定候选池大小： k~ideal∼UniformInt(ktrain,kideal)\tilde{k}_{\text{ideal}} \sim \text{UniformInt}(k_{\text{train}}, k_{\text{ideal}})

  2. 选择候选专家： Sk~ideal(x)←TopKIndices(h(x),k~ideal)\mathcal{S}_{\tilde{k}_{\text{ideal}}}(x) \leftarrow \text{TopKIndices}(h(x), \tilde{k}_{\text{ideal}})

  3. 采样子集用于计算： Sco-act(x)∼UniformSample(Sk~ideal(x),ktrain)\mathcal{S}_{\text{co-act}}(x) \sim \text{UniformSample}(\mathcal{S}_{\tilde{k}_{\text{ideal}}}(x), k_{\text{train}})

  4. 计算MoE输出： yco-act(x)=∑i∈Sco-act(x)exp⁡(hi(x))∑j∈Sco-act(x)exp⁡(hj(x))⋅Ei(x;θi)y_{\text{co-act}}(x) = \sum_{i \in \mathcal{S}_{\text{co-act}}(x)} \frac{\exp(h_i(x))}{\sum_{j \in \mathcal{S}_{\text{co-act}}(x)} \exp(h_j(x))} \cdot E_i(x; \theta_i)

关键创新：

  • 训练时只计算 ktraink_{\text{train}} 个专家（保持训练效率）
  • 但从大小为 kidealk_{\text{ideal}} 的候选池中采样（覆盖多样化组合）
  • 多个训练步骤提供随机近似，捕捉多样化专家组合

层次化路由器损失

问题：路由器在低推理预算下无法做出有效选择，因为专家权重接近均匀分布。

解决方案：鼓励路由器分布远离均匀分布。

LHR=−DKL(q(x)∥U)=−∑i=1Nqi(x)log⁡(qi(x)1/N)\mathcal{L}_{\text{HR}} = -D_{KL}(q(x) \| \mathcal{U}) = -\sum_{i=1}^{N} q_i(x) \log\left(\frac{q_i(x)}{1/N}\right)

其中 q(x)=softmax(h(x))q(x) = \text{softmax}(h(x))，U\mathcal{U} 为均匀分布。

为什么用反向KL而非前向KL：

  • 前向KL梯度：−∂∂qiDKL(U∥q(x))=1Nqi(x)-\frac{\partial}{\partial q_i} D_{KL}(\mathcal{U} \| q(x)) = \frac{1}{N q_i(x)}，当 qi(x)q_i(x) 很小时梯度爆炸
  • 反向KL梯度：∂LHR/∂qi=−log⁡(qi(x)N)−1\partial \mathcal{L}_{\text{HR}} / \partial q_i = -\log(q_i(x) N) - 1，更平滑

效果：

  • 产生清晰、层次化的专家排名
  • 少量专家或大量专家激活时都能可靠工作
  • 熵从7.09降至5.78（数学），从8.63降至7.87（常识）

总损失函数

L=Lce+αLb+λLHR\mathcal{L} = \mathcal{L}_{\text{ce}} + \alpha \mathcal{L}_b + \lambda \mathcal{L}_{\text{HR}}

其中：

  • Lce\mathcal{L}_{\text{ce}}：交叉熵损失
  • Lb\mathcal{L}_b：负载均衡损失
  • LHR\mathcal{L}_{\text{HR}}：层次化路由器损失

四、核心创新

创新点说明理论/实验依据
推理时扩展墙首次发现标准MoE无法弹性扩展的现象共现矩阵分析
随机共激活采样解耦训练预算与组合覆盖蒙特卡洛采样思想
层次化路由器损失鼓励果断的专家排名反向KL正则化
零额外训练开销EMoE与标准Top-k训练成本相同仅在非密集计算部分引入
通用性适用于LoRA和FFN两种MoE场景4种架构验证

五、实验结果

实验设置

模型：

  • LoRA-based：LLaMA2-7B + 32 LoRA专家，ktrain=2k_{\text{train}}=2
  • FFN-based：
    • OLMoE-1B-7B-0924：ktrain=4k_{\text{train}}=4
    • DeepSeek-V2-Lite：ktrain=6+2k_{\text{train}}=6+2（含共享专家）
    • ERNIE-4.5-21B-A3B

基准测试（9个）： ARC-c, ARC-e, GSM8K, HellaSwag, HumanEval, NQ, TriviaQA, WinoGrande, MMLU

基线：

  • 标准Top-k MoE
  • Top-p routing
  • AdaMoE

主要结果

Table 1: EMoE vs Top-k 在不同模型上

模型方法k’=1k’=2k’=4k’=6
LoRAMoETop-k45.4847.8848.0547.57 ↓
LoRAMoEEMoE45.6848.2249.0049.50 ↑
OLMoETop-k--35.31-
OLMoEEMoE--37.81-

关键发现：

  • Top-k模型在 k′>kk' > k 时性能下降（推理时扩展墙）
  • EMoE性能随 k′k' 单调递增
  • 在所有模型架构上一致有效

动态路由方法对比

Table 2: 与Top-p和AdaMoE对比

方法k’=1k’=2k’=4k’=6
Top-k45.4847.8848.0547.57
Top-p46.0947.0448.1047.48
AdaMoE41.3447.1848.5448.52
EMoE45.6848.2249.0049.50

关键发现：

  • 动态路由方法（Top-p, AdaMoE）在固定预算内优化，但无法突破扩展墙
  • EMoE专为可变和可扩展计算资源设计

消融实验

Table 3: EMoE组件消融

方法k’=1k’=2k’=4k’=6
Top-k45.4847.8848.0547.57
EMoE45.6848.2249.0049.50
w/o co-act.45.8348.0348.6848.15
w/o LHR\mathcal{L}_{\text{HR}}45.1947.7948.8149.08

关键发现：

  • 去掉co-activation sampling：k′=6k'=6 时性能下降（协作失败）
  • 去掉 LHR\mathcal{L}_{\text{HR}}：k′=1k'=1 时性能下降（排名不稳定）
  • 两者缺一不可

kidealk_{\text{ideal}} 超参数分析

k_ideal分析

关键发现：

  • kideal=2×ktraink_{\text{ideal}} = 2 \times k_{\text{train}} 已能获得显著增益
  • 最佳范围：2−4×ktrain2-4 \times k_{\text{train}}
  • 过大的 kidealk_{\text{ideal}} 导致收益递减

训练效率

Table 9: 训练开销对比

方法ktraink_{\text{train}}训练时间GPU内存
EMoE210.92h43.4GB
Top-k210.92h43.4GB

关键发现：EMoE训练开销与标准Top-k完全相同。

与大 ktraink_{\text{train}} 训练对比

Table 4: EMoE vs 大 ktraink_{\text{train}} 训练

方法k’=1k’=2k’=4k’=6GPU内存
Top-k (ktrain=6k_{\text{train}}=6)42.8746.2648.3049.1476.0GB
EMoE (ktrain=2k_{\text{train}}=2)45.6848.2249.0049.5043.4GB

关键发现：

  • 大 ktraink_{\text{train}} 训练缺乏弹性（k′=2k'=2 时性能下降）
  • 训练成本高3×
  • EMoE在更低成本下实现更好性能

六、相关工作

方法特点EMoE优势
标准Top-k MoE固定激活专家数无弹性
动态路由 (Top-p, AdaMoE)token间重分配计算固定预算内优化
专家剪枝提取子网络精度损失大
推理时深度扩展增加推理链长度不同维度
EMoE弹性专家激活2-3×扩展范围

七、总结

核心贡献

  1. 发现推理时扩展墙：首次揭示标准MoE模型无法弹性扩展的现象，并通过专家共现分析诊断根因
  2. 随机共激活采样：将训练预算与组合覆盖解耦，用标准训练成本覆盖多样化专家组合
  3. 层次化路由器损失：通过反向KL正则化鼓励果断的专家排名，确保不同预算下的可靠性能
  4. 零额外训练开销：EMoE与标准Top-k训练成本完全相同
  5. 广泛验证：在4种MoE架构（7B-21B）和9个基准测试上一致有效

技术影响

  • 为MoE模型在异构硬件和波动工作负载下的部署提供了实用解决方案
  • 单个训练模型可服务多种推理预算，无需维护多个变体
  • 与现有MoE优化技术（量化、蒸馏等）正交，可进一步结合
  • 为MoE模型的推理时scaling研究开辟了新方向

局限性

  • 弹性范围有界：kidealk_{\text{ideal}} 过大时收益递减
  • 在超大规模模型（>100B）上的验证有限
  • 极端情况下（kideal=N,ktrain=2k_{\text{ideal}}=N, k_{\text{train}}=2）退化为随机选择

八、关键图片索引

图片说明文件名
Figure 1专家共现矩阵（k=2 vs k=6）expert-cooccurrence-k2.jpg, expert-cooccurrence-k6.jpg
Figure 2推理时扩展墙现象performance-scaling-wall.jpg
Figure 3共现距离与性能关系cooccurrence-distance.jpg
Figure 4EMoE架构对比architecture-overview.jpg
Figure 5kidealk_{\text{ideal}} 超参数分析effect-of-k-ideal.jpg
Figure 7OLMoE性能与延迟olmoe-performance-latency.jpg
Figure 8不同训练预算training-budgets.jpg
Figure 9DeepSeek-V2的 kidealk_{\text{ideal}} 分析deepseek-k-ideal.jpg
Figure 10LHR\mathcal{L}_{\text{HR}} 对专家选择稳定性影响cooccurrence-math.jpg, cooccurrence-commonsense.jpg
Figure 11训练动态（epoch 1-4）training-dynamics-epoch{1,2,3,4}.jpg

九、参考资源