dMoE : dLLMs avec des experts de blocs apprenables

Résumé

Les Modèles de Langage de Grande Taille à Diffusion (dLLMs) sont récemment apparus comme une alternative prometteuse aux modèles autorégressifs, offrant des performances compétitives tout en supportant naturellement le décodage parallèle. Cependant, à mesure que les dLLMs sont intégrés à des architectures Mixture d'Experts (MoE) pour augmenter la capacité des modèles, une inadéquation fondamentale émerge entre le décodage parallèle par blocs et la sélection d'experts au niveau des tokens. Plus précisément, chaque passe avant d'un dLLM traite plusieurs tokens avec des dépendances bidirectionnelles, tandis que les couches MoE conventionnelles acheminent chaque token indépendamment. Cette inadéquation augmente considérablement le nombre d'experts activés de manière unique, rendant l'inférence de plus en plus limitée par la mémoire. Pour y remédier, nous proposons dMoE, un cadre MoE simple mais efficace au niveau des blocs. L'idée centrale de dMoE est d'agréger les distributions d'experts au niveau des tokens au sein de chaque bloc en une distribution d'experts unifiée au niveau du bloc, qui est ensuite utilisée pour guider le routage des experts de manière plus cohérente. Ainsi, dMoE réduit considérablement le nombre d'experts activés de manière unique lors de l'inférence sans sacrifier les performances, atténuant ainsi le goulot d'étranglement lié à la mémoire. Des expériences approfondies sur divers benchmarks démontrent l'efficacité de dMoE. En moyenne, dMoE réduit le nombre d'experts activés de manière unique de 69,5 à 14,6 tout en conservant 99,11 % de la performance originale. Parallèlement, il réduit l'utilisation mémoire de 76,64 % à 79,84 % et atteint une accélération de la latence de bout en bout de 1,14x à 1,66x. Le code est disponible à l'adresse : 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