L'Instradamento a Scelta Esperta Abilita il Calcolo Adattivo nei Modelli Linguistici di Diffusione

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

April 2, 2026

Autori: 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

Abstract

I modelli linguistici di diffusione (DLM) abilitano la generazione di testo parallela e non autoregressiva, tuttavia gli esistenti modelli DLM con miscela di esperti (MoE) ereditano l'instradamento basato sulla scelta del token (TC) dai sistemi autoregressivi, causando squilibrio di carico e allocazione rigida del calcolo. Dimostriamo che l'instradamento a scelta dell'esperto (EC) è più adatto ai DLM: fornisce per progettazione un bilanciamento deterministico del carico, producendo una maggiore velocità di elaborazione e una convergenza più rapida rispetto al TC. Basandoci sulla proprietà per cui la capacità EC è esternamente controllabile, introduciamo una capacità esperto dipendente dal passo temporale, che varia l'allocazione degli esperti in base allo step di denoising. Rileviamo che allocare maggiore capacità ai passi con basso rapporto di mascheramento raggiunge costantemente le migliori prestazioni a parità di FLOP, e forniamo una spiegazione meccanicistica: i token in contesti con basso rapporto di mascheramento mostrano un'efficienza di apprendimento di un ordine di grandezza superiore, quindi concentrare il calcolo su questi passi produce il rendimento marginale maggiore. Infine, dimostriamo che i DLM TC preaddestrati esistenti possono essere adattati all'EC sostituendo solo il router, raggiungendo una convergenza più rapida e una precisione migliorata in varie attività downstream. Complessivamente, questi risultati stabiliscono l'instradamento EC come paradigma superiore per i modelli DLM MoE e dimostrano che il calcolo nei DLM può essere trattato come una politica adattiva piuttosto che una costante architetturale fissa. Il codice è disponibile all'indirizzo 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.