BEAM : Masquage binaire d'activation des experts pour le routage dynamique dans les MoE

Résumé

Les architectures Mixture-of-Experts (MoE) améliorent l'efficacité des grands modèles de langage en n'activant qu'un sous-ensemble d'experts par jeton. Cependant, les MoE standards utilisent une stratégie de routage Top-K fixe, ce qui entraîne des calculs redondants et une latence d'inférence sous-optimale. Les méthodes d'accélération existantes nécessitent soit un réentraînement coûteux avec des modifications architecturales, soit souffrent d'une forte baisse de performance à haute parcimonie en raison d'un décalage entre l'entraînement et l'inférence. Pour remédier à ces limitations, nous proposons BEAM (Binary Expert Activation Masking), une méthode novatrice qui apprend une sélection d'experts adaptative aux jetons via des masques binaires entraînables. Grâce à un estimateur straight-through et une perte de régularisation auxiliaire, BEAM induit une parcimonie dynamique des experts par un apprentissage de bout en bout tout en préservant la capacité du modèle. Nous implémentons également un kernel CUDA personnalisé et efficace pour BEAM, assurant une intégration transparente avec le cadre d'inférence vLLM. Les expériences montrent que BEAM conserve plus de 98 % des performances du modèle original tout en réduisant les FLOPs des couches MoE jusqu'à 85 %, atteignant un décodage jusqu'à 2,5 fois plus rapide et un débit 1,4 fois plus élevé, démontrant ainsi son efficacité en tant que solution pratique et prête à l'emploi pour une inférence MoE efficiente.

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.5times faster decoding and 1.4times higher throughput, demonstrating its effectiveness as a practical, plug-and-play solution for efficient MoE inference.