BEAM让专家模型动态选择专家,提速近2.5倍且性能损失低于2%。
BEAM: Binary Expert Activation Masking for Dynamic Routing in MoE

- 用可学习的二值掩码实现按令牌动态选专家。
- 保持98%以上原模型性能,摩尔层计算量减少85%。
- 无需重训练,适配vLLM框架,适合部署优化场景。
混合专家(MoE)架构通过每标记仅激活部分专家来提升大语言模型效率。然而,标准MoE采用固定的Top-K路由策略,导致冗余计算和推理延迟偏高。现有加速方法要么需昂贵的重新训练与结构修改,要么在高稀疏度下因训练-推理不一致而出现严重性能下降。为此,我们提出BEAM(二值专家激活掩码),一种通过可训练二值掩码实现令牌自适应专家选择的新方法。结合直通估计器与辅助正则化损失,BEAM在端到端训练中诱导动态专家稀疏性,同时保持模型能力。我们进一步开发了针对BEAM的高效自定义CUDA内核,确保与vLLM推理框架无缝集成。实验表明,BEAM在保留超过98%原始模型性能的同时,将MoE层浮点运算量减少高达85%,实现最高2.5倍的解码速度提升与1.4倍更高的吞吐量,证明其作为高效、即插即用的MoE推理解决方案的有效性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。