Source-linked AI summary
dMoE: dLLMs with Learnable Block Experts
Sicheng Feng, Zigeng Chen, Gongfan Fang, Xinyin Ma, Xinchao Wang
TL;DR
MoE dLLMs route tokens independently even though diffusion decoding processes blocks in parallel, activating many unique experts and creating a memory bottleneck. dMoE aggregates token-level expert scores into adaptive block-level routing, reducing expert activation while preserving performance across benchmarks.
Problem
MoE dLLMs independently route tokens within parallel blocks, causing many unique experts to activate and making memory access a primary inference bottleneck.
Method
dMoE aggregates token-level expert scores into block-level scores and uses a top-p-selected coreset to guide token routing adaptively.
Results
Across four benchmarks, dMoE reduces uniquely activated experts by 4.21× while retaining 99.55% of the original performance.
Takeaways & Limitations
dMoE establishes a block-level routing baseline that enables aggressive expert compression with almost no performance degradation in MoE dLLMs.
Takeaways & Limitations
The evaluation is limited to language-modality MoE dLLMs, while extensions to image and video are proposed as future work.
Abstract
from arXiv · showhide
Diffusion Large Language Models (dLLMs) have recently emerged as a promising alternative to autoregressive models, offering competitive performance while naturally supporting parallel decoding. However, as dLLMs are increasingly integrated with Mixture-of-Experts (MoE) architectures to scale model capacity, a fundamental mismatch arises between block parallel decoding and token-level expert selection. Specifically, each dLLM forward pass processes multiple tokens with bidirectional dependencies, whereas conventional MoE layers route each token independently. This mismatch substantially increases the number of uniquely activated experts, making inference increasingly memory-bound. To address this, we propose dMoE, a simple yet effective block-level MoE framework. The central idea of dMoE is to aggregate token-level expert distributions within each block into a unified block-level expert distribution, which is then used to guide expert routing in a more coherent manner. In this way, dMoE substantially reduces the number of uniquely activated experts during inference without sacrificing performance, thereby mitigating the memory-bound bottleneck. Extensive experiments across a variety of benchmarks demonstrate the effectiveness of dMoE. On average, dMoE reduces the number of uniquely activated experts from 69.5 to 14.6 while retaining 99.11% of the original performance. Meanwhile, it reduces memory usage by 76.64% to 79.84% and achieves 1.14$\times$ to 1.66$\times$ end-to-end latency speedup. Code is available at: https://github.com/fscdc/dMoE
1 Introduction
dLLMs support parallel block decoding, but conventional token-level MoE routing creates an efficiency mismatch because multiple tokens are processed together while experts are selected independently. dMoE addresses this by aggregating token-level expert scores into block-level scores to guide routing, reducing uniquely activated experts while preserving performance.
- Motivation: dLLMs iteratively refine mask tokens and naturally support parallel decoding beyond autoregressive left-to-right generation.Recent dLLMs are increasingly integrated with MoE architectures to scale capacity while keeping active parameters manageable.
- Problem: Token-level MoE routing independently selects experts for each token, despite dLLMs processing multiple tokens in one forward pass, creating an efficiency challenge.This mismatch is especially relevant when block diffusion decoding processes a whole token block together.
- Method: dMoE aggregates token-level expert scores into block-level scores and uses them to dynamically guide the original routing process.The design relies on token-level scores as informative expert-importance signals and on expert concentration varying across denoising steps and blocks.
- Results: 4.21×: dMoE reduces the number of uniquely activated experts while retaining 99.55% of the original performance across four benchmarks.Evaluation uses LLaDA2.0-mini on MATH500, GSM8K, ARC-C, and MMLU, with no reported performance degradation.
2 Related Work
Related work spans diffusion language models that generate through masked-token denoising and efficient MoE strategies developed primarily for autoregressive decoding. Recent studies begin adapting MoE efficiency techniques to dLLMs through routing and expert reuse across denoising steps.
- Diffusion Language Models: Masked-token denoising formulates diffusion directly in token space, enabling parallel generation and relaxing autoregressive left-to-right dependencies.This formulation underlies recent diffusion language models (dLLMs).
- Efficient Mixture-of-Expert Strategies: Autoregressive MoE efficiency methods primarily target token-wise expert computation or serving overhead for left-to-right decoding.They broadly include pre-execution compression approaches such as expert pruning and expert merging.
- Efficient Mixture-of-Expert Strategies: MoE efficiency in dLLMs remains underexplored, with recent work providing initial evidence of its potential.The cited approaches address routing choices, load balancing, and expert reuse during denoising.
- Efficient Mixture-of-Expert Strategies: EC-DLM replaces token-choice routing with expert-choice routing and dynamically adjusts expert capacity across denoising steps to improve load balancing.This combines an alternative routing formulation with step-dependent capacity adjustment.
- Efficient Mixture-of-Expert Strategies: TEAM exploits temporal and spatial consistency in expert routing to reuse experts across denoising steps, reducing the overall number of activated experts.Its strategy specifically targets repeated expert activation during the denoising process.
3 Preliminaries
dLLMs generate by iteratively refining masked sequences, enabling non-autoregressive parallel decoding of multiple positions. In MoE dLLMs, simultaneously routing many tokens can activate numerous unique experts, making memory overhead a primary inference bottleneck.
- dLLM Parallel Decoding: dLLMs iteratively refine partially masked sequences, predicting multiple unknown positions simultaneously rather than generating strictly left to right.This supports parallel generation by updating sets of masked tokens in a single denoising step.
- Masked Diffusion Language Models: Masked diffusion language models begin from a corrupted sequence produced by independently masking tokens from a clean sequence.The masking intensity is represented by t ∈[0, 1].
- dLLM Parallel Decoding: All masked positions can be reconstructed in parallel within each denoising step, giving dLLMs non-autoregressive parallel decoding.The denoising model predicts original tokens at masked positions.
- MoE dLLM Bottleneck: MoE dLLM forward passes route multiple tokens simultaneously, such as a full block under block diffusion decoding, activating many unique experts.The activated experts must be iteratively loaded during inference, creating substantial memory overhead.
4 Methods
dMoE reduces block-level expert activation by first aggregating token-level expert scores into a block-level distribution, then selecting an adaptive expert coreset for token routing. Its design is motivated by varying expert concentration across denoising steps and blocks, with training and inference sharing the same self-distilled routing process.
- Design observations: Expert concentration varies substantially across denoising steps and blocks, with consistently large variation in unique expert count across benchmarks.The analysis examines Layer 10 during inference and reports that the variation is not tied to a specific dataset or task.
- dMoE routing: dMoE uses a coarse-to-fine strategy: it selects a coreset from the full expert pool before performing token-level expert selection within that coreset.This targets fewer block-level unique experts while retaining a fixed number of selected experts for each token.
- dMoE routing: Token-level expert scores are directly aggregated across the current block to produce block-level expert scores.The aggregation is motivated by the observed positive relationship between token-level router weights and expert importance.
- Adaptive coreset selection: A normalized top-p criterion selects the block’s expert coreset instead of enforcing a fixed coreset size.The adaptive rule yields smaller coresets for concentrated score distributions and larger coresets for dispersed distributions.
- Training and inference: Training uses self-distillation with the proposed routing in the forward pass, and inference uses the same forward process to maintain alignment.This shared routing procedure is also summarized in the overview of dMoE.
5 Experiments
Across mathematical, scientific, and general reasoning benchmarks, dMoE compresses activated experts while largely preserving performance. It also reduces MoE memory usage and delivers end-to-end inference speedups, including under extreme compression.
- Evaluation: dMoE is evaluated on MATH500, GSM8K, ARC-C, and MMLU using zero-shot step-by-step reasoning prompts.These benchmarks span mathematical, scientific, and general high-level reasoning tasks.
- Main Results: 83.575% average performance retains 99.55% of the original while reducing average unique expert count from 69.5 to 16.575.The average unique expert count falls by 76.15%, while dMoE achieves a 55.11% to 59.62% greater reduction than baselines at comparable performance.
- Efficiency: 76.64% to 79.84% lower memory usage is achieved versus the original model, alongside a 59.18% to 67.64% reduction versus baselines.Memory is measured as the average footprint of uniquely activated MoE parameters across all layers.
- Efficiency: 1.14× to 1.66× end-to-end latency speedup is achieved across four benchmarks, mainly through reduced MoE latency.The reported end-to-end latency includes MoE latency.
- Extreme Compression: 99.11% of original performance is retained at p = 0.6 under compression levels comparable to DES-S and DES-V, while those baselines degrade substantially.At p = 0.5, dMoE retains 97.50% of the original performance.
- Ablations: 76.15%, 72.97%, 68.46%, and 53.97% reductions in average unique expert count occur at block sizes 32, 24, 16, and 8, respectively.The ablation reports consistently strong performance across these block sizes.
6 Conclusion · Appendix
The paper introduces dMoE for block-level expert routing in MoE dLLMs, enabling aggressive expert compression with almost no performance degradation. The appendices cover limitations, future work, impact, licensing, and computing resources.
- 6 Conclusion: dMoE introduces block-level expert routing for MoE dLLMs.The method is presented as a novel strategy for routing experts at the block level.
- 6 Conclusion: dMoE enables aggressive expert compression with almost no performance degradation.This is the paper’s central efficiency claim.
- 6 Conclusion: Comprehensive evaluations across diverse benchmarks demonstrate dMoE’s effectiveness.The conclusion summarizes results across multiple benchmark settings.
- 6 Conclusion: The work lays a foundation for block-level routing in MoE dLLMs.The authors position dMoE as groundwork for this routing paradigm.
- 6 Conclusion: The paper identifies block-level routing as a promising direction for improving MoE dLLM efficiency.This frames the broader significance of the proposed approach.
- Appendix: Appendix A discusses limitations and future work.The appendix explicitly addresses both constraints and subsequent research directions.
- Appendix: Appendix B provides the impact statement, while Appendix C provides the license statement.These appendices document the paper’s impact and licensing information.
- Appendix: Appendix D reports the computing resources used.The computing-resource disclosure is included as a separate appendix.
A Limitations & Future Work … A Limitations & Future Work
The method is presented as applicable beyond language-only MoE dLLMs, with natural extensions to image and video modalities and related multimodal tasks. The authors identify block-level expert routing as a promising direction for future MoE dLLM research.
- A Limitations & Future Work: The method is, in principle, applicable to any MoE dLLM.
- A Limitations & Future Work: The evaluation in this work focuses on the language modality.
- A Limitations & Future Work: The same idea can naturally extend to image and video modalities.
- A Limitations & Future Work: Potential downstream applications include visual question answering.
- A Limitations & Future Work: Potential downstream applications also include visual reasoning.
- A Limitations & Future Work: The results suggest that block-level expert routing is a promising direction for future MoE dLLMs.
B Impact Statement
This work improves inference efficiency for MoE dLLMs through block-level expert routing, reducing uniquely activated experts and alleviating memory bottlenecks. The resulting reductions in memory overhead and inference cost support more practical deployment in resource- and latency-sensitive settings.
- Impact Statement: Block-level expert routing reduces the number of uniquely activated experts during MoE dLLM decoding.This targets the memory bottleneck in inference.
- Impact Statement: The method reduces memory overhead, lowers inference cost, and alleviates memory bottlenecks in MoE dLLMs.These improvements focus on making inference more efficient.
- Impact Statement: These efficiency improvements make large-scale model deployment more practical in resource-constrained, memory-sensitive, or latency-sensitive settings.The impact is especially relevant where memory and latency limit deployment feasibility.
C License Statement
The model and datasets are publicly available, and experiments comply with their respective licenses. The paper lists the applicable licenses and follows official model guidance for fair, reproducible inference evaluation.
- License compliance: The model and datasets are publicly available, and all experiments comply with their respective licenses.The paper states that specific model and dataset licenses are listed below.
- Model license: LLaDA2.0-mini5 is released under the Apache 2.0 License.
- Dataset licenses: GSM8K, MATH500, and MMLU use the MIT License, while ARC-C uses the cc-by-sa-4.0 License.
- Evaluation practice: Inference procedures follow each model’s official documentation and recommended practices to support fair and reproducible evaluation.
D Computing Resources
The experiments used a high-memory server with four NVIDIA H100 GPUs and data parallelism, completing evaluation within 12 hours and training within 144 hours.
- Hardware and execution: The experiments ran on two AMD EPYC 9654 96-Core processors, 1.5 TiB system memory, and four NVIDIA H100 GPUs with 80 GB VRAM each.Data parallelism was employed to speed up evaluation.
- Hardware and execution: 12 hours was the maximum duration for each evaluation run, while each training run finished within 144 hours.