BEAM: Enmascaramiento Binario de Activación de Expertos para Enrutamiento Dinámico en MoE

Resumen

Las arquitecturas de Mezcla de Expertos (MoE) mejoran la eficiencia de los modelos de lenguaje grandes al activar solo un subconjunto de expertos por token. Sin embargo, el MoE estándar emplea una estrategia de enrutamiento Top-K fija, lo que genera cómputo redundante y una latencia de inferencia subóptima. Los métodos de aceleración existentes requieren un reentrenamiento costoso con cambios arquitectónicos o sufren una grave caída del rendimiento en alta dispersión debido al desajuste entre entrenamiento e inferencia. Para abordar estas limitaciones, proponemos BEAM (Binary Expert Activation Masking, enmascaramiento binario de activación de expertos), un método novedoso que aprende una selección de expertos adaptativa por token mediante máscaras binarias entrenables. Con un estimador de paso directo y una pérdida de regularización auxiliar, BEAM induce una dispersión dinámica de expertos a través del entrenamiento de extremo a extremo mientras mantiene la capacidad del modelo. Además, implementamos un kernel CUDA personalizado y eficiente para BEAM, garantizando una integración fluida con el marco de inferencia vLLM. Los experimentos muestran que BEAM retiene más del 98% del rendimiento del modelo original mientras reduce los FLOPs de la capa MoE hasta en un 85%, logrando una decodificación hasta 2.5 veces más rápida y un rendimiento 1.4 veces mayor, lo que demuestra su eficacia como una solución práctica y plug-and-play para la inferencia MoE eficiente.

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.