BEAM: Binary Expert Activation Masking for Dynamic Routing in MoE

📄 arXiv: 2605.14438v1 📥 PDF

作者: Juntong Wu, Jialiang Cheng, Qishen Yin, Yue Dai, Yuliang Yan, Fuyu Lv, Ou Dan, Li Yuan

分类: cs.AI

发布日期: 2026-05-14

备注: 22 pages, 12 figures

💡 一句话要点

提出BEAM：一种二元专家激活掩码方法，用于MoE中的动态路由，提升推理效率。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 动态路由 二元掩码 模型压缩 推理加速

📋 核心要点

  1. 传统MoE模型采用固定Top-K路由策略，导致计算冗余和推理延迟，现有加速方法要么需要昂贵的重训练，要么在高稀疏度下性能显著下降。
  2. BEAM通过学习token自适应的二元掩码来实现动态专家选择，利用straight-through estimator和辅助损失函数在训练过程中诱导专家稀疏性。
  3. 实验表明，BEAM在显著降低FLOPs的同时，保持了模型性能，并实现了更快的解码速度和更高的吞吐量，验证了其有效性。

📝 摘要（中文）

本文提出了一种名为BEAM（Binary Expert Activation Masking）的新方法，旨在解决MoE模型中固定Top-K路由策略导致的冗余计算和次优推理延迟问题。BEAM通过可训练的二元掩码学习token自适应的专家选择。利用straight-through estimator和辅助正则化损失，BEAM在端到端训练中诱导动态专家稀疏性，同时保持模型能力。此外，本文还为BEAM实现了一个高效的自定义CUDA kernel，确保与vLLM推理框架的无缝集成。实验结果表明，BEAM在减少高达85%的MoE层FLOPs的同时，保留了原始模型超过98%的性能，实现了高达2.5倍的解码速度提升和1.4倍的吞吐量提升，证明了其作为高效MoE推理的实用即插即用解决方案的有效性。

🔬 方法详解

问题定义：MoE模型中的Top-K路由策略在推理时存在计算冗余，导致效率低下。现有的加速方法，如剪枝或量化，通常需要重新训练模型，或者在高稀疏度下由于训练和推理之间的不匹配而导致性能下降。因此，需要一种能够在不显著降低模型性能的前提下，动态地选择专家并减少计算量的方法。

核心思路：BEAM的核心思路是通过学习二元掩码来动态地选择每个token激活的专家。这种方法允许模型根据输入token的重要性自适应地选择专家，从而减少不必要的计算。通过引入可训练的二元掩码，BEAM能够在训练过程中学习到最优的专家选择策略。

技术框架：BEAM方法主要包含以下几个关键模块：1) 二元掩码生成器：该模块负责为每个token生成一个二元掩码，用于指示哪些专家应该被激活。2) 专家激活模块：该模块根据二元掩码选择激活的专家，并进行后续的计算。3) Straight-through Estimator (STE)：由于二元掩码的离散性，直接反向传播梯度会存在问题。STE通过在正向传播时使用二元值，在反向传播时传递连续梯度来解决这个问题。4) 辅助正则化损失：为了鼓励掩码的稀疏性，引入辅助正则化损失，促使模型选择更少的专家。

关键创新：BEAM的关键创新在于引入了可训练的二元掩码，实现了token自适应的专家选择。与传统的Top-K路由策略相比，BEAM能够更精细地控制每个token激活的专家数量，从而减少计算冗余。此外，BEAM通过STE和辅助正则化损失，实现了端到端的训练，避免了训练和推理之间的不匹配问题。

关键设计：BEAM的关键设计包括：1) 二元掩码生成器：可以使用简单的线性层或更复杂的神经网络来生成二元掩码。2) Straight-through Estimator：在反向传播时，梯度直接通过sigmoid函数的梯度。3) 辅助正则化损失：可以使用L1正则化或KL散度等方法来鼓励掩码的稀疏性。具体来说，论文中可能使用了L1正则化损失，其系数需要根据具体任务进行调整。4) CUDA Kernel优化：为了进一步提高推理效率，论文还专门设计了CUDA Kernel，以加速二元掩码的计算和专家激活过程。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，BEAM在减少高达85%的MoE层FLOPs的同时，保留了原始模型超过98%的性能。在推理速度方面，BEAM实现了高达2.5倍的解码速度提升和1.4倍的吞吐量提升。这些结果表明，BEAM是一种高效且实用的MoE推理加速方法，能够在不显著降低模型性能的前提下，显著提高推理效率。

🎯 应用场景

BEAM方法可广泛应用于各种需要高效MoE推理的场景，如大型语言模型部署、在线翻译、语音识别等。通过减少计算量和提高推理速度，BEAM能够降低部署成本，提升用户体验，并促进MoE模型在资源受限环境中的应用。未来，BEAM可以进一步扩展到其他类型的模型和任务中，例如视觉Transformer等。

📄 摘要（原文）

Mixture-of-Experts (MoE) architectures enhance the efficiency of large language models by activating only a subset of experts per token. However, standard MoE employs a fixed Top-K routing strategy, leading to redundant computation and suboptimal inference latency. Existing acceleration methods either require costly retraining with architectural changes or suffer from severe performance drop at high sparsity due to train-inference mismatch. To address these limitations, we propose BEAM (Binary Expert Activation Masking), a novel method that learns token-adaptive expert selection via trainable binary masks. With a straight-through estimator and an auxiliary regularization loss, BEAM induces dynamic expert sparsity through end-to-end training while maintaining model capability. We further implement an efficient custom CUDA kernel for BEAM, ensuring seamless integration with the vLLM inference framework. Experiments show that BEAM retains over 98\% of the original model's performance while reducing MoE layer FLOPs by up to 85\%, achieving up to 2.5$\times$ faster decoding and 1.4$\times$ higher throughput, demonstrating its effectiveness as a practical, plug-and-play solution for efficient MoE inference.