mlsanityData to 7 Oct 2026

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

arXiv 2604.01622 ↗First posted 2 Apr 2026cs.LG cs.CLcited by 3

On the map: Diffusion models → Discrete and text diffusion

also in 1 more
Abstract (arXiv)

Diffusion language models (DLMs) enable parallel, non-autoregressive text generation, yet existing DLM mixture-of-experts (MoE) models inherit token-choice (TC) routing from autoregressive systems, leading to load imbalance and rigid computation allocation. We show that expert-choice (EC) routing is a better fit for DLMs: it provides deterministic load balancing by design, yielding higher throughput and faster convergence than TC. Building on the property that EC capacity is externally controllable, we introduce timestep-dependent expert capacity, which varies expert allocation according to the denoising step. We find that allocating more capacity to low-mask-ratio steps consistently achieves the best performance under matched FLOPs, and provide a mechanistic explanation: tokens in low-mask-ratio contexts exhibit an order-of-magnitude higher learning efficiency, so concentrating compute on these steps yields the largest marginal return. Finally, we show that existing pretrained TC DLMs can be retrofitted to EC by replacing only the router, achieving faster convergence and improved accuracy across diverse downstream tasks. Together, these results establish EC routing as a superior paradigm for DLM MoE models and demonstrate that computation in DLMs can be treated as an adaptive policy rather than a fixed architectural constant. Code is available at https://github.com/zhangshuibai/EC-DLM.

Key facts. Cited by 3 papers in the map (data through 7 Oct 2026). Placed in Discrete and text diffusion, one of 14 directions in Diffusion models. Builds on 55 of its 60 references that are in the map. First posted on arXiv 2 Apr 2026.

Builds on

55 of 60 references.

Closest ancestors · 11

Same direction, newest first

  1. 2025WeDLM: Reconciling Diffusion Language Models with Standard Causal Attention for Fast Inference28 Dec
  2. 2025Fine-Tuning Masked Diffusion for Provable Self-Correction1 Oct
  3. 2025Planner Aware Path Learning in Diffusion Language Models Training27 Sep
  4. 2025Dream 7B: Diffusion Large Language Models21 Aug
  5. 2025Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models12 Mar
  6. 2025Large Language Diffusion Models14 Feb
  7. 2025Path Planning for Masked Diffusion Model Sampling5 Feb
  8. 2024Scaling Diffusion Language Models via Adaptation from Autoregressive Models23 Oct
  9. 2024Simple and Effective Masked Diffusion Language Models11 Jun
  10. 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution25 Oct
  11. 2022DiffusionBERT: Improving Generative Masked Language Models with Diffusion Models28 Nov

Other references · 44

  1. 2021Training Verifiers to Solve Math Word Problemscited by 9,743
  2. 2021Evaluating Large Language Models Trained on Codecited by 8,053
  3. 2020Measuring Massive Multitask Language Understandingcited by 7,768
  4. 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
  5. 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,413
  6. 2021Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsitycited by 3,014
  7. 2019ALBERT: A Lite BERT for Self-supervised Learning of Language Representationscited by 2,772
  8. 2021Structured Denoising Diffusion Models in Discrete State-Spacescited by 1,665
  9. 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601
  10. 2023Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generationcited by 1,332
  11. 2020What Disease does this Patient Have? A Large-scale Open Domain Question Answering Dataset from Medical Examscited by 1,324
  12. 2024DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Modelscited by 834
32 more
  1. 2022Mixture-of-Experts with Expert Choice Routingcited by 611
  2. 2025Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decodingcited by 404
  3. 2025Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approachcited by 372
  4. 2021BASE Layers: Simplifying Training of Large, Sparse Modelscited by 315
  5. 2023Common Diffusion Noise Schedules and Sample Steps are Flawedcited by 308
  6. 2024LayerSkip: Enabling Early Exit Inference and Self-Speculative Decodingcited by 302
  7. 2022Confident Adaptive Language Modelingcited by 297
  8. 2024OLMoE: Open Mixture-of-Experts Language Modelscited by 291
  9. 2023Efficient Diffusion Training via Min-SNR Weighting Strategycited by 242
  10. 2024Mixture-of-Depths: Dynamically allocating compute in transformer-based language modelscited by 242
  11. 2024From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Stepcited by 225
  12. 2025Mercury: Ultra-Fast Language Models Based on Diffusioncited by 225
  13. 2024Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Expertscited by 191
  14. 2024Nemotron-CC: Transforming Common Crawl into a Refined Long-Horizon Pretraining Datasetcited by 153
  15. 2025Fast-dLLM v2: Efficient Block-Diffusion LLMcited by 128
  16. 2025Mixture-of-Recursions: Learning Dynamic Recursive Depths for Adaptive Token-Level Computationcited by 120
  17. 2025Accelerating Diffusion LLMs via Adaptive Parallel Decodingcited by 108
  18. 2022Block-Recurrent Transformerscited by 107
  19. 2024Looped Transformers for Length Generalizationcited by 91
  20. 2024Relaxed Recursive Transformers: Effective Parameter Sharing with Layer-wise LoRAcited by 72
  21. 2024Dynamic Diffusion Transformercited by 60
  22. 2025dParallel: Learnable Parallel Decoding for dLLMscited by 59
  23. 2025OpenCodeInstruct: A Large-scale Instruction Tuning Dataset for Code LLMscited by 57
  24. 2024ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routingcited by 53
  25. 2024A Closer Look at Time Steps is Worthy of Triple Speed-Up for Diffusion Model Trainingcited by 30
  26. 2024EC-DIT: Scaling Diffusion Transformers with Adaptive Expert-Choice Routingcited by 29
  27. 2024Denoising Task Difficulty-based Curriculum for Training Diffusion Modelscited by 26
  28. 2025DiffMoE: Dynamic Token Selection for Scalable Diffusion Transformerscited by 23
  29. 2025Free Draft-and-Verification: Toward Lossless Parallel Decoding for Diffusion Large Language Modelscited by 19
  30. 2026Scalable Training of Mixture-of-Experts Models with Megatron Corecited by 16
  31. 2023AdaDiff: Accelerating Diffusion Models through Step-Wise Adaptive Computationcited by 14
  32. 2026Expert Threshold Routing for Autoregressive Language Modeling with Dynamic Computation Allocation and Load Balancingcited by 3

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2026: 3.

3 papers.

  1. 2026dMoE: dLLMs with Learnable Block Expertscited by 2
  2. 2026Fixed-Point Masked Generative Modelingcited by 2
  3. 2026LLaDA MoE v2: Scaling Mixture-of-Experts Diffusion Language Modelscited by 1

Similar, not cited

Close in wording, no citation either way

  1. 2026REFLEX: Rethinking MoE Inference as Refinement-Aware Compute Allocation in Diffusion Language Models
  2. 2026Dynamic Expert Sharing: Decoupling Memory from Parallelism in Mixture-of-Experts Diffusion LLMs
  3. 2026MMOE: Modernizing Diffusion Transformers with Efficient Expert Design
  4. 2026Adaptation to Intrinsic Dependence in Diffusion Language Models
  5. 2026TEAM: Temporal-Spatial Consistency Guided Expert Activation for MoE Diffusion Language Model Acceleration
  6. 2026ALoDLM: Adaptively Looped Diffusion Language Models
  7. 2025On Powerful Ways to Generate: Autoregression, Diffusion, and Beyond
  8. 2026From Position Risks to Block Survival: Faster Generation for Diffusion Language Models