MoE-Spec: Expert Budgeting for Efficient Speculative Decoding

作者/机构: Bradley McDanel 1 2, Steven Li 2, Sruthikesh Surineni 2, Harshit Khaitan 2

A1 主要贡献

本文旨在解决在混合专家(Mixture-of-Experts, MoE)模型上应用推测解码(Speculative Decoding)时遇到的一个严重瓶颈:在并行验证多个草稿(draft)词元时,大型的草稿树会激活大量不同的专家(expert),这极大地增加了内存压力,从而削弱了推测解码相对于自回归解码(autoregressive decoding)的速度优势。

核心问题与观察
对于MoE模型,验证成本不再是恒定的,而是与草案树中专家激活的多样性成正比。每个词元独立地路由到其自己的一组专家,目标模型必须加载所有被选中专家的并集来验证整个树(如图2a所示)。随着草稿规模的增长,唯一专家的数量也随之增加:例如,在OLMoE-1B-7B模型上,一个包含127个词元的树会激活每层64个专家中的54个,几乎接近整个模型(如图2b所示)。这导致验证过程可能比自回归解码更昂贵,从而抵消了稀疏激活和推测解码带来的好处。然而,研究发现并非所有专家对验证的贡献都相同。草稿树中的专家激活遵循一个重尾分布(heavy-tailed distribution),即少数专家占据了绝大部分的路由权重(如图2c所示)。

研究目标与创新点
本文提出了一种名为 MoE-Spec 的方法,它在验证期间强制执行一个严格的专家预算(expert budget),从而将验证成本与草稿树的复杂性解耦。具体方法是:对于一个给定的草稿树,MoE-Spec会聚合所有词元的路由概率,并选择得分最高的B个专家形成一个候选列表(shortlist)。然后,只加载这B个专家进行验证,而丢弃那些使用频率较低、但会带来带宽开销的长尾专家。对于其自然路由超出候选列表的词元,通过截断(omitting missing experts)或替换(replacing them with available alternatives)策略进行处理。该方法无需训练,可直接集成到现有的推测解码流程中。

图1. MoE-Spec 概述。标准的 MoE 推测解码会加载草稿树中激活的所有 N 个专家(N=8)(上图路径)。MoE-Spec 按聚合路由概率对专家进行排序,并强制执行一个预算 B(此处 B=4),只加载得分最高的专家(下图路径)。词元根据其自然路由中缺失专家的数量进行颜色编码。在此示例中,尽管 MoE-Spec 加载的专家数量减半,但两条路径接受了相同的三个词元。
图1. MoE-Spec 概述。标准的 MoE 推测解码会加载草稿树中激活的所有 N 个专家(N=8)(上图路径)。MoE-Spec 按聚合路由概率对专家进行排序,并强制执行一个预算 B(此处 B=4),只加载得分最高的专家(下图路径)。词元根据其自然路由中缺失专家的数量进行颜色编码。在此示例中,尽管 MoE-Spec 加载的专家数量减半,但两条路径接受了相同的三个词元。
图2. 专家预算的动机。(a) 在验证期间,每个词元路由到 k 个专家;目标模型必须加载草稿树中所有唯一的专家 E。(b) 随着树大小 M 的增长,每层的唯一专家数接近 N,从而抵消了稀疏激活的好处。(c) 路由概率呈重尾分布:对于大小为 63 的树,排名前 32 的专家(共 64 个)占据了 93% 的路由权重。
图2. 专家预算的动机。(a) 在验证期间,每个词元路由到 k 个专家;目标模型必须加载草稿树中所有唯一的专家 E。(b) 随着树大小 M 的增长,每层的唯一专家数接近 N,从而抵消了稀疏激活的好处。(c) 路由概率呈重尾分布:对于大小为 63 的树,排名前 32 的专家(共 64 个)占据了 93% 的路由权重。

主要贡献如下:
* 揭示了重尾分布:证明了在MoE模型验证过程中,专家激活是重尾的,一小部分专家占据了草稿树中的大部分路由权重。
* 提出MoE-Spec方法:提出了一种简单而有效的MoE-Spec方法,通过在验证期间选择一个有预算的顶级专家子集来强制执行专家预算。该方法无需训练,可直接集成到现有的推测解码流程中。
* 验证性能提升:在三种不同规模和专家配置的MoE架构上,实验证明MoE-Spec相比EAGLE-3在吞吐量上有显著提升。
* 揭示优化潜力:通过Oracle分析表明,改进的排序方法有可能比当前基于路由器的启发式方法减少25%的专家需求,为未来的研究指明了方向。

A3 背景知识

2.1. 混合专家模型(Mixture-of-Experts)

MoE模型通过稀疏激活实现参数效率。混合专家(MoE)模型通过仅为每个词元激活模型权重的一个子集来达到参数效率。在一个MoE层中,标准的前馈网络被N个并行的专家网络{E1, ..., EN}和一个学习到的路由器(router)所取代,该路由器为每个词元选择k个专家进行处理(其中k ≪ N)。

路由器计算专家概率分布。对于一个给定的词元表示 $h \in \mathbb{R}^d$,路由器通过一个线性投影和softmax函数来计算专家的概率分布:

$$g_{i}(\mathbf{h})=\frac{\exp(\mathbf{w}_{i}^{\top}\mathbf{h})}{\sum_{j=1}^{N}\exp(\mathbf{w}_{j}^{\top}\mathbf{h})},$$

其中 $w_i \in \mathbb{R}^d$ 是专家i的路由权重向量。令 $T_k(h) \subset \{1, ..., N\}$ 表示按路由概率排名的前k个专家的索引。

MoE层的输出是所选专家的加权组合。MoE层的输出是所选专家的加权组合:

$$\mathrm{MoE}(\mathbf{h}) = \sum_{i \in \mathcal{T}_k(\mathbf{h})} \bar{g}_i(\mathbf{h}) \cdot E_i(\mathbf{h}),$$

其中 $\bar{g}_i(h)$ 表示混合权重。这些权重可以是原始的路由概率 $g_i(h)$,也可以是重新归一化的值 $g_i(h) / \sum_{j \in T_k(h)} g_j(h)$,具体取决于模型架构(例如,OLMoE进行重新归一化;Qwen3不进行)。

稀疏激活模式的代价。这种稀疏激活模式意味着推理成本与活动参数的数量成正比,而不是总参数。对于单个词元,只需要从内存中加载k个专家(通常k∈{2, 8})。然而,当联合处理多个词元时,模型必须加载所有词元选择的专家的并集,随着词元数量的增加,这个并集的大小可能接近N个专家的全集。

2.2. 使用草稿树的推测解码

推测解码通过并行验证加速自回归生成。推测解码通过起草多个候选词元并并行验证它们来加速自回归生成【11, Y. Leviathan et al., Fast inference from transformers via speculative decoding, 2023, ICML】【3, C. Chen et al., Accelerating large language model decoding with speculative sampling, 2023, arXiv】。一个轻量级的草稿模型提出候选的续写,然后完整的目标模型在单次前向传播中验证这些候选。被接受的词元会立即输出;被拒绝的词元会从修正后的分布中重新采样。对于贪婪解码,这个过程是可证明无损的,因为其输出分布与标准自回归解码相匹配。

现代方法采用树状草稿。现代方法如EAGLE【13, Y. Li et al., Eagle: Speculative sampling requires rethinking feature uncertainty, 2024, arXiv】【14, Y. Li et al., Eagle-2: Faster inference of language models with dynamic draft trees, 2024, EMNLP】将此扩展到树状结构的草稿,其中在每个位置提出多个候选词元。目标模型使用树状注意力机制同时验证所有候选,并接受一个与目标分布一致的词元路径。

稠密模型与MoE模型的验证成本差异。对于稠密模型,无论草稿树的大小如何,验证成本实际上是恒定的:所有模型参数在每次前向传播中加载一次,在内存受限的情况下,额外词元的计算量可以忽略不计。这使得激进的推测变得有吸引力——更大的树增加了每次迭代中预期接受的词元数量。然而,对于MoE模型,这个假设不成立;我们将在3.1节中分析这个问题。

A2 方法细节

3.1. 树状解码中的验证开销

形式化定义MoE验证成本。如第1节所述,MoE的验证成本随着草稿树中专家多样性的增加而增加。我们将其形式化如下。令 $D = \{h_1, ..., h_M\}$ 表示一个包含M个词元表示的草稿树。为了验证该树,目标模型必须加载所有被选中专家的并集:

$$\mathcal{E}(\mathcal{D})=\bigcup_{t=1}^{M} \mathcal{T}_{k}\left(\mathbf{h}_{t}\right).$$

专家路由的重尾分布是降低开销的关键。如图2b所示,对于大型树,$|\mathcal{E}|$ 的大小接近N,这抵消了稀疏激活的好处。然而,专家路由是重尾的(图2c):通过强制设定一个预算 $B < |\mathcal{E}|$ 并只加载得分最高的专家,我们可以在保持验证质量的同时减少带宽消耗。

3.2. 专家重要性排序

专家排序策略概述。给定一个专家预算 $B < N$,我们在验证期间选择一个包含B个专家的候选列表S进行加载。我们考虑三种排序策略。

静态排序 (Static Ranking)。最简单的方法是通过计算校准数据中的选择频率来预先计算一个固定的排序:

 c_i = \sum_{\mathbf{h} \in \text{calib}} \mathbf{1}[i \in \mathcal{T}_k(\mathbf{h})].
c_i = \sum_{\mathbf{h} \in \text{calib}} \mathbf{1}[i \in \mathcal{T}_k(\mathbf{h})].

候选列表由具有最高 $c_i$ 的B个专家组成。这种方法运行时开销为零,并且可以实现永久性的模型剪枝——例如,部署64个专家中的32个可将内存需求减半。在适度的预算下($B \ge N/2$),静态排序在大多数任务上的表现与自适应方法相当(见4.6节)。但在激进的预算下,差距会拉大,此时依赖于输入的路由变得更有价值。

基于路由器的排序 (Router-Based Ranking)。我们的主要方法是聚合草稿树中所有词元的路由概率:

$$s_i = \sum_{t=1}^M g_i(\mathbf{h}_t).$$

候选列表S由具有最高 $s_i$ 的B个专家组成。这不需要超出路由器已产生计算的额外计算,并且能适应每棵树的路由模式。

Oracle排序 (Oracle Ranking)。作为性能上限,我们考虑一个可以访问目标模型路由决策和专家输出的“先知”(Oracle)。这个Oracle会贪婪地选择专家,以最小化与未剪枝输出之间的重构误差。设 $O_t^* = \sum_{i \in \mathcal{T}_k(h_t)} g_i(h_t) \cdot E_i(h_t)$ 表示黄金标准输出。从 $S = \emptyset$ 开始,Oracle迭代地添加能最小化以下误差的专家 $i^*$:

$$i^* = \arg \min_{i \notin \mathcal{S}} \sum_{t=1}^{M} \left\| O_t^* - \sum_{j \in \mathcal{S} \cup \{i\}} g_j(\mathbf{h}_t) \cdot E_j(\mathbf{h}_t) \right\|_2^2$$

这种方法需要运行所有专家,因此对于加速不切实际,但它为可实现的质量提供了一个上限。

3.3. 专家覆盖

处理自然路由超出候选列表的词元。给定候选列表S,我们必须处理那些自然路由中包含S之外专家的词元。我们考虑两种策略。

截断 (Truncation)。每个词元使用其自然的top-k路由,但S之外的专家贡献为零:

$$\mathrm{MoE}_{\mathcal{S}}^{\text {trunc }}\left(\mathbf{h}_{t}\right)=\sum_{i \in \mathcal{T}_{k}\left(\mathbf{h}_{t}\right) \cap \mathcal{S}} \bar{g}_{i}\left(\mathbf{h}_{t}\right) \cdot E_{i}\left(\mathbf{h}_{t}\right).$$

概率 $\bar{g}_i$ 是在原始的top-k集合上计算的(而不是为缩减后的集合重新计算)。那些自然路由与S部分重叠的词元会接收到少于k个的专家。在极端情况下,如果 $T_k(h_t) \cap S = \emptyset$,MoE的输出为零,只有残差连接会通过。这实际上是在该层为该词元跳过了MoE计算。

替换 (Substitution)。Top-k选择被限制在S内。令 $T_k^S(h_t)$ 表示在S内按 $g_i(h_t)$ 排序的前k个专家:

$$\mathrm{MoE}_{\mathcal{S}}^{\text{sub}}(\mathbf{h}_t) = \sum_{i \in \mathcal{T}_k^{\mathcal{S}}(\mathbf{h}_t)} \bar{g}_i^{\mathcal{S}}(\mathbf{h}_t) \cdot E_i(\mathbf{h}_t),$$

其中 $\bar{g}^S$ 是在受限选择上的路由概率。如果一个词元的自然top-k包括了S之外的专家,那些专家将被候选列表内得分最高的可用备选专家所取代。每个词元都能接收到恰好k个专家,尽管它们可能与自然路由不同。我们在4.4节中对这两种策略进行了实证比较。

A4 实验环境

  • 模型架构:

    • OLMoE-1B-7B: 64个专家,k=8,总参数7B,活动参数1B。使用 EAGLE-3 草稿模型。
    • Qwen3-30B-A3B: 128个专家,k=8,总参数30B,活动参数3B。使用 EAGLE-3 草稿模型。
    • Mixtral-8x7B: 8个专家,k=2,总参数47B,活动参数13B。使用 EAGLE (Li et al., 2024b) 草稿模型。
  • 数据集与任务:

    • 数学推理: GSM8K (多步算术), MATH500 (竞赛级数学),评估指标为准确率 (Accuracy)。
    • 代码生成: HumanEval, MBPP (函数合成),评估指标为 pass@1。
    • 摘要: CNN/DailyMail (抽取式摘要),评估指标为 ROUGE-L F1。
  • 硬件配置:

    • GPU: NVIDIA A100 (80GB显存)。OLMoE和Qwen3使用1个GPU,Mixtral使用2个GPU。
  • 软件配置:

    • 代码实现: 基于开源的 EAGLE-3 代码库进行扩展。
    • 依赖库: PyTorch 2.1, CUDA 12.1。
    • 精度: FP16。
  • 实验设置:

    • 基线: 自回归解码 (AR) 和 EAGLE (当前SOTA的树状推测解码方法)。
    • MoE-Spec配置: 主要结果使用基于路由器的专家选择和替换策略。
    • 草稿树大小: 63个词元。
    • 解码温度: 0 (贪心解码) 和 1 (采样)。

A4 实验结果

主要结果

MoE-Spec在保持质量的同时显著提升了吞吐量。如表1所示,在30个模型-基准测试配置中,MoE-Spec有27个配置的加速比优于EAGLE,同时保持了相当的质量。例如,在Mixtral模型的MBPP任务上(T=0),MoE-Spec实现了2.3倍的加速,而EAGLE为1.7倍,在任务准确率相同的情况下,相对提升了35%。

更大模型从专家预算中获益更多。Mixtral-8x7B(14B活动参数)的平均相对加速增益为27%;Qwen3-30B-A3B(3B活动参数)为16%;而OLMoE-1B-7B(1B活动参数)为6%。MoE-Spec的平均接受长度与EAGLE相似(平均仅降低1.4%),因此加速完全来自于验证成本的降低。

专家预算提供了质量与延迟的权衡。图3展示了MoE-Spec提供了一个连续的帕累托曲线,允许用户根据质量要求选择不同的操作点。Qwen3模型展现了最佳的权衡,即使在2.0倍加速下也能保持约95%的质量。Mixtral的质量下降适中,而OLMoE在激进的加速目标下质量下降最快。

表格1. 相对于自回归解码的加速比、平均接受长度(τ)以及在T=0和T=1下的任务指标。MoE-Spec使用替换策略,并选择了在保持与EAGLE-3质量相当(在一个标准差内,根据5次T=1运行测量)的情况下最大化吞吐量的专家预算。T=1的结果报告了5个随机种子的平均值。
表格1. 相对于自回归解码的加速比、平均接受长度(τ)以及在T=0和T=1下的任务指标。MoE-Spec使用替换策略,并选择了在保持与EAGLE-3质量相当(在一个标准差内,根据5次T=1运行测量)的情况下最大化吞吐量的专家预算。T=1的结果报告了5个随机种子的平均值。

图3. T=1时质量与加速的权衡,在五个基准测试上取平均。对每个基准,质量和加速都相对于AR(100%)进行归一化,然后取平均。误差棒结合了跨基准的方差和5次运行的种子间方差。每个基准的单独曲线见附录E。
图3. T=1时质量与加速的权衡,在五个基准测试上取平均。对每个基准,质量和加速都相对于AR(100%)进行归一化,然后取平均。误差棒结合了跨基准的方差和5次运行的种子间方差。每个基准的单独曲线见附录E。

加速分析

MoE-Spec通过限制专家数量来降低验证成本。图4展示了在OLMoE-1B-7B上,MoE-Spec(B=32)和EAGLE-3在不同草稿树大小下的表现。

  • 专家激活数量(图4a):随着树大小增加,EAGLE-3激活的专家数量逼近全部64个,而MoE-Spec则饱和在预算值32。
  • 加速比与树大小(图4b):EAGLE-3的加速比在树大小为31时达到峰值后下降,因为加载“掉队”(straggler)专家的成本超过了收益。而MoE-Spec的加速比持续增长。
  • 加速比与激活专家数(图4c):该图显示存在一个约32个专家的饱和点。EAGLE-3在大的树尺寸下会加载超过饱和点的专家(如54个),导致效率降低,而MoE-Spec通过预算限制避免了这个问题。

方法鲁棒性来源:在B=32时,被排除的专家总共只接收了草稿树中7%的路由权重,因此对结果影响较小。选择开销(计算路由概率和排序)仅占总生成时间的2-3%。

图4. OLMoE-1B-7B上的专家激活与加速比。(a) 随着树大小增加,验证过程中激活的唯一专家数量。EAGLE-3在大型树上加载超过50个(共64个)专家;带有B=32的MoE-Spec在预算处饱和。(b) 相对于AR的加速比。EAGLE-3在树大小为31时达到峰值,然后随着专家加载成为主导而下降;MoE-Spec在更大的树上持续改进。(c) 加速比与激活专家数。在树大小为255时,EAGLE-3加载54个专家获得1.85倍加速;MoE-Spec加载32个专家获得2.1倍加速。
图4. OLMoE-1B-7B上的专家激活与加速比。(a) 随着树大小增加,验证过程中激活的唯一专家数量。EAGLE-3在大型树上加载超过50个(共64个)专家;带有B=32的MoE-Spec在预算处饱和。(b) 相对于AR的加速比。EAGLE-3在树大小为31时达到峰值,然后随着专家加载成为主导而下降;MoE-Spec在更大的树上持续改进。(c) 加速比与激活专家数。在树大小为255时,EAGLE-3加载54个专家获得1.85倍加速;MoE-Spec加载32个专家获得2.1倍加速。

消融研究

专家选择策略对比(图5a):
* 静态(Static)排序:在低预算(B≤16)下性能很差,因为它无法适应输入的变化。
* 基于路由器(Router-based)排序:在整个预算范围内表现稳健,紧密跟随Oracle的性能。
* Oracle排序:作为性能上限,即使在B=24(37.5%的专家)时也接近AR的质量,表明选择策略仍有改进空间。

专家覆盖策略对比(图5b):
* 截断(Truncation)替换(Substitution) 两种策略在所有基准和预算下表现相当。本文选择替换策略作为默认设置。

图5. OLMoE-1B-7B在T=1时的设计消融研究,相对于自回归(AR)基线进行归一化。(a) 选择方法:静态排序在低预算下失败,而基于路由器的排序则紧跟Oracle上限。(b) 覆盖策略:截断和替换在各种预算下表现相当。
图5. OLMoE-1B-7B在T=1时的设计消融研究,相对于自回归(AR)基线进行归一化。(a) 选择方法:静态排序在低预算下失败,而基于路由器的排序则紧跟Oracle上限。(b) 覆盖策略:截断和替换在各种预算下表现相当。

重构误差分析

Oracle与Router-based方法的差距来源。图6显示,Oracle方法可以用比Router-based方法少25%的专家达到相同的重构误差。这源于两种方法的预算分配方式不同:Router-based方法独立地为每个专家评分,忽略了专家间的交互作用。如果两个专家经常共同激活,它们对输出的贡献是相关的,同时选择两者会产生边际效益递减。而Oracle通过贪婪地最小化重构误差,避免了这种冗余。分析表明,某些专家对的共激活频率比随机选择高10-30倍,这解释了Oracle的优势。

图6. OLMoE-1B-7B上不同选择方法的重构误差,在所有MoE层和5个基准上取平均。Oracle使用比Router少25%的专家即可达到相同的误差。
图6. OLMoE-1B-7B上不同选择方法的重构误差,在所有MoE层和5个基准上取平均。Oracle使用比Router少25%的专家即可达到相同的误差。

跨数据集的专家覆盖率

不同任务的路由集中度不同。图7显示,代码生成任务(MBPP, HumanEval)的路由非常集中,前8个专家就捕获了超过90%的路由概率。而摘要任务(CNN/DM)的路由则更分散,需要24个专家才能达到90%的覆盖率。这解释了为什么静态排序对某些任务效果好,而对另一些任务效果差。基于路由器的选择通过对每个草稿树进行自适应聚合,从而回避了这个问题。

图7. OLMoE-1B-7B上,树大小为63时,不同专家预算在各个数据集上的路由概率覆盖率。数据集按覆盖率从好到坏排序:MBPP用8个专家达到95%的覆盖率,而CNN/DM需要24个专家才能达到90%。
图7. OLMoE-1B-7B上,树大小为63时,不同专家预算在各个数据集上的路由概率覆盖率。数据集按覆盖率从好到坏排序:MBPP用8个专家达到95%的覆盖率,而CNN/DM需要24个专家才能达到90%。

A7 补充细节

5. 相关工作

推测解码。现有的推测解码方法,如EAGLE系列【13, Y. Li et al., Eagle: Speculative sampling requires rethinking feature uncertainty, 2024, arXiv】【14, Y. Li et al., Eagle-2: Faster inference of language models with dynamic draft trees, 2024, EMNLP】【15, Y. Li et al., Eagle-3: Scaling up inference acceleration of large language models via training-time test, 2025, NeurIPS】和Medusa【2, T. Cai et al., Medusa: Simple llm inference acceleration framework with multiple decoding heads, 2024, arXiv】,都假设验证成本是恒定的,这一假设在MoE模型上不成立。

混合专家模型。MoE架构通过稀疏路由降低计算量【20, N. Shazeer et al., Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017, arXiv】,Switch Transformers【5, W. Fedus et al., Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022, JMLR】确立了现代的top-k路由范式。系统方面的工作如MegaBlocks【7, T. Gale et al., Megablocks: Efficient sparse training with mixture-ofexperts, 2023, MLSys】专注于优化MoE的GPU计算核。本文工作针对的是不同的瓶颈:推测验证过程中的内存带宽成本。

动态稀疏性。Deja Vu【16, Z. Liu et al., Deja vu: Contextual sparsity for efficient llms at inference time, 2023, ICML】和PowerInfer【21, Y. Song et al., Powerinfer: Fast large language model serving with a consumer-grade gpu, 2024, SOSP】等工作利用了稠密LLM中的激活重尾分布来跳过部分计算。它们针对的是稠密模型中的单token推理。MoE-Spec将同样原理应用于稀疏模型中的批处理验证,通过聚合整个草稿树的路由概率来选择专家,并且无需额外的学习预测器。

MoE压缩。MoE-Spec在推理时进行专家预算,与永久性减小模型尺寸的方法(如MoE-Pruner【23, Y. Xie et al., Moe-pruner: Pruning mixture-of-experts large language model using the hints from its router, 2024, arXiv】、专家合并【12, P. Li et al., Merge, then compress: Demystify efficient SMoe with hints from its routing policy, 2024, ICLR】、量化【10, Y. J. Kim et al., Mixture of quantized experts (moqe): Complementary effect of low-bit quantization and robustness, 2023, arXiv】【6, E. Frantar & D. Alistarh, Qmoe: Practical sub-1-bit compression of trillion-parameter models, 2023, arXiv】)是正交的。

针对MoE的推测解码。先前的工作,如Saxena等人【19, A. Saxena et al., Utility-driven speculative decoding for mixture-ofexperts, 2025, arXiv】通过动态调整推测长度来应对验证开销,但这主要针对短链式草稿(K≤3)。对于EAGLE-3等使用大型树状草稿(60+ token)的方法,这种方法效果有限,因为专家激活很快饱和。MoE-Spec不减少推测,而是直接降低验证成本。其他工作如SP-MoE【4, L. Chen et al., Sp-moe: Speculative decoding and prefetching for accelerating moe-based model inference, 2025, arXiv】和SpecMoEOff【22, Z. Wang et al., Accelerating mixture-of-experts inference by hiding offloading latency with speculative decoding, 2025, arXiv】关注于专家卸载(offloading)场景,解决PCIe瓶颈,而MoE-Spec解决的是专家在GPU显存内的HBM带宽瓶颈。

6. 局限性

引入超参数。MoE-Spec引入了一个专家预算超参数,需要根据具体部署进行选择。激进的预算可能导致质量下降。

少量开销。基于路由器的排序方法在验证前需要计算所有专家的路由概率,这会增加2-3%的开销,对于非常小的草稿树可能不划算。

推理场景限制。本文评估的是单请求推理。在批处理推理场景下,多个请求可以共享专家,加载额外专家的边际成本降低,因此专家预算带来的增益可能会减小。

路由机制假设。本文实验基于top-k softmax路由,对于其他路由方案(如sigmoid)可能需要调整聚合策略。

未来优化方向。本文没有对每层进行预算自适应调整,并且基于路由器的选择方法未考虑专家间的共激活模式,这些是未来可以改进的方向。

A5 结论

本文的研究表明,MoE模型在验证阶段的专家路由呈现重尾分布:大部分带宽被用于加载对输出质量贡献甚微的专家。MoE-Spec利用这一特性,通过强制执行专家预算,只加载得分最高的专家,从而将验证成本与草稿树的复杂性解耦。实验结果显示,在三种不同的MoE架构上,该方法比EAGLE-3的吞吐量高出10-30%。

该方法无需额外的学习预测器或辅助模型,因为路由器本身已经提供了重要性分数。通过聚合这些分数,即可识别出需要加载的关键专家。本文的Oracle分析表明,通过对专家共激活模式进行建模,而不是独立评分,有可能进一步降低预算需求。这为未来的工作指明了一个自然延伸方向:开发利用专家路由中相关性结构的选择方法。

A6 附录

A. 超参数

通用实验设置和模型特定配置。表2总结了通用的实验设置。表3提供了模型特定的配置。注意,在T=1的实验中使用了较保守的预算范围,而在T=0的实验中探索了表中所示的更完整范围。

草稿树大小。遵循EAGLE-3的配置,默认草稿树大小为63个词元。图4中的树大小扫描(3-511)表明,MoE-Spec的加速效果随树大小的增加而增长,因为标准验证下更大的树会激活更多的唯一专家。选择63是在基线的接受率和验证开销之间做的权衡。

专家预算选择。预算范围根据模型的专家配置而定。对于OLMoE(64专家,k=8),我们扫描了 B ∈ {8, 16, 24, 32, 40, 48, 56};最小值为8确保至少有一组完整的活动专家。对于Qwen3(128专家,k=8),我们使用 B ∈ {16, 24, 32, 48, 64}。Mixtral的专家池较小(8专家,k=2),范围限制在 B ∈ {2, 3, 4, 5, 6, 7}。主要结果(表1)报告的是在保持质量与EAGLE相当的情况下最大化加速的预算。

可复现性。对于T=1的实验,我们报告了5个随机种子的平均值和标准差。所有模型均使用FP16精度。我们使用了PyTorch 2.1和CUDA 12.1。模型权重从HuggingFace Hub加载;草稿模型适配器来自官方EAGLE仓库。

表格2. 通用实验设置。
表格2. 通用实验设置。

表格3. 模型特定配置。
表格3. 模型特定配置。

B. MoE-Spec 算法

算法伪代码。算法1提供了验证时专家预算的伪代码。该算法分为三个阶段:(1)计算草稿树中所有词元的路由概率,(2)通过对路由概率求和来聚合每个专家的重要性分数,(3)仅使用得分最高的B个专家来应用词元策略(截断或替换)。聚合步骤除了路由器已有的计算外,不需要额外的模型执行。

算法1 MoE-Spec专家预算
算法1 MoE-Spec专家预算

C. 批处理MoE计算

实现细节。我们使用开源的EAGLE-3实现,未修改其树构建或验证逻辑。唯一的改动是MoE计算:参考的EAGLE实现是按顺序处理专家的,这种方法简单但效率低下,因为它会发出许多小的矩阵乘法。我们实现了批处理专家计算,使用批处理矩阵乘法并行处理所有活动专家。这减少了内核启动开销,并更好地利用了GPU的并行性。

性能对比。表4比较了两种执行模式在OLMoE-1B-7B上的表现。批处理实现使吞吐量提高了1.4-1.5倍,而任务质量差异可忽略不计。为保证公平比较,我们将批处理执行统一应用于所有方法,包括AR和EAGLE-3基线。

表格4. 在OLMoE-1B-7B上(EAGLE-3,无专家预算)顺序与批处理MoE计算的比较。质量差异在噪声范围内;吞吐量提高了1.4-1.5倍。
表格4. 在OLMoE-1B-7B上(EAGLE-3,无专家预算)顺序与批处理MoE计算的比较。质量差异在噪声范围内;吞吐量提高了1.4-1.5倍。

D. 专家共激活分析

分析动机。为了理解为什么Oracle选择优于基于路由器的选择(4.5节),我们分析了推理过程中的专家共激活模式。共激活衡量专家对为同一词元被同时选择的频率。

方法论。我们在每个基准测试的50个样本上运行OLMoE-1B-7B推理,并记录每个词元在验证过程中的top-k选择。对于每个MoE层,我们构建一个64×64的共激活矩阵。

集中度度量。我们将路由集中度定义为最频繁共激活对的计数与均匀随机选择下期望计数之间的比率。

$$\text{Concentration} = \frac{\max_{i,j} \text{coact}(i, j)}{\text{expected per pair}}$$

结果。图8显示了所有16个MoE层的共激活热力图。在各层中,某些专家组的共激活频率远高于随机水平:集中度比率从10倍到33倍不等(平均23倍)。这种集中性解释了Oracle与Router之间的差距。基于路由器的选择会倾向于从同一个高频集群中选择多个专家,而Oracle选择通过贪婪地最小化重构误差来避免这种冗余。

图8. OLMoE-1B-7B中所有16个MoE层的专家共激活热力图。每个单元格显示归一化的成对共激活频率。专家通过层次聚类重新排序。明亮的对角块表示经常一起激活的专家组,其集中度比率比均匀随机期望高10-33倍。
图8. OLMoE-1B-7B中所有16个MoE层的专家共激活热力图。每个单元格显示归一化的成对共激活频率。专家通过层次聚类重新排序。明亮的对角块表示经常一起激活的专家组,其集中度比率比均匀随机期望高10-33倍。

E. 各基准的质量-加速权衡

详细权衡曲线。图9和图10展示了在温度T=0和T=1下,所有模型-基准组合的完整质量-加速权衡曲线。菱形标记代表EAGLE,连接的点代表不同专家预算下的MoE-Spec。

观察结论。在两种设置下,MoE-Spec都将帕累托前沿扩展到了更高的加速比,只有在激进的预算下才会出现质量下降。总的来说,我们发现代码生成任务(HumanEval, MBPP)比推理任务(GSM8K, MATH500)更能容忍更紧的预算,这与4.6节的覆盖率分析一致。摘要任务(CNN/DM)介于两者之间。

图9. T=0时,三种模型(行)和五个基准(列)的质量与加速比关系图。菱形标记代表EAGLE;连接的点表示不同专家预算下的MoE-Spec。
图9. T=0时,三种模型(行)和五个基准(列)的质量与加速比关系图。菱形标记代表EAGLE;连接的点表示不同专家预算下的MoE-Spec。
图10. T=1时,三种模型(行)和五个基准(列)的质量与加速比关系图。误差棒显示了5个种子的标准差。
图10. T=1时,三种模型(行)和五个基准(列)的质量与加速比关系图。误差棒显示了5个种子的标准差。