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 this https URL.
Abstract (translated)
扩散语言模型（DLMs）支持并行、非自回归的文本生成，但现有的DLM混合专家（MoE）模型沿用了自回归系统中的令牌选择（TC）路由机制，导致负载不均衡和计算分配僵化。我们证明专家选择（EC）路由更适用于DLMs：其设计上即提供确定性负载均衡，相比TC能实现更高吞吐量和更快的收敛。基于EC容量可外部调控的特性，我们引入了随时间步变化的专家容量，使专家分配能根据去噪步骤动态调整。研究发现，在相同浮点运算量（FLOPs）下，为低掩码比步骤分配更多容量 consistently 能实现最佳性能，并提供机理解释：低掩码比上下文中的令牌学习效率高出一个数量级，因此将计算集中于这些步骤能获得最大边际收益。最后，我们证明现有预训练的TC DLMs仅通过替换路由器即可改造为EC路由，在多种下游任务中实现更快收敛和更高精度。综上，这些结果确立了EC路由作为DLM MoE模型的优越范式，并表明DLM中的计算可被视为自适应策略而非固定架构常量。代码已公开于此https链接。
URL
https://arxiv.org/abs/2604.01622
