Papers
arxiv:2604.01622

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

Published on Apr 2
· Submitted by
Shuibai Zhang
on Apr 8
Authors:
,
,
,
,
,
,
,
,
,

Abstract

Expert-choice routing improves diffusion language model mixture-of-experts by providing deterministic load balancing and adaptive computation allocation based on denoising steps.

AI-generated summary

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.

Community

Paper author Paper submitter

We show that expert-choice (EC) routing is a better fit for diffusion language models (DLMs) than token-choice routing: it gives deterministic load balancing, 2× faster convergence, and enables timestep-dependent adaptive computation. Pretrained TC models can be retrofitted to EC by replacing only the router. Code: https://github.com/zhangshuibai/EC-DLM

Sign up or log in to comment

Get this paper in your agent:

hf papers read 2604.01622
Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash

Models citing this paper 0

No model linking this paper

Cite arxiv.org/abs/2604.01622 in a model README.md to link it from this page.

Datasets citing this paper 0

No dataset linking this paper

Cite arxiv.org/abs/2604.01622 in a dataset README.md to link it from this page.

Spaces citing this paper 0

No Space linking this paper

Cite arxiv.org/abs/2604.01622 in a Space README.md to link it from this page.

Collections including this paper 0

No Collection including this paper

Add this paper to a collection to link it from this page.