Le Routage par Choix d'Expert Permet un Calcul Adaptatif dans les Modèles de Langage par Diffusion
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
Résumé
Les modèles de langage par diffusion (DLM) permettent une génération de texte parallèle et non autorégressive. Cependant, les modèles de type mixture d'experts (MoE) existants pour les DLM héritent du routage par choix de token (TC) des systèmes autorégressifs, ce qui entraîne un déséquilibre de charge et une allocation de calcul rigide. Nous démontrons que le routage par choix d'expert (EC) est mieux adapté aux DLM : il offre par conception un équilibrage de charge déterministe, procurant un débit supérieur et une convergence plus rapide que le TC. En nous appuyant sur la propriété selon laquelle la capacité des experts dans l'EC est contrôlable extérieurement, nous introduisons une capacité d'expert dépendante du pas de temps, qui fait varier l'allocation des experts en fonction de l'étape de débruitage. Nous constatons qu'allouer plus de capacité aux étapes où le ratio de masquage est faible permet systématiquement d'atteindre les meilleures performances à complexité de calcul (FLOPs) équivalente, et nous en proposons une explication mécaniste : les tokens dans des contextes à faible ratio de masquage présentent une efficacité d'apprentissage supérieure d'un ordre de grandeur, donc concentrer le calcul sur ces étapes génère le rendement marginal le plus important. Enfin, nous montrons que les DLM-TC pré-entraînés existants peuvent être adaptés à l'EC en ne remplaçant que le routeur, obtenant ainsi une convergence accélérée et une précision améliorée sur diverses tâches en aval. Ensemble, ces résultats établissent le routage EC comme un paradigme supérieur pour les modèles DLM MoE et démontrent que le calcul dans les DLM peut être traité comme une politique adaptive plutôt que comme une constante architecturale fixe. Le code est disponible à l'adresse 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.