O Roteamento de Escolha do Especialista Permite Computação Adaptativa em Modelos de Linguagem de Difusão
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
Resumo
Os modelos de linguagem de difusão (DLMs) permitem geração de texto paralela e não autorregressiva, porém os modelos existentes de mistura de especialistas (MoE) para DLMs herdam o roteamento por escolha de token (TC) de sistemas autorregressivos, levando a desequilíbrio de carga e alocação computacional rígida. Demonstramos que o roteamento por escolha de especialista (EC) é mais adequado para DLMs: ele fornece balanceamento de carga determinístico por projeto, resultando em maior taxa de transferência e convergência mais rápida que o TC. Com base na propriedade de que a capacidade do EC é controlável externamente, introduzimos capacidade de especialista dependente do passo temporal, que varia a alocação de especialistas de acordo com a etapa de remoção de ruído. Descobrimos que alocar mais capacidade para etapas com baixa taxa de mascaramento consistentemente alcança o melhor desempenho sob FLOPs equivalentes, e fornecemos uma explicação mecanicista: tokens em contextos de baixa taxa de mascaramento exibem uma eficiência de aprendizado uma ordem de grandeza maior, portanto concentrar computação nessas etapas produz o maior retorno marginal. Finalmente, mostramos que DLMs TC pré-treinados existentes podem ser adaptados para EC substituindo apenas o roteador, alcançando convergência mais rápida e precisão melhorada em diversas tarefas downstream. Juntos, esses resultados estabelecem o roteamento EC como um paradigma superior para modelos MoE de DLM e demonstram que a computação em DLMs pode ser tratada como uma política adaptativa em vez de uma constante arquitetural fixa. O código está disponível em 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.