Expert-Choice Routing Maakt Adaptieve Verwerking Mogelijk in Diffusion Taalmodellen
Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models
April 2, 2026
Auteurs: Shuibai Zhang, Caspian Zhuang, Chihan Cui, Zhihan Yang, Fred Zhangzhi Peng, Yanxin Zhang, Haoyue Bai, Zack Jia, Yang Zhou, Guanhua Chen, Ming Liu
cs.AI
Samenvatting
Diffusie-taalmmodellen (DLM's) maken parallelle, niet-autoregressieve tekstgeneratie mogelijk, maar bestaande DLM mixture-of-experts (MoE) modellen erven token-choice (TC) routering van autoregressieve systemen, wat leidt tot belastingsonbalans en rigide compute-toewijzing. Wij tonen aan dat expert-choice (EC) routering beter geschikt is voor DLM's: het biedt door ontwerp deterministische belastingsbalancering, wat een hogere doorvoer en snellere convergentie oplevert dan TC. Voortbouwend op de eigenschap dat EC-capaciteit extern aanstuurbaar is, introduceren we tijdstap-afhankelijke expertcapaciteit, die de experttoewijzing varieert volgens de denoisingsstap. Wij constateren dat het toewijzen van meer capaciteit aan stappen met een lage maskerratio consequent de beste prestaties bereikt bij gelijke FLOPs, en geven een mechanistische verklaring: tokens in contexten met een lage maskerratio vertonen een leerefficiëntie die een orde van grootte hoger ligt, waardoor het concentreren van compute op deze stappen het grootste marginale rendement oplevert. Ten slotte tonen we aan dat bestaande voorgetrainde TC DLM's kunnen worden omgebouwd naar EC door alleen de router te vervangen, wat snellere convergentie en verbeterde nauwkeurigheid oplevert bij diverse downstreamtaken. Samen vestigen deze resultaten EC-routering als een superieur paradigma voor DLM MoE-modellen en demonstreren ze dat compute in DLM's kan worden behandeld als een adaptief beleid in plaats van een vaste architecturale constante. Code is beschikbaar op https://github.com/zhangshuibai/EC-DLM.
English
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.