BEAM: Mascaramento Binário de Ativação de Especialistas para Roteamento Dinâmico em MoE

Resumo

Arquiteturas Mistura de Especialistas (MoE) melhoram a eficiência de modelos de linguagem grandes ao ativar apenas um subconjunto de especialistas por token. No entanto, o MoE padrão emprega uma estratégia de roteamento Top-K fixa, levando a computação redundante e latência de inferência subótima. Métodos de aceleração existentes ou exigem um retreinamento custoso com mudanças arquiteturais, ou sofrem uma queda severa de desempenho em alta esparsidade devido à incompatibilidade entre treino e inferência. Para superar essas limitações, propomos BEAM (Máscara Binária de Ativação de Especialistas), um método inovador que aprende a seleção de especialistas adaptativa a tokens por meio de máscaras binárias treináveis. Com um estimador de passagem direta e uma perda de regularização auxiliar, o BEAM induz esparsidade dinâmica de especialistas através de treinamento de ponta a ponta, mantendo a capacidade do modelo. Implementamos ainda um kernel CUDA personalizado e eficiente para o BEAM, garantindo integração perfeita com o framework de inferência vLLM. Experimentos mostram que o BEAM retém mais de 98% do desempenho do modelo original enquanto reduz os FLOPs da camada MoE em até 85%, alcançando uma decodificação até 2,5 vezes mais rápida e um throughput 1,4 vezes maior, demonstrando sua eficácia como uma solução prática plug-and-play para inferência 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.