Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

2026-04-02

AI Summary

This work addresses the limitations of existing Mixture-of-Experts (MoE) systems in diffusion language models, which inherit token-based routing from autoregressive models, leading to imbalanced expert utilization and inflexible computation allocation. We introduce Expert Choice (EC) routing—adapted for the first time to diffusion language models—along with a timestep-dependent dynamic expert capacity strategy that enables balanced load distribution and adaptive computation. Our analysis reveals that training stages with lower masking ratios exhibit higher learning efficiency, allowing strategic reallocation of computational resources to significantly improve training returns. Experiments demonstrate that EC routing substantially accelerates convergence and enhances final performance over conventional token routing under identical FLOPs budgets. Notably, these gains are achieved by simply replacing the router in a pretrained model, offering a plug-and-play improvement with minimal architectural modification.

Abstract

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.

Problem

• Diffusion Language Models
• Mixture-of-Experts
• Routing
• Load Imbalance
• Adaptive Computation

Innovation

• Expert-Choice Routing
• Diffusion Language Models
• Adaptive Computation
• Mixture-of-Experts
• Timestep-Dependent Capacity