dMoE: dLLMs with Learnable Block Experts

📄 arXiv: 2605.30876v1 📥 PDF

作者: Sicheng Feng, Zigeng Chen, Gongfan Fang, Xinyin Ma, Xinchao Wang

分类: cs.CL

发布日期: 2026-05-29

备注: Working in progress. Code is available at: \url{https://github.com/fscdc/dMoE}

🔗 代码/项目: GITHUB


💡 一句话要点

dMoE:提出可学习块专家机制,解决扩散语言模型中专家选择与块并行解码的失配问题。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 扩散语言模型 混合专家模型 并行解码 块级专家选择 内存优化 推理加速 模型压缩

📋 核心要点

  1. 扩散语言模型结合MoE扩展模型容量时,token级专家选择与块并行解码不匹配,导致大量专家被激活,推理受限于内存。
  2. dMoE通过将块内token级专家分布聚合为块级分布,指导专家路由,从而减少激活专家数量,缓解内存瓶颈。
  3. 实验表明,dMoE在保持性能的同时,显著减少了激活专家数量,降低内存使用,并加速推理过程。

📝 摘要(中文)

扩散大型语言模型(dLLMs)最近作为自回归模型的有希望的替代方案出现,在提供有竞争力的性能的同时,自然地支持并行解码。然而,随着dLLMs越来越多地与混合专家(MoE)架构集成以扩展模型容量,块并行解码和token级专家选择之间出现了根本的不匹配。具体来说,每个dLLM前向传播处理具有双向依赖关系的多个token,而传统的MoE层独立地路由每个token。这种不匹配大大增加了唯一激活专家的数量,使得推理越来越受内存限制。为了解决这个问题,我们提出了dMoE,一个简单而有效的块级MoE框架。dMoE的核心思想是将每个块内的token级专家分布聚合为统一的块级专家分布,然后使用该分布以更连贯的方式指导专家路由。通过这种方式,dMoE在不牺牲性能的情况下,大大减少了推理过程中唯一激活专家的数量,从而缓解了内存限制的瓶颈。在各种基准上的大量实验证明了dMoE的有效性。平均而言,dMoE将唯一激活专家的数量从69.5减少到14.6,同时保留了99.11%的原始性能。同时,它将内存使用量减少了76.64%到79.84%,并实现了1.14倍到1.66倍的端到端延迟加速。

🔬 方法详解

问题定义:论文旨在解决扩散语言模型(dLLMs)与混合专家模型(MoE)结合时,由于token级别的专家选择与dLLMs的块并行解码方式不匹配,导致推理过程中需要激活大量的专家,从而使得内存成为瓶颈的问题。现有方法中,每个token独立选择专家,忽略了块内token之间的依赖关系,造成了资源浪费。

核心思路:dMoE的核心思路是将一个块内的所有token的专家选择分布进行聚合,得到一个块级别的专家选择分布。然后,使用这个块级别的分布来指导整个块的专家路由。这样做的目的是减少需要激活的专家数量,因为一个块内的token更有可能共享相同的专家。

技术框架:dMoE框架主要包含以下几个步骤:1. 对于每个块,计算其中每个token的专家选择分布(例如,通过一个门控网络)。2. 将块内所有token的专家选择分布进行聚合,得到一个块级别的专家选择分布。聚合方法可以是简单的平均,也可以是更复杂的加权平均。3. 使用块级别的专家选择分布来选择该块需要激活的专家。4. 将块内的所有token路由到选定的专家进行处理。

关键创新:dMoE的关键创新在于提出了块级别的专家选择机制,而不是传统的token级别的专家选择。这种方法能够更好地利用块内token之间的依赖关系,从而减少需要激活的专家数量,降低内存需求。与现有方法的本质区别在于,dMoE考虑了块的整体信息,而传统方法只关注单个token。

关键设计:dMoE的关键设计包括:1. 如何有效地聚合块内的token级专家分布。论文中可能探讨了不同的聚合方法,例如平均、加权平均等。2. 如何根据块级别的专家选择分布来选择需要激活的专家。这可能涉及到一些阈值设置或者Top-K选择策略。3. 如何平衡性能和内存使用。减少激活专家数量可能会导致性能下降,因此需要在两者之间进行权衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,dMoE在保持99.11%原始性能的同时,将唯一激活的专家数量从69.5平均减少到14.6,内存使用量减少了76.64%到79.84%,端到端延迟加速了1.14倍到1.66倍。这些数据表明dMoE在降低内存需求和加速推理方面具有显著优势。

🎯 应用场景

dMoE技术可应用于各种需要大规模语言模型进行并行推理的场景,例如机器翻译、文本摘要、对话生成等。通过降低内存需求和加速推理速度,dMoE使得在资源受限的设备上部署大型语言模型成为可能,并有望推动AI技术在边缘计算和移动设备上的应用。

📄 摘要(原文)

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