Alloc-MoE: Budget-Aware Expert Activation Allocation for Efficient Mixture-of-Experts Inference¶

会议: ACL 2026
arXiv: 2604.08133
代码: 无
领域: LLM 效率 / MoE 推理
关键词: Mixture-of-Experts, 专家激活预算, 动态规划, Token 级再分配, 推理加速

一句话总结¶

把 MoE 推理的"激活专家个数"抽象成全局预算 \(B\)，先用动态规划在层间做最优 Top-K 分配（Alloc-L），再用全局 Top-(K·T) 选择在 token 间重分配（Alloc-T），在 DeepSeek-V2-Lite 上把激活预算砍掉一半还能保持精度，prefill 加速 1.15×、decode 加速 1.34×。

研究背景与动机¶

领域现状：稀疏 MoE（DeepSeek-V2/V3、Qwen-MoE、OLMoE 等）把每个 token 路由到 Top-K 个专家。推理延迟近似随激活专家数线性增长，因此降低 K 是天然的加速思路。

现有痛点：两条主流压缩路线都把"减少激活"当目标，但忽视了对模型性能的影响。Token 级方法（XMoE 的 Top-P、Dynamic-MoE、NAEE、AdapMoE）大多依赖训练期校准或固定阈值；层级方法（LExI 等启发式）对不同层一视同仁。论文实测：DeepSeek-V2-Lite 把每 token 激活从 6 砍到 3，性能掉 17%；砍到 2 时直接掉到接近 40%。

核心矛盾：激活配额是稀缺资源，但既不应在所有层均摊，也不应在所有 token 均摊——不同层对稀疏化的敏感度差异显著，不同 token 的路由分布也疏密不一。已有工作要么只动层、要么只动 token，从未把两者作为同一个预算优化问题来解。

本文目标：在固定全局激活预算 \(B\) 下，找到层级分配 \(\mathbf{K}^{\ast}=[K_0,\dots,K_{L-1}]\) 和 token 级激活分布，使整体性能损失最小。

切入角度：作者发现层敏感度可以被一个端到端 PPL 指标准确量化（关键技巧：profile 第 \(i\) 层时把它后面所有层 固定到 Top-1、前面所有层保持 \(K_{\text{orig}}\)，剥离深层补偿效应）；token 级再分配则可以转化成在 \(T\times K_{\text{orig}}\) 的候选集合上做全局 top-selection。

核心 idea：把"激活预算"显式建模成可调度的资源，用 DP 做层间最优分配 + 用全局 top-selection 做 token 间再分配，两者正交、可叠加。

方法详解¶

整体框架¶

输入：预训练 MoE 模型 \(M\)（含 \(L\) 个 MoE 层、原始 Top-K \(K_{\text{orig}}\)）、校准数据 \(D_{\text{calib}}\)、全局激活预算 \(B\)。离线阶段（Alloc-L）：① 对每个层 \(i\) 跑"隔离式敏感度 profiling"，构建敏感度矩阵 \(\mathbf{S}\in\mathbb{R}^{L\times K_{\text{orig}}}\)；② 把"层间预算分配"建成分组背包，用 DP 求出最优 \(\mathbf{K}^{\ast}\)。在线阶段（Alloc-T）：在每一层 \(l\) 内部，按 \(K_l^{\ast}\) 重新分配 token 之间的激活——先给每个 token 保底 \(K_{\text{base}}\) 个专家，剩下的 \((K_l-K_{\text{base}})\cdot T\) 个名额在所有 token 的候选 (token, expert) 对中按路由分数全局 top-K 抢占。

关键设计¶

Alloc-L：隔离式敏感度 profiling + 背包 DP:
- 功能：在不重训的前提下，为每一层算出"该层每种 Top-K 取值下的相对损失"，并求出预算约束下的最优层级分配。
- 核心思路：profiling 时按层从深到浅扫描，第 \(i\) 层取 \(k\in\{K_{\text{orig}},\dots,1\}\) 时，强制所有 \(j>i\) 的深层都用 Top-1（剥离它们的补偿），所有 \(j<i\) 的浅层保持原配置，记 \(\mathbf{S}[i,k]=\mathrm{PPL}\)。求解 \(\arg\min_{\mathbf{K}}\sum_i\mathbf{S}[i,K_i]\ \text{s.t.}\ \sum_i K_i\le B\) 时套分组背包 DP：\(\mathrm{DP}[i,b]=\min_{k\le b}(\mathrm{DP}[i-1,b-k]+\mathbf{S}[i,k])\)，时间复杂度 \(O(L\cdot B\cdot K_{\text{orig}})\)，因 \(B\le L K_{\text{orig}}\) 实际很快。
- 设计动机：直接 grid-search 所有 \(\mathbf{K}\) 组合是 \(K_{\text{orig}}^L\) 量级；不剥离深层补偿则浅层敏感度被掩盖（深层会用更多专家"擦屁股"）。隔离 profiling + DP 同时解决"测得准"与"算得快"两个问题。
Alloc-T：保底 + 全局 top-selection 的 token 级再分配:
- 功能：在不增加任何推理开销的前提下，把固定的层级预算 \(T\cdot K_l\) 个激活机会，按路由分数的疏密重新分配到 token 上——分布越平（高熵）就多给几个专家，分布越尖（高置信）就少给。
- 核心思路：把每层 token-expert 选择写成 0-1 整数规划 \(\max\sum z_{t,e}w_{t,e}\)，约束 \(\sum z_{t,e}\le T\cdot K_l\) 且每个 token 至少分 \(K_{\text{base}}\) 个专家。实操等价于：① 先对每个 token 保住 Top-\(K_{\text{base}}\) 个专家，② 把剩下的 \(K_{\text{orig}}-K_{\text{base}}\) 列分数拍平到一个 \(T\cdot(K_{\text{orig}}-K_{\text{base}})\) 的池里，③ 全局选 top \((K_l-K_{\text{base}})\cdot T\) 个。整套操作就是两次 mask + 一次 top-K，几乎零开销。
- 设计动机：标准 Top-K routing 对每个 token 切一刀，完全忽视了"有些 token 第一个专家分数 0.9、有些 token 前四个分数都 0.25"的客观差异；本设计把激活作为可调度资源，让"路由不自信"的 token 多用配额。当 \(K_{\text{base}}=K_l\) 时退化为标准 Top-K，构成严格的泛化。
Alloc-MoE：层级与 token 级正交叠加:
- 功能：把 Alloc-L 和 Alloc-T 拼成一个统一框架——Alloc-L 决定每层的平均预算 \(K_l^{\ast}\)，Alloc-T 在该 \(K_l^{\ast}\) 下再做 token 级重排。
- 核心思路：两者作用的资源维度互不重叠（层维 vs token 维），可以乘法组合；论文消融显示在 DeepSeek-V2-Lite 全预算区间叠加均最优或近最优。
- 设计动机：单独跑 Alloc-L 在 4 个预算上平均 +0.4%，单独跑 Alloc-T 平均 +0.93%；说明 token 级再分配在激进稀疏化场景下增益更大，但两者并不冲突，组合后取得最高均值 45.19% vs Uniform 的 44.19%。

损失函数 / 训练策略¶

完全无需重训：Alloc-L 一次性离线 profile \(L\cdot K_{\text{orig}}\) 次 PPL 即可拿到 \(\mathbf{S}\)；Alloc-T 是纯推理期 mask + top-K，无任何额外参数。默认 \(K_{\text{base}}=1\)（在所有模型预算上都接近最优）。

实验关键数据¶

主实验¶

在 DeepSeek-V2-Lite（\(L=26\), \(K_{\text{orig}}=6\)）/ Qwen1.5-MoE-A2.7B / OLMoE-1B-7B 三个模型 + 20 个 NLU / Reasoning / Math benchmark 上评估，与 Uniform / LExI / Dynamic-MoE / NAEE 对比。下表是 DeepSeek 半预算（\(B=78\)，即每 token 平均 3 个专家）下的"任务组均值"。

任务组	Uniform	Dynamic-MoE	NAEE	LExI	Alloc-MoE
NLU	~63.0	~62.5	~62.6	~63.4	~63.5
Reasoning	~37.2	~37.9	~38.0	~37.5	~38.6
Math	~31.6	~32.3	~31.7	~31.0	~34.4

整体上 Alloc-MoE 在 12 个 budget×task 配置里赢了 10 个，且任务越难提升越明显：NLU 平均 +0.05%、Reasoning +0.70%、Math +2.15%。推理效率上半预算下 prefill 1.15× / decode 1.34× 加速，与 LExI 同档但精度高一截。

消融实验（DeepSeek-V2-Lite 任务组平均）¶

配置	\(B=52\)	\(B=78\)	\(B=104\)	\(B=130\)	均值
Uniform	41.27	43.90	45.51	46.06	44.19
+Alloc-L	41.53	44.61	45.76	46.47	44.59
+Alloc-T	42.83	45.42	45.93	46.30	45.12
+Alloc-L +Alloc-T	43.09	45.48	46.01	46.17	45.19

关键发现¶

预算越紧（\(B=52\)，即每层均摊 2 个专家）时 Alloc-T 贡献越大（+1.56），说明在激进稀疏化下 token 级再分配比层级再分配更值钱；预算宽松时两者贡献趋平。
\(K_{\text{base}}\) 消融：从 0 到 5 扫，\(K_{\text{base}}=1\) 在所有预算上几乎都是最优，\(K_{\text{base}}\ge 2\) 在紧预算下显著掉点——保得越多，可重分配的灵活性就越小。
校准数据集（WikiText2 / C4 / Pile）对 Alloc-L 几乎无影响，4 个预算的均值在 44.59 ~ 44.68 之间，说明敏感度 profiling 鲁棒。
Alloc-MoE 的层级分配可视化显示：浅层留多、深层砍狠，与"浅层提取通用特征更敏感"的直觉一致；token 级则在分配数和路由熵之间呈现 \(>0.7\) 的强正相关。

亮点与洞察¶

"激活预算"这一抽象是最大亮点：把原本散落在层级 / token 级的稀疏化决策统一成可调度资源，让 DP 与全局 top-K 这两件常规武器立刻可用。这种"先把问题换个名字"的思路可以迁移到 KV 缓存预算、注意力头预算、并行专家切分等场景。
隔离式敏感度 profiling 的小技巧很关键：直接独立量化每层敏感度时，深层会"补偿"浅层的稀疏化，让浅层看起来"无所谓"；强行把后面所有层钉死到 Top-1 反而能测到真实的敏感度。这是个可复用的 profiling 范式。
Alloc-T 的零开销实现优雅：把整数规划等价成"两次 mask + 一次全局 top-K"，没有引入任何额外内核，可直接嵌入现有 MoE 路由器。
expert load 分析做得稳：分配后专家负载分布与原模型 Spearman 相关 0.93~0.99，JS 散度 <0.014，说明并不破坏专家专精——对分布式 MoE 部署友好。

局限与展望¶

作者承认未考虑硬件感知（expert 放置、通信开销），分布式 MoE 部署收益可能进一步放大。
Alloc-L 需要离线跑 \(L\cdot K_{\text{orig}}\) 次 PPL，对 100B+ 模型的 profiling 成本不可忽视；可以考虑用代理小模型估计敏感度。
框架完全 post-hoc，未与训练阶段联动。如果在训练时就引入"激活预算感知"的辅助 loss，可能进一步提升鲁棒性。
仅评了三类任务和三种模型，未覆盖 long-context、code 这种激活模式可能很不同的场景。

评分¶

新颖性: ⭐⭐⭐⭐ "激活预算"的统一抽象 + 隔离敏感度 profiling 是清新组合，但 DP / 全局 top-K 这两件武器都不新。
实验充分度: ⭐⭐⭐⭐ 三模型 × 四预算 × 20 数据集，消融、\(K_{\text{base}}\) 扫、校准集鲁棒性、负载分析齐全；缺超大模型实测。
写作质量: ⭐⭐⭐⭐ 公式与算法清楚，图表（Figure 3-8）信息量大；术语 Alloc-L / Alloc-T 名字直观。
价值: ⭐⭐⭐⭐ 纯后处理、零训练、能直接落地到现有 MoE 框架（vLLM 等），半预算下 34% decode 加速对实际部署吸引力强。