dMoE: dLLMs with Learnable Block Experts

Technical Summary: dMoE – Diffusion Large Language Models with Learnable Block Experts

Problem Statement
Diffusion Large Language Models (dLLMs) have emerged as a competitive alternative to autoregressive models, offering the advantage of parallel decoding. To scale capacity, dLLMs are increasingly integrated with Mixture-of-Experts (MoE) architectures. However, a fundamental mismatch exists between the block-parallel decoding nature of dLLMs and the token-level expert selection paradigm inherited from autoregressive MoE models. In dLLMs, a single forward pass processes a block of tokens with bidirectional dependencies, whereas conventional MoE layers route each token independently. This independence causes a dramatic increase in the number of uniquely activated experts within a single forward pass. Empirical analysis reveals that this explosion in unique expert activation makes memory access the primary inference bottleneck, with MoE latency dominating end-to-end inference time and increasing linearly with the number of unique experts.

Methodology: dMoE
The authors propose dMoE, a block-level MoE framework designed to reduce the number of uniquely activated experts without altering the number of experts selected per token. The method is grounded in two key empirical observations:

1. Token-level expert scores (router weights) serve as effective indicators of expert importance.
2. Expert concentration varies significantly across different denoising steps and blocks.

Based on these observations, dMoE employs a coarse-to-fine routing strategy:

• Aggregation: Instead of routing tokens individually, dMoE first aggregates token-level expert scores within a block to generate a unified block-level expert distribution.
• Adaptive Coreset Selection: Rather than enforcing a fixed coreset size, the method normalizes the block-level scores and applies a top-p criterion to select an adaptive expert coreset. This allows the coreset size to dynamically adjust based on the concentration of the expert distribution (smaller coreset for concentrated distributions, larger for dispersed ones).
• Restricted Routing: The original token-level routing is then constrained to occur only within this selected coreset.
• Training Strategy: During training, dMoE utilizes a self-distillation paradigm, where the model generates supervision targets for itself using the same block-level routing procedure, ensuring alignment between training and inference.

Key Contributions

• Identification of the Bottleneck: The paper establishes that memory-bound latency caused by excessive unique expert activation is the primary efficiency barrier in MoE dLLMs, distinct from the challenges in autoregressive models.
• Block-Level Routing Framework: dMoE introduces a novel mechanism to shift from token-level to block-level expert selection, effectively compressing the unique expert count while maintaining the computational load per token.
• Dynamic Adaptability: By utilizing a top-p criterion, the method dynamically adapts to varying routing characteristics across different denoising steps and blocks, avoiding the rigidity of fixed-size expert selection.

Experimental Results
Evaluated on the state-of-the-art open-source model LLaDA2.0-mini across four benchmarks (MATH500, GSM8K, ARC-C, and MMLU), dMoE demonstrates significant efficiency gains with minimal performance loss:

• Expert Compression: On average, dMoE reduces the number of uniquely activated experts from 69.5 to 14.6 (a ~4.77× reduction), achieving a 79.04% reduction in unique expert count.
• Performance Retention: The method retains 99.11% of the original model's performance on average. Even under aggressive compression (reducing experts to near the per-token limit of 8), performance degradation remains negligible.
• Efficiency Gains:
  • Memory Usage: Reduced by 76.64% to 79.84% compared to the original model.
  • Latency: Achieves an end-to-end latency speedup of 1.14× to 1.66×, primarily driven by the reduction in MoE kernel latency.
• Comparison: dMoE outperforms existing baselines (including Top-4 selection and DES variants) in the performance-efficiency trade-off, particularly at high compression ratios where other methods suffer significant accuracy drops.

Significance
The paper positions dMoE as a foundational step toward efficient block-level routing in MoE dLLMs. It demonstrates that the memory bottleneck in diffusion-based MoE models can be effectively mitigated by aggregating routing decisions at the block level. The work establishes a strong baseline for future research in scaling dLLMs, proving that aggressive expert compression is achievable without compromising the model's reasoning capabilities. The method is presented as tunable, allowing practitioners to adjust the number of activated experts to fit specific hardware constraints or application requirements.