dMoE: dLLMs com Especialistas de Bloco Aprendíveis

Resumo

Modelos de Linguagem de Grande Escala por Difusão (dLLMs) surgiram recentemente como uma alternativa promissora aos modelos autorregressivos, oferecendo desempenho competitivo enquanto suportam naturalmente a decodificação paralela. No entanto, à medida que os dLLMs são cada vez mais integrados a arquiteturas de Mistura de Especialistas (MoE) para escalar a capacidade do modelo, surge uma incompatibilidade fundamental entre a decodificação paralela por bloco e a seleção de especialistas em nível de token. Especificamente, cada passagem direta do dLLM processa múltiplos tokens com dependências bidirecionais, enquanto camadas MoE convencionais roteiam cada token de forma independente. Essa incompatibilidade aumenta substancialmente o número de especialistas ativados de forma única, tornando a inferência cada vez mais limitada por memória. Para resolver isso, propomos o dMoE, um framework MoE em nível de bloco simples, porém eficaz. A ideia central do dMoE é agregar as distribuições de especialistas em nível de token dentro de cada bloco em uma distribuição unificada de especialistas em nível de bloco, que é então usada para guiar o roteamento de especialistas de maneira mais coerente. Dessa forma, o dMoE reduz substancialmente o número de especialistas ativados de forma única durante a inferência sem sacrificar o desempenho, mitigando assim o gargalo limitado por memória. Extensos experimentos em uma variedade de benchmarks demonstram a eficácia do dMoE. Em média, o dMoE reduz o número de especialistas ativados de forma única de 69,5 para 14,6, mantendo 99,11% do desempenho original. Enquanto isso, reduz o uso de memória em 76,64% a 79,84% e atinge uma aceleração de latência ponta a ponta de 1,14× a 1,66×. O código está disponível em: 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