dMoE: dLLMs с обучаемыми блочными экспертами

Аннотация

Недавно появившиеся диффузионные большие языковые модели (dLLM) стали многообещающей альтернативой авторегрессионным моделям, предлагая конкурентоспособную производительность при естественной поддержке параллельного декодирования. Однако по мере интеграции dLLM с архитектурами на основе смеси экспертов (MoE) для масштабирования ёмкости модели возникает фундаментальное несоответствие между блочным параллельным декодированием и выбором экспертов на уровне токенов. В частности, каждый прямой проход dLLM обрабатывает несколько токенов с двунаправленными зависимостями, тогда как стандартные слои MoE маршрутизируют каждый токен независимо. Это несоответствие существенно увеличивает количество уникально активированных экспертов, делая инференс всё более ограниченным по памяти. Для решения этой проблемы мы предлагаем dMoE — простую, но эффективную блочную структуру MoE. Основная идея dMoE заключается в агрегации распределений экспертов на уровне токенов внутри каждого блока в единое блочное распределение экспертов, которое затем используется для более согласованной маршрутизации. Таким образом, dMoE значительно сокращает количество уникально активированных экспертов во время инференса без потери производительности, снижая тем самым узкое место по памяти. Обширные эксперименты на различных тестовых наборах демонстрируют эффективность dMoE. В среднем dMoE сокращает количество уникально активированных экспертов с 69,5 до 14,6, сохраняя при этом 99,11% исходной производительности. Одновременно это снижает использование памяти на 76,64%–79,84% и обеспечивает ускорение сквозной задержки в 1,14–1,66 раза. Код доступен по адресу: https://github.com/fscdc/dMoE

English

Diffusion Large Language Models (dLLMs) have recently emerged as a promising alternative to autoregressive models, offering competitive performance while naturally supporting parallel decoding. However, as dLLMs are increasingly integrated with Mixture-of-Experts (MoE) architectures to scale model capacity, a fundamental mismatch arises between block parallel decoding and token-level expert selection. Specifically, each dLLM forward pass processes multiple tokens with bidirectional dependencies, whereas conventional MoE layers route each token independently. This mismatch substantially increases the number of uniquely activated experts, making inference increasingly memory-bound. To address this, we propose dMoE, a simple yet effective block-level MoE framework. The central idea of dMoE is to aggregate token-level expert distributions within each block into a unified block-level expert distribution, which is then used to guide expert routing in a more coherent manner. In this way, dMoE substantially reduces the number of uniquely activated experts during inference without sacrificing performance, thereby mitigating the memory-bound bottleneck. Extensive experiments across a variety of benchmarks demonstrate the effectiveness of dMoE. On average, dMoE reduces the number of uniquely activated experts from 69.5 to 14.6 while retaining 99.11% of the original performance. Meanwhile, it reduces memory usage by 76.64% to 79.84% and achieves 1.14times to 1.66times end-to-end latency speedup. Code is available at: https://github.com/fscdc/dMoE