BEAM: Binary Expert Activation Masking for Dynamic Routing in 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.5times faster decoding and 1.4times higher throughput, demonstrating its effectiveness as a practical, plug-and-play solution for efficient MoE inference.

📋 한 줄 요약

**[MoE 효율 추론 / 동적 라우팅]** 학습 가능한 이진 마스크로 토큰 적응적 전문가 선택을 학습해 MoE LLM 추론을 가속하는 BEAM 제안.

🎯 핵심 기여도

표준 MoE의 고정 Top-K 라우팅이 불필요한 연산과 추론 지연을 유발한다는 문제를 진단하고, 고비용 재학습이나 train-inference 불일치 없는 새 접근을 정의.
학습 가능한 이진 마스크 + straight-through estimator + 보조 정규화 손실을 결합한 BEAM(Binary Expert Activation Masking)으로 종단간 학습 가능한 동적 전문가 sparsity 유도.
효율적 커스텀 CUDA 커널을 구현해 vLLM 추론 프레임워크와 원활히 통합되는 실용성 확보.
원본 모델 성능의 98% 이상 유지, MoE 레이어 FLOPs 최대 85% 감소, 디코딩 2.5배 가속, 처리량 1.4배 향상.

💡 핵심 아이디어

"토큰마다 필요한 전문가 수가 다르다"는 직관을 학습 가능한 이진 마스크로 명시화하고, STE로 미분 가능하게 학습하면 train-inference 불일치 없이 동적 sparsity를 얻을 수 있다.

🔬 기술적 접근법

**모델/방법론**: 표준 MoE의 라우팅을 트레이너블 이진 마스크로 대체하는 BEAM, vLLM 호환 CUDA 커널 구현.
**핵심 기법**: 학습 가능한 마스크에 STE 적용해 미분 가능 학습, 보조 정규화로 sparsity·표현력 균형, 종단간 학습으로 train-inference 일치 보장.

📊 주요 결과

원본 모델 성능 ≥98% 유지.
MoE 레이어 FLOPs 최대 85% 감소.
디코딩 최대 2.5×, 처리량 최대 1.4× 향상.

💭 의의 및 한계

**의의**: 별도의 아키텍처 변경 없이 기존 MoE 모델에 plug-and-play로 적용 가능한 실용적 추론 가속 솔루션 제공. **한계**: 매우 높은 sparsity에서 일부 도메인의 미세 성능 저하 가능성, 학습 추가 비용은 존재.

🚀 실용적 활용

대규모 MoE LLM 서빙 비용 절감.
온디바이스·엣지 환경에서 MoE 추론.
동적 sparsity 연구의 기반 알고리즘.