dMoE: dLLMs mit lernbaren Block-Experten

Zusammenfassung

Diffusion Large Language Models (dLLMs) haben sich in letzter Zeit als vielversprechende Alternative zu autoregressiven Modellen erwiesen, da sie eine wettbewerbsfähige Leistung bieten und gleichzeitig auf natürliche Weise parallele Dekodierung unterstützen. Mit der zunehmenden Integration von dLLMs in Mixture-of-Experts-Architekturen (MoE) zur Skalierung der Modellkapazität tritt jedoch ein grundlegender Konflikt zwischen block-paralleler Dekodierung und tokenweiser Expertenauswahl auf. Konkret verarbeitet jeder Vorwärtsdurchlauf eines dLLM mehrere Tokens mit bidirektionalen Abhängigkeiten, während konventionelle MoE-Schichten jedes Token unabhängig routen. Diese Diskrepanz erhöht die Anzahl der einzigartig aktivierten Experten erheblich, wodurch die Inferenz zunehmend speichergebunden wird. Um dieses Problem zu adressieren, schlagen wir dMoE vor, ein einfaches, aber effektives MoE-Framework auf Blockebene. Die zentrale Idee von dMoE besteht darin, Token-Ebene-Expertenverteilungen innerhalb jedes Blocks zu einer einheitlichen Expertenverteilung auf Blockebene zu aggregieren, die dann verwendet wird, um das Expertenrouting kohärenter zu steuern. Auf diese Weise reduziert dMoE die Anzahl der während der Inferenz einzigartig aktivierten Experten erheblich, ohne die Leistung zu beeinträchtigen, und mildert somit den speichergebundenen Engpass. Umfangreiche Experimente über verschiedene Benchmarks hinweg belegen die Wirksamkeit von dMoE. Im Durchschnitt reduziert dMoE die Anzahl der einzigartig aktivierten Experten von 69,5 auf 14,6, während 99,11 % der ursprünglichen Leistung erhalten bleiben. Gleichzeitig wird der Speicherverbrauch um 76,64 % bis 79,84 % reduziert und eine End-to-End-Latenzbeschleunigung um den Faktor 1,14 bis 1,66 erreicht. Der Code ist verfügbar unter: 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