El Enrutamiento de Elección Experta Permite la Computación Adaptativa en Modelos de Lenguaje de Difusión
Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models
April 2, 2026
Autores: 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
Resumen
Los modelos de lenguaje de difusión (DLM) permiten la generación de texto paralela y no autorregresiva; sin embargo, los modelos existentes de mezcla de expertos (MoE) para DLM heredan el enrutamiento por elección de token (TC) de los sistemas autorregresivos, lo que provoca desequilibrio de carga y una asignación computacional rígida. Demostramos que el enrutamiento por elección de experto (EC) es más adecuado para los DLM: ofrece un equilibrio de carga determinista por diseño, logrando un mayor rendimiento y una convergencia más rápida que TC. Basándonos en la propiedad de que la capacidad de EC es externamente controlable, introducimos una capacidad de experto dependiente del paso temporal, que varía la asignación de expertos según el paso de eliminación de ruido. Descubrimos que asignar más capacidad a los pasos con baja proporción de enmascaramiento logra sistemáticamente el mejor rendimiento con FLOPs equivalentes, y ofrecemos una explicación mecanicista: los tokens en contextos con baja proporción de enmascaramiento exhiben una eficiencia de aprendizaje un orden de magnitud mayor, por lo que concentrar el cómputo en estos pasos produce el mayor retorno marginal. Finalmente, demostramos que los DLM con TC preentrenados existentes pueden adaptarse a EC reemplazando únicamente el enrutador, logrando una convergencia más rápida y una precisión mejorada en diversas tareas posteriores. En conjunto, estos resultados establecen el enrutamiento EC como un paradigma superior para los modelos MoE de DLM y demuestran que el cómputo en los DLM puede tratarse como una política adaptativa en lugar de una constante arquitectónica fija. El código está disponible en 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.