BEAM: Бинарное маскирование активации экспертов для динамической маршрутизации в MoE.

Аннотация

Архитектуры смеси экспертов (MoE) повышают эффективность больших языковых моделей, активируя только подмножество экспертов для каждого токена. Однако стандартная MoE использует фиксированную стратегию маршрутизации Top-K, что приводит к избыточным вычислениям и неоптимальной задержке инференса. Существующие методы ускорения либо требуют дорогостоящего переобучения с изменением архитектуры, либо страдают от значительного падения производительности при высокой разреженности из-за несоответствия между обучением и инференсом. Для устранения этих ограничений мы предлагаем BEAM (Binary Expert Activation Masking) — новый метод, который обучает адаптивное к токенам выделение экспертов с помощью обучаемых бинарных масок. Используя оценщик прямого прохода (straight-through estimator) и вспомогательную функцию потерь регуляризации, BEAM обеспечивает динамическую разреженность экспертов в ходе сквозного обучения, сохраняя при этом возможности модели. Мы также реализовали эффективное пользовательское ядро CUDA для BEAM, обеспечивающее бесшовную интеграцию с фреймворком инференса vLLM. Эксперименты показывают, что BEAM сохраняет более 98% производительности исходной модели, сокращая при этом FLOPs в слоях MoE до 85%, что позволяет достичь до 2,5-кратного ускорения декодирования и в 1,4 раза более высокой пропускной способности, демонстрируя свою эффективность как практическое решение типа «plug-and-play» для эффективного инференса MoE.

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.