arXiv:2605.30876cs.CL2026-05被引 4

解决扩散语言模型中专家路由的内存瓶颈问题

dMoE: dLLMs with Learnable Block Experts

论文配图:dMoE: dLLMs with Learnable Block Experts
图 1 · 摘自论文原文
  • 按块聚合专家分布,统一路由决策
  • 激活专家数从69.5降至14.6,性能损失仅0.89%
  • 适合大规模扩散模型推理优化场景

扩散语言模型(dLLMs)作为自回归模型的替代方案,具备并行解码优势。但将其与专家混合(MoE)架构结合时,块级并行解码与逐令牌专家选择之间存在根本性矛盾:每个前向传播处理多个具有双向依赖的令牌,而传统MoE独立路由每个令牌,导致唯一激活专家数量激增,使推理严重受限于内存。为此,我们提出dMoE,一种简单的块级MoE框架。其核心思想是将每个块内令牌级专家分布聚合为统一的块级分布,以更连贯地指导专家路由。该方法显著减少推理过程中唯一激活的专家数量,同时保持性能。在多个基准测试中,dMoE将唯一激活专家数从69.5降至14.6,保留99.11%原始性能,内存使用降低76.64%至79.84%,端到端延迟提速1.14×至1.66×。代码已开源。

原文摘要 · Abstract (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.14$\times$ to 1.66$\times$ end-to-end latency speedup. Code is available at: https://github.com/fscdc/dMoE

扩散模型MoE推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。