Экспертное выборочное маршрутизация обеспечивает адаптивные вычисления в диффузионных языковых моделях

Expert-Choice Routing Enables Adaptive Computation in Diffusion Language Models

April 2, 2026

Авторы: 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

Аннотация

Диффузионные языковые модели (DZM) обеспечивают параллельную неавторегрессионную генерацию текста, однако существующие DZM с механизмом смеси экспертов (MoE) наследуют маршрутизацию на основе выбора токенов (token-choice, TC) из авторегрессионных систем, что приводит к дисбалансу нагрузки и жесткому распределению вычислений. Мы показываем, что маршрутизация на основе выбора экспертом (expert-choice, EC) лучше подходит для DZM: она обеспечивает детерминированное балансирование нагрузки по своей конструкции, обеспечивая более высокую пропускную способность и ускоренную сходимость по сравнению с TC. Опираясь на свойство внешней управляемости пропускной способности EC, мы вводим зависящую от временного шага пропускную способность экспертов, которая варьирует распределение экспертов в соответствии с шагом денойзинга. Мы обнаруживаем, что выделение большей пропускной способности на шаги с низким уровнем маскирования последовательно обеспечивает наилучшую производительность при сопоставимых FLOPs, и даем механистическое объяснение: токены в контекстах с низким уровнем маскирования демонстрируют на порядок более высокую эффективность обучения, поэтому концентрация вычислений на этих шагах дает наибольшую предельную отдачу. Наконец, мы показываем, что существующие предобученные DZM с TC можно адаптировать к EC, заменив только маршрутизатор, добиваясь ускоренной сходимости и повышения точности в разнообразных downstream-задачах. В совокупности эти результаты утверждают EC-маршрутизацию как превосходную парадигму для MoE-моделей DZM и демонстрируют, что вычисления в DZM можно рассматривать как адаптивную политику, а не как фиксированную архитектурную константу. Код доступен по адресу 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.