dMoE: dLLMs con Bloques de Expertos Aprendibles

Resumen

Los Modelos de Lenguaje de Gran Escala de Difusión (dLLMs) han surgido recientemente como una alternativa prometedora a los modelos autorregresivos, ofreciendo un rendimiento competitivo y soportando de forma natural la decodificación paralela. Sin embargo, a medida que los dLLMs se integran cada vez más con arquitecturas de Mezcla de Expertos (MoE) para escalar la capacidad del modelo, surge un desajuste fundamental entre la decodificación paralela por bloques y la selección de expertos a nivel de token. Específicamente, cada pase hacia adelante de un dLLM procesa múltiples tokens con dependencias bidireccionales, mientras que las capas MoE convencionales enrutan cada token de forma independiente. Este desajuste incrementa sustancialmente el número de expertos activados de forma única, haciendo que la inferencia esté cada vez más limitada por la memoria. Para abordar esto, proponemos dMoE, un marco de Mezcla de Expertos a nivel de bloque simple pero efectivo. La idea central de dMoE es agregar las distribuciones de expertos a nivel de token dentro de cada bloque en una distribución unificada a nivel de bloque, que luego se utiliza para guiar el enrutamiento de expertos de una manera más coherente. De esta forma, dMoE reduce sustancialmente el número de expertos activados de forma única durante la inferencia sin sacrificar el rendimiento, mitigando así el cuello de botella limitado por la memoria. Experimentos exhaustivos en una variedad de puntos de referencia demuestran la efectividad de dMoE. En promedio, dMoE reduce el número de expertos activados de forma única de 69,5 a 14,6, manteniendo el 99,11% del rendimiento original. Al mismo tiempo, reduce el uso de memoria entre un 76,64% y un 79,84%, y logra una aceleración de latencia de extremo a extremo de 1,14 a 1,66 veces. El código está disponible en: 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