PaperlayerView exact revision on arXiv

Source-linked AI summary

BEAM: Binary Expert Activation Masking for Dynamic Routing in MoE

Juntong Wu, Jialiang Cheng, Qishen Yin, Yue Dai, Yuliang Yan, Fuyu Lv, Ou Dan, Li Yuan

arXiv:2605.14438v1cs.AI

TL;DR

Fixed Top-K routing gives all tokens the same expert budget, creating redundant computation, while existing alternatives have practical sparsity or high-sparsity limitations. BEAM learns binary masks over Top-K candidates with STE and sparsity regularization, and integrates with vLLM through a CUDA kernel. It retains over 98% of performance while reducing MoE-layer FLOPs by up to 85%, reaching 2.5× faster decoding and 1.4× higher throughput.

  • Problem

    Fixed Top-K routing assigns a uniform expert budget to tokens, causing redundancy, while existing dynamic methods have sparsity, routing, or train-inference limitations.

  • Method

    BEAM uses a learnable binary mask router over Top-K candidates, trained with a straight-through estimator and auxiliary sparsity regularization.

  • Results

    98% performance retention accompanies up to 85% lower MoE-layer FLOPs, 2.5× faster decoding, and 1.4× higher throughput.

  • Takeaways & Limitations

    BEAM provides a practical, plug-and-play approach for token-adaptive sparse MoE inference through vLLM integration.

  • Takeaways & Limitations

    BEAM is evaluated on three MoE architectures, requires post-training SFT, and has speedups constrained by shared-expert ratios and single-GPU evaluation.

Abstract

from arXiv · show

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

1 Introduction

BEAM addresses redundant computation from fixed Top-K routing by learning token-adaptive masks over candidate experts. It targets high sparsity while preserving performance and supports practical vLLM deployment.

  • Fixed Top-K assigns every token the same expert budget, causing redundant computation for simple tokens and limiting faster inference.
  • Existing dynamic-routing approaches modify logits, add special experts, or statically merge and prune experts, but face sparsity, tuning, adaptability, or performance limitations.
  • BEAM uses a learnable mask router to deactivate redundant experts within the primary router’s Top-K candidates for token-adaptive computation.
  • BEAM separates sparsity control from expert selection: the primary router handles load balancing and choice, while the mask router determines activation count.
  • BEAM integrates into vLLM through a custom CUDA kernel with minimal code changes, providing a practical deployment path.
  • 98% performance retention accompanies up to 85% lower MoE-layer FLOPs, 1.4× higher throughput, and 2.5× faster decoding.

2 Related Work

Related work pursues dynamic expert counts through routing-logit changes, special zero-computation experts, or static expert merging and pruning. These strategies trade off direct sparsity control, simplicity, adaptability, or high-sparsity performance.

  • Routing Logits Modification: Routing-logit methods use cumulative-probability thresholds or logit gaps to select token-adaptive expert counts.
  • Special Experts: Special-expert methods reduce FLOPs with null or zero-computation experts, but introduce extra hyperparameters and indirect sparsity control.
  • Static Expert Merging and Pruning: Static merging and pruning methods reduce redundancy without inference-time adaptation to input complexity.

3 Method

BEAM preserves standard Top-K candidate selection while adding a learnable binary mask that selectively deactivates redundant experts for each token. End-to-end training uses STE and sparsity regularization to learn token-adaptive computation while separating routing, load balancing, and sparsification.

  • Motivation: Fixed Top-K assigns every token the same computational budget, creating redundancy for simple tokens.
  • BEAM: Binary Expert Activation Masking: BEAM applies a lightweight mask router to standard Top-K candidates, selectively deactivating redundant experts through a binary mask.The primary router produces Top-K weights, while the mask router processes the same token to estimate expert necessity.
  • BEAM: Binary Expert Activation Masking: BEAM can reduce activated experts per token to zero because a zero mask disables an expert regardless of its Top-K status.
  • BEAM: Binary Expert Activation Masking: BEAM separates expert selection and load balancing from redundancy elimination, allowing the mask router to focus exclusively on sparsification.The primary router handles selection and load balancing, while the mask router selectively suppresses redundant candidates.
  • Training Strategy: Training combines the language-modeling, load-balancing, and sparsity losses, with STE enabling backpropagation through binary masking.The sparsity regularizer applies the L1 norm of raw mask values within the Top-K candidate set.
  • Theoretical Analysis: Only Top-K experts receive mask-router learning signals, while non-selected experts receive zero gradient.For selected experts, task loss and sparsity pressure jointly determine retention or pruning; β controls this trade-off.

4 Experiments

BEAM is evaluated across three MoE models, multiple sparsity levels, downstream benchmarks, and inference settings. It preserves accuracy while achieving substantial sparsity and consistent serving acceleration, with ablations supporting its training design.

  • Experimental Setup: BEAM is evaluated on Qwen1.5-MoE-A2.7B, DeepSeekV2-Lite, and Qwen3-30B-A3B using shared supervised fine-tuning configurations.The evaluation uses the Tulu 3 SFT Mixture Dataset and identical training settings for baselines and BEAM.
  • Evaluation: The study measures accuracy across eight reasoning, coding, knowledge, and common-sense benchmarks, alongside TPOT, TTFT, and throughput under vLLM serving.Acceleration tests use fixed 128/32-token input/output lengths and 5,000 samples on one NVIDIA H20 GPU.
  • Performance Comparison: BEAM preserves over 98% of original accuracy at mid sparsity while reducing average activated experts per token by 47%–61% across all three models.At high sparsity, average activated experts fall as low as 14% of the original while retaining over 95% accuracy.
  • Performance Comparison: At extreme sparsity, BEAM outperforms Top-K Reduced by 32.49% on DeepSeekV2 and 33.29% on Qwen3, while Qwen1.5 retains 85% performance at Avg-K = 0.11.The Qwen1.5 result indicates that most tokens bypass routed experts.
  • Ablation Study: Increasing β smoothly improves sparsity with gradual accuracy loss, while L1 regularization and STE-based binary masking outperform alternative training variants.Removing regularization increases expert activation, L2 regularization is inferior to L1, and soft-mask training underperforms binary masking.
  • Acceleration Comparison: BEAM achieves at least 1.3× better TPOT and over 1.1× gains in TTFT and throughput, reaching up to 2.5× faster decoding on DeepSeek-V2-Lite at QPS=24.Qwen3 permits up to 85% MoE FLOPs reduction because it has no shared experts, whereas Qwen1.5 has a 50% architectural limit.

5 Analysis

BEAM allocates different numbers of experts to different tokens and positions. The analysis links higher activation to semantic richness and identifies layer- and phase-dependent usage patterns.

  • Token-wise Activation: Content words such as nouns and verbs trigger more experts than function words and punctuation.The analysis characterizes this pattern as alignment between activation and semantic richness.
  • Layer-wise Analysis: DeepSeek uses nearly identical expert counts during prefill and decoding, whereas Qwen1.5 and Qwen3 use more experts during decoding.The comparison is based on average activated experts per layer across 1,000 randomly sampled inputs.
  • Layer-wise Analysis: Qwen1.5 and Qwen3 develop an encoder-decoder-like pattern in which shallower layers primarily support knowledge storage and deeper layers allocate more expert capacity.The supplied analysis describes this pattern across prefill and decode measurements.

6 Conclusion

BEAM provides token-adaptive sparsity by selectively deactivating redundant experts within the Top-K set without modifying model architecture. Its vLLM integration delivers substantial decoding and throughput gains while retaining accuracy.

  • Conclusion: BEAM uses a lightweight mask router to deactivate redundant Top-K experts and enable token-adaptive sparsity without architectural modification.The framework is integrated into vLLM through an efficient CUDA kernel.
  • Conclusion: BEAM delivers up to 2.5× faster decoding and 1.4× higher throughput while retaining over 98% accuracy.The conclusion presents these results as evidence of practical MoE inference acceleration.

A.2 Limitations

BEAM’s evaluation and implementation span three MoE architectures, with zero activation handled through shared-expert computation or residual bypass. Its vLLM integration masks experts before token grouping, while several scope and deployment limitations remain.

  • Limitations: BEAM’s effectiveness remains unvalidated on MoE designs with different gating mechanisms or expert granularities.This is the first limitation identified by the authors.
  • Limitations: BEAM requires post-training SFT, and its speedup depends on the shared-expert ratio and has been benchmarked only on single-GPU settings.The authors identify additional training cost, reduced benefit from shared experts, and unresolved multi-GPU interactions.
  • Zero activation: When all routed experts are skipped, shared-expert architectures perform shared-expert-only computation, while architectures without shared experts bypass the MoE layer through the residual path.This zero-activation behavior corresponds to dynamic layer skipping.
  • Inference implementation: BEAM’s vLLM implementation marks masked expert indices as −1 and excludes them during expert-wise token grouping and block alignment.The changes extend two existing CUDA kernels while preserving compatibility with vLLM optimizations.
  • Evaluation scope: BEAM is evaluated on Qwen1.5-MoE-A2.7B, DeepSeekV2-Lite, and Qwen3-30B-A3B, covering different routed and shared-expert configurations.The models use 60, 64, and 128 experts, respectively, with distinct numbers of activated shared and routed experts.

B.1.3 Benchmarks

The benchmark suite evaluates accuracy across reasoning, knowledge, and commonsense tasks, while acceleration tests measure latency and throughput under fixed workloads on a single GPU.

  • Accuracy benchmarks: Accuracy evaluation uses OpenCompass tasks spanning reasoning, knowledge, and commonsense domains.The listed tasks include MATH, GSM8K, Human Eval, MMLU, CEVAL, CMMLU, CommonsenseQA, and BoolQ.
  • Acceleration benchmarks: Acceleration evaluation records TPOT across QPS levels, TTFT at 32 QPS, and offline throughput using vLLM.Each model runs on a single GPU with fixed 128-token inputs, 32-token outputs, and 5,000 samples.

B.2 Training Dynamics

During BEAM training, language-modeling loss converges near the standard SFT baseline while expert activity rapidly sparsifies, and DynMoE shows instability in post-training comparison.

  • Training dynamics: BEAM’s language-modeling loss converges to a level comparable to standard SFT across all three models.The authors interpret this as evidence that masking and sparsity regularization do not compromise model capacity.
  • Training dynamics: The active rate drops sharply from near 100% to a stable plateau during approximately the first 0.5 epoch.Later training focuses on optimizing the language-modeling objective under the learned sparsity pattern.
  • Baseline comparison: DynMoE activates 61.66 experts on Qwen3-30B-A3B, 30.06 on Qwen1.5-MoE-A2.7B, and 30.50 on DeepSeekV2-Lite, exceeding their Top-K budgets.The corresponding Top-K values are 8, 4, and 6.
  • Baseline comparison: DeepSeekV2-Lite accuracy falls from 55.15 to 3.59 under DynMoE, while BEAM retains over 98% of original-model performance across all three models.The passage attributes DynMoE’s instability to replacing the original router architecture in post-training.

B.4 Layer-wise Masking Rank Analysis

BEAM’s masking decisions are token-dependent rather than determined solely by routing rank: highly ranked experts may be pruned, while lower-ranked experts may be retained. Expert utilization remains relatively balanced after fine-tuning.

  • Layer-wise masking rank analysis: Across all layers and models, the minimum masked rank reaches 1–3, showing that highly ranked experts can be pruned for individual tokens.Masking therefore does not simply remove the lowest-ranked experts.
  • Layer-wise masking rank analysis: The maximum kept rank extends to the lower end of the Top-K range, indicating that low-ranked experts can be retained when critical.The shaded overlap zone represents token-dependent masking decisions.
  • Expert load balance: BEAM applies relatively uniform masking across experts and maintains balanced expert loads, including in Qwen3-30B-A3B with 128 experts.The comparison uses expert utilization rates before and after BEAM fine-tuning.

B.6 Task-specific Inference Speed Analysis

BEAM delivers consistent inference speedups across evaluation tasks, with visualizations showing token- and layer-dependent expert activation patterns that differ across MoE models.

  • The task-specific acceleration results are summarized in Table 9 for Qwen3-30B-A3B.
  • BEAM achieves consistent speedups across all evaluated tasks on Qwen3-30B-A3B using vLLM on a single NVIDIA H20 GPU.The evaluation measures inference speed for both the BEAM-augmented model and its baseline, covering prefill and decoding.
  • Activation patterns differ significantly across Qwen1.5-MoE-A2.7B, DeepSeekV2-Lite, and Qwen3-30B-A3B across tokens and layers.The visualizations use the same prompt across models and layers.
  • Chat-template tokens such as “You are a helpful assistant” activate almost no experts across all three models.
  • Figures 10–12 encode the number of activated experts for each token at each layer as heatmap cells.Tokens are shown vertically and layers horizontally.
Loading 2605.14438v1…