BEAM: 이진 전문가 활성화 마스킹을 통한 MoE의 동적 라우팅

BEAM: Binary Expert Activation Masking for Dynamic Routing in MoE

Mixture-of-Experts (MoE) 아키텍처는 각 토큰에 대해 일부 전문가만 활성화하여 대규모 언어 모델의 효율성을 향상시킵니다. 그러나 기존의 MoE는 고정된 Top-K 라우팅 전략을 사용하므로, 불필요한 계산이 발생하고 최적의 추론 지연 시간이 확보되지 않습니다. 기존의 가속화 방법은 비싼 재학습과 아키텍처 변경을 요구하거나, 학습-추론 불일치로 인해 높은 희소성에서 심각한 성능 저하를 초래합니다. 이러한 한계를 해결하기 위해, 우리는 토큰에 적응적인 전문가 선택을 학습하는 학습 가능한 이진 마스크를 사용하는 새로운 방법인 BEAM (Binary Expert Activation Masking)을 제안합니다. Straight-through estimator와 보조 정규화 손실을 사용하여, BEAM은 모델의 능력을 유지하면서 엔드-투-엔드 학습을 통해 동적인 전문가 희소성을 유도합니다. 또한, BEAM을 위한 효율적인 사용자 정의 CUDA 커널을 구현하여 vLLM 추론 프레임워크와의 원활한 통합을 보장합니다. 실험 결과, BEAM은 원래 모델 성능의 98% 이상을 유지하면서 MoE 레이어의 FLOPs를 최대 85%까지 줄이고, 최대 2.5배 빠른 디코딩 속도와 1.4배 더 높은 처리량을 달성하여, 효율적인 MoE 추론을 위한 실용적이고 간편하게 적용 가능한 솔루션임을 입증합니다.

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.