BEAM：用于MoE动态路由的二值专家激活掩码

# BEAM：面向MoE动态路由的二元专家激活掩码
来源：https://arxiv.org/html/2605.14438
Juntong Wu¹,²,*, Jialiang Cheng¹,*, 🖂, Qishen Yin², Yue Dai¹, & Yuliang Yan¹, Fuyu Lv¹, Ou Dan¹, Li Yuan², 🖂
¹阿里巴巴淘天集团
²北京大学深圳研究生院
通信邮箱：jichen\.cjl@alibaba\-inc\.com (https://arxiv.org/html/2605.14438v1/mailto:[email protected]), yuanli\-ece@pku\.edu\.cn (https://arxiv.org/html/2605.14438v1/mailto:[email protected])

###### 摘要
混合专家（MoE）架构通过仅对每个令牌激活一部分专家，提升了大型语言模型的效率。然而，标准 MoE 采用固定的 Top-K 路由策略，导致冗余计算和次优的推理延迟。现有的加速方法要么需要伴随架构改变的昂贵重训练，要么因训练与推理不匹配而在高稀疏度下性能严重下降。为解决这些限制，我们提出 BEAM（二元专家激活掩码），一种通过可训练二元掩码学习令牌自适应专家选择的新方法。利用直通估计器和辅助正则化损失，BEAM 通过端到端训练诱导动态专家稀疏性，同时保持模型能力。我们还为 BEAM 实现了高效的自定义 CUDA 内核，确保与 vLLM 推理框架的无缝集成。实验表明，BEAM 在保留原始模型超过 98% 性能的同时，将 MoE 层的 FLOPs 最多减少 85%，实现了高达 2.5 倍的解码加速和 1.4 倍的吞吐量提升，证明了其作为高效 MoE 推理的实用即插即用解决方案的有效性。BEAM 的代码实现可在 https://github.com/Time-Rune/BEAM 找到。

## 1 引言
混合专家（MoE）通过稀疏激活实现高效扩展，其中每个令牌仅由一小部分专门的前馈网络（FFN）专家处理（Yang et al., 2025a (https://arxiv.org/html/2605.14438#bib.bib1); Liu et al., 2024a (https://arxiv.org/html/2605.14438#bib.bib2); Jiang et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib3)）。
参考图注图1：Qwen3-30B-A3B 上 BEAM 与基线方法的性能-稀疏度权衡。

专家选择的主流范式是固定 Top-K 路由机制，它为每个令牌选择路由 logits 最高的 K 个专家（Shazeer et al., 2017 (https://arxiv.org/html/2605.14438#bib.bib4); Lepikhin et al., 2020 (https://arxiv.org/html/2605.14438#bib.bib5)）。虽然简单且被广泛采用，但它忽略了令牌级别的复杂度，导致对简单令牌的冗余计算（Huang et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib6); Zeng et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib7)）。这种低效率最终限制了 MoE 模型更快推理的潜力。

为了解决固定 Top-K 路由的低效问题，近期工作探索了动态专家激活，分为三类。第一类修改路由 logits 以实现令牌自适应的专家数量（Huang et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib6); Lu et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib8); Yang et al., 2024b (https://arxiv.org/html/2605.14438#bib.bib9); Aghdam et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib10); Guo et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib11)），但未能跳过冗余的高权重专家，并强制设置最小激活阈值，限制了可实现的稀疏度。第二类引入特殊专家，例如零计算空专家来控制稀疏度（Zeng et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib7); Jin et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib12); Gui et al., 2025 (https://arxiv.org/html/2605.14438#bib.bib13)），但需要额外的超参数和复杂的微调过程，并且只能实现被动、间接的稀疏度控制。第三类静态合并或剪枝专家（Chen et al., 2025 (https://arxiv.org/html/2605.14438#bib.bib14); Liu et al., 2024b (https://arxiv.org/html/2605.14438#bib.bib15); Yang et al., 2024a (https://arxiv.org/html/2605.14438#bib.bib16)），但无法在推理时适应输入复杂度，且在高稀疏度下常常性能严重下降。
参考图注图2：普通 Top-K 与 BEAM 对比：BEAM 在 Top-K 候选集上学习二元掩码以实现令牌自适应激活。

在这项工作中，我们提出 BEAM（Binary Expert Activation Masking），一种新颖的动态路由框架，旨在实现 MoE 模型中的极致专家稀疏性和推理加速。如图 2 (https://arxiv.org/html/2605.14438#S1.F2) 所示，BEAM 引入一个轻量级可学习掩码路由器，生成应用于主路由器 Top-K 候选专家的二元掩码，选择性停用冗余专家。通过辅助正则化损失鼓励稀疏性，并使用直通估计器（STE）（Bengio et al., 2013 (https://arxiv.org/html/2605.14438#bib.bib18)）通过二元掩码传播梯度。关键是，BEAM 将稀疏度控制与专家选择解耦。主路由器仍负责负载均衡和专家选择，而掩码路由器仅决定激活数量。这种分离避免了冲突，并在 Top-K 候选集内实现更多激活模式，提供固定 Top-K 或基于 logits 的方法无法实现的细粒度、令牌自适应的稀疏度控制。

为了展示实际影响，我们通过自定义 CUDA 内核将 BEAM 集成到 vLLM（Kwon et al., 2023 (https://arxiv.org/html/2605.14438#bib.bib17)）中，仅需一行代码更改，即可带来显著的现实加速，使 BEAM 成为高效 MoE 部署的实用即插即用解决方案。我们的贡献总结如下：
- • 我们提出 BEAM，一种新颖的动态路由框架，通过可学习掩码路由器实现极致的专家稀疏性。它直接从 Top-K 集中剪枝冗余专家以实现令牌自适应计算，与现有的间接或事后方法形成对比。
- • 我们通过自定义 CUDA 内核将 BEAM 集成到 vLLM 中，提供了一个实用、即插即用的部署方案，仅需极少的代码更改。
- • 大量实验表明，BEAM 在保留超过 98% 性能的同时，将 MoE 层 FLOPs 最多减少 85%（图 1 (https://arxiv.org/html/2605.14438#S1.F1)），实现了 1.4 倍的吞吐量提升和 2.5 倍的解码加速。

## 2 相关工作
**路由 Logits 修改**
这类方法修改路由 logits 以实现令牌自适应专家数量。MoE-Dynamic (Huang et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib6)) 和 XMoE (Yang et al., 2024b (https://arxiv.org/html/2605.14438#bib.bib9)) 激活专家直到累积概率超过阈值。DTop-p (Jin et al., 2025 (https://arxiv.org/html/2605.14438#bib.bib20)) 通过用可学习稀疏度控制器替换固定阈值来改进 MoE-Dynamic。Adaptive Gating (Li et al., 2023b (https://arxiv.org/html/2605.14438#bib.bib19)) 和 NAEE (Lu et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib8)) 根据前两个 logits 之间的差距动态切换 Top-1 和 Top-2。DA-MoE (Aghdam et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib10)) 从注意力分数计算令牌重要性来分配动态 Top-K。DynMoE (Guo et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib11)) 用每个专家的 sigmoid 门控替换 softmax 路由器。MaskMoE (Su et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib37)) 采用基于预训练数据分布的静态词汇表掩码来改进稀有令牌的专家分配。然而，大多数方法依赖于未经验证的启发式假设——路由 logits 的低熵意味着所需专家更少，它们无法跳过冗余的高权重专家，并且要求至少有一个活跃专家，从而阻碍了加速。

**特殊专家**
这类方法通过将令牌路由到不产生计算的专家来减少 FLOPs。AdaMoE (Zeng et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib7)) 引入输出为零的空专家。LongCat (Gui et al., 2025 (https://arxiv.org/html/2605.14438#bib.bib13)) 使用零计算专家，将输入直接作为输出返回。MoE++ (Jin et al., 2024 (https://arxiv.org/html/2605.14438#bib.bib12)) 用三种类型的零计算专家扩展了这一思想。然而，这些方法引入了额外的超参数，并且通过被动的占位符路由而非显式专家最小化间接实现稀疏度，损害了即插即用性。

**静态专家合并与剪枝**
这些免训练方法通过合并或剪枝专家来减少冗余。DEK (Zhang et al., 2025 (https://arxiv.org/html/2605.14438#bib.bib21)) 在特征空间中分组相似专家并合并每组内的专家。EEP (Liu et al., 2024b (https://arxiv.org/html/2605.14438#bib.bib15)) 利用无梯度进化搜索确定剪枝和合并模式。MC-SMoE (Li et al., 2023c (https://arxiv.org/html/2605.14438#bib.bib22)) 利用路由统计指导专家合并，并将合并后的专家分解为低秩和结构化稀疏替代方案。HC-SMoE (Chen et al., 2025 (https://arxiv.org/html/2605.14438#bib.bib14)) 对专家输出应用层次聚类来合并专家。然而，这些方法无法在推理时适应输入令牌的不同复杂度，并且在高压縮下常常性能下降。

## 3 方法
### 3.1 预备知识与动机
MoE 用 N 个专家网络 {E₁, ..., E_N} 和一个路由器 R 替换密集 FFN 层。对于输入令牌 x ∈ ℝ^{d_h}，路由器计算 logits r = R(x) ∈ ℝ^N，并通过 softmax 转换为路由权重。在标准 Top-K 路由下，只有路由 logits 最大的 K 个专家被激活。具体来说，Top-K(·) 操作符保留 r 中最大的 K 个值，并将其余条目设为 −∞，得到路由权重：
g_i = Softmax(Top-K(r))_i.   (1)
MoE 输出是专家输出的加权和：
y = ∑_{i=1}^N g_i · E_i(x),   (2)
其中每个专家 E_i 通常遵循门控线性单元（GLU）结构：
E_i(x) = (δ(x W_gate^(i)) ⊙ (x W_up^(i))) W_down^(i).   (3)

尽管 Top-K 路由实现了可扩展训练，但它为所有令牌分配了统一的计算预算，导致简单令牌的冗余。现有的动态路由方法试图解决这个问题，但在实践中仍然受限。首先，这些方法隐含地将路由排名视为专家重要性的代理。然而，对于给定令牌，排名较低的专家可能仍然关键，而高权重专家可能是冗余的——这在第 5.2 节 (https://arxiv.org/html/2605.14438#S5.SS2) 和附录 B.4 (https://arxiv.org/html/2605.14438#A2.SS4) 中得到了实验验证。其次，累积概率阈值和空专家无法主动剪枝冗余专家，限制了压缩比（第 4.2 节 (https://arxiv.org/html/2605.14438#S4.SS2)）。第三，这些方法将专家选择、负载均衡和稀疏度控制纠结在单个路由器中，造成固有的梯度冲突，从而降低了模型容量（第 4.2 节 (https://arxiv.org/html/2605.14438#S4.SS2)）。

### 3.2 BEAM：二元专家激活掩码
参考图注图3：我们提出的 BEAM 方法示意图，以 4 个专家且 K=3 为例。

上述动机催生了 BEAM，它通过引入一个轻量级可学习的掩码路由器，生成二元掩码以从标准 Top-K 候选集中选择性停用冗余专家，实现令牌自适应的专家激活，如图 3 (https://arxiv.org/html/2605.14438#S3.F3) 所示。形式上，给定输入令牌嵌入 x ∈ ℝ^{d_h}，BEAM 分四步操作。

**步骤 1：标准 Top-K 路由。** 主路由器 R 计算 logits r = R(x) ∈ ℝ^N，其中 N 是专家总数。Top-K(·) 操作符保留最大的 K 个值，并将其余设为 −∞。归一化的路由权重计算如下：
g_i = Softmax(Top-K(r))_i,   i = 1, ..., N,   (4)
其中仅对前 K 个专家有 g_i > 0，且 ∑_{i=1}^N g_i = 1。

**步骤 2：原始掩码生成。** 一个轻量级辅助掩码路由器，参数为 W_m ∈ ℝ^{d_h × N}，处理相同的输入令牌 x 以生成原始掩码 m̂。我们应用 Sigmoid 激活 σ 将原始掩码值约束到 (0, 1) 范围：
m̂ = σ(x W_m).   (5)
m̂ 反映了模型对当前令牌是否需要每个专家的置信度。

**步骤 3：二元掩码化。** 我们使用固定阈值 τ = 0.5 对原始掩码 m̂ 进行二值化，得到离散掩码 m ∈ {0, 1}^N：
m_i = { 1, 如果 m̂_i ≥ 0.5, 0, 否则 }.   (6)
因为 m_i = 0 会禁用专家 i（无论其 Top-K 状态如何），每个令牌激活的专家数量可能减少到 0。

**步骤 4：掩码聚合。** 最终路由权重 ĝ 通过 Top-K 权重 g 与二元掩码 m 的元素级相乘获得：
ĝ = g ⊙ m,   (7)
层输出通过聚合掩码激活计算：
y = ∑_{i=1}^N ĝ_i · E_i(x).   (8)

这种设计提供了三个关键优势。第一，它将路由与稀疏化解耦，即主路由器处理专家选择和负载均衡，而掩码路由器仅专注于冗余消除，避免了冲突的优化目标。第二，专家稀疏性通过端到端学习实现，无需手动调整，从而在保持模型能力的同时实现激进的专家缩减。第三，二元掩码提供了硬件友好的信号，可直接由自定义 CUDA 内核利用，便于高效的实际部署。

### 3.3 训练策略
BEAM 通过两个关键组件进行端到端训练。第一个是直通估计器（STE），用于处理不可微的二值化操作。第二个是辅助稀疏性正则化损失，添加到标准 MoE 目标中，以联合优化任务性能、专家负载均衡和计算效率。

#### 3.3.1 直通估计