dMoE: 学習可能なブロックエキスパートを備えたdLLM

要旨

拡散大規模言語モデル（dLLMs）は、自己回帰モデルに代わる有望な選択肢として最近登場し、競争力のある性能を提供すると同時に、自然に並列デコードをサポートしています。しかし、dLLMsがモデル容量を拡大するために専門家混合（MoE）アーキテクチャとますます統合されるにつれて、ブロック並列デコードとトークンレベルの専門家選択の間に根本的な不一致が生じます。具体的には、各dLLMの順伝播は双方向依存関係を持つ複数のトークンを処理する一方、従来のMoE層は各トークンを独立にルーティングします。この不一致により、一意に活性化される専門家数が大幅に増加し、推論がますますメモリ律速となります。これに対処するため、我々はシンプルでありながら効果的なブロックレベルのMoEフレームワークであるdMoEを提案します。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