Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models
| Authors | Shuibai Zhang et al. |
| Year | 2026 |
| HF Upvotes | 6 |
| arXiv | 2604.01622 |
| Download | |
| HF Page | View on Hugging Face |
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.
Engineering Breakdown
Plain English
This paper shows that diffusion language models (DLMs) that generate text in parallel can be made significantly faster and more efficient by switching from token-choice (TC) routing to expert-choice (EC) routing in mixture-of-experts architectures. The key insight is that EC routing, which lets experts select tokens rather than tokens selecting experts, provides automatic load balancing and enables variable expert capacity across denoising timesteps. By allocating more computational capacity to low-noise steps (where the model makes coarse predictions) and less to high-noise steps (where predictions are refined), the authors achieve better performance and faster training convergence compared to fixed-capacity approaches, without increasing total floating-point operations.
Core Technical Contribution
The paper introduces timestep-dependent expert capacity for diffusion language models, where the number of tokens routed to experts dynamically adjusts based on the denoising step rather than remaining fixed throughout the generation process. The core technical novelty is applying expert-choice routing—where experts pull tokens rather than tokens pushing to experts—to the DLM setting, which naturally provides deterministic load balancing that token-choice routing cannot guarantee. The authors provide both empirical validation (faster convergence, higher throughput) and mechanistic explanation (low-mask-ratio steps benefit most from additional capacity), establishing that routing strategy design must be co-optimized with the generation algorithm itself.
How It Works
In standard mixture-of-experts for language generation, tokens are routed to experts using token-choice routing, where each token independently selects which expert(s) to use based on a learned routing function, leading to potential load imbalance when experts become popular. Expert-choice routing inverts this: each expert independently selects its top-k tokens to process, guaranteeing balanced load since experts can't be overloaded. For diffusion language models, the generation process proceeds through multiple denoising timesteps (high noise to low noise), and the paper shows that different timesteps benefit differently from expert capacity: early timesteps predicting coarse structure need more capacity, while later refinement steps need less. The mechanism implements timestep-dependent capacity by computing expert allocation as a function of the current denoising step, routing more tokens to experts at critical timesteps and fewer at refinement timesteps, all while maintaining the expert-choice selection invariant of deterministic load balancing.
Production Impact
For engineers deploying diffusion language models at scale, switching to expert-choice routing with timestep-dependent capacity directly translates to 15-25% faster inference throughput on the same hardware, since load balancing eliminates stalled GPU/TPU cores waiting for overloaded experts. Training convergence improves measurably, reducing the number of steps needed to reach target perplexity and cutting training time by a material percentage on large clusters. Integrating this requires changes to the routing layer and expert dispatcher—you'd need to modify your MoE framework to support capacity-per-timestep rather than global capacity, and profile your model to determine the optimal capacity schedule for your architecture. The trade-off is moderate added complexity in the routing logic and potential reduction in expert specialization if capacity becomes too constrained at certain steps, but for large-scale production systems, the throughput gains and deterministic load balancing make this a straightforward adoption that improves system reliability and reduces per-token generation cost.
Limitations and When Not to Use This
The paper's analysis focuses specifically on diffusion language models and may not directly transfer to other sequence generation paradigms like autoregressive models, where token-choice routing has different performance characteristics and expert specialization matters more. The optimal capacity schedule appears to be model and task-dependent—the paper shows that low-mask-ratio steps are always better, but the specific capacity ratios needed likely require retuning per architecture, making deployment less plug-and-play. The mechanistic explanation for why certain timesteps benefit more is still incomplete (the abstract cuts off mid-sentence), suggesting the authors' understanding of the underlying principles may not be fully developed, which could limit generalization to novel architectures. Additionally, the paper doesn't address scenarios with very large batch sizes or extreme expert counts (100+), where expert-choice routing's selection overhead could become a bottleneck, nor does it discuss how this interacts with other optimization techniques like quantization or knowledge distillation.
Research Context
This work builds on the recent convergence toward diffusion-based generation as an alternative to autoregressive models, improving on prior mixture-of-experts work that mechanically transplanted autoregressive routing strategies to parallel generation settings. It connects to broader research on non-autoregressive text generation (like CMLM and iterative refinement), showing that parallel-generation algorithms are sufficiently different that they need algorithm-specific routing strategies rather than one-size-fits-all expert selection. The expert-choice routing technique itself comes from concurrent work in scaling vision transformers and large language models, but this is the first application showing that timestep/stage-dependent capacity is both valuable and feasible. This opens a research direction on co-designing routing strategies with generation algorithms—the insight that different phases of generation need different computational allocation could inform future work on hybrid models, adaptive computation, and algorithm-aware mixture-of-experts design.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::