← 返回 2026-06-01

dMoE：可学习块专家的扩散大语言模型 dMoE: dLLMs with Learnable Block Experts

Sicheng Feng, Zigeng Chen, Gongfan Fang, Xinyin Ma, Xinchao Wang 📅 2026-05-29 👍 38 2026-07-13 08:36
Diffusion LLM MoE 内存优化 块级路由 模型效率
arXiv HuggingFace PDF

块级路由优化MoE扩散模型内存效率

前置知识

混合专家模型

混合专家模型是一种稀疏激活的神经网络架构，通过路由器将每个token分配给少数几个专家进行计算。MoE的核心思想是增加模型总参数量但保持每次前向传播的计算量可控。在LLaDA2.0-mini中，有8个专家，每个token选择其中4个进行计算。路由器通过计算token表示与专家之间的相似度分数来决定激活哪些专家。

本文的核心问题是MoE在dLLM中的效率瓶颈，必须理解MoE的工作原理才能理解为什么需要块级路由来减少专家激活数量。

扩散大语言模型

扩散大语言模型将扩散过程引入文本生成领域，通过迭代式的去噪过程生成文本。与自回归模型逐token生成不同，dLLM使用掩码token机制，可以在一次前向传播中同时预测多个位置。具体来说，dLLM通过一个噪声调度器逐步添加掩码，然后使用去噪模型预测原始token。这种并行解码能力是dLLM相对于自回归模型的关键优势。

本文解决的核心问题正是源于dLLM的并行解码特性——块并行处理导致MoE专家选择存在根本性不匹配，理解dLLM的解码机制是理解问题本质的前提。

块扩散解码

块扩散解码是dLLM的一种高效推理策略，它将序列划分为多个块，每个块包含多个token（如32个）。在每次去噪步骤中，模型同时处理整个块内的所有掩码位置，而不是逐个token处理。这种机制充分利用了dLLM的并行解码能力，显著提升了推理速度。然而，当与MoE结合时，这意味着每次前向传播需要为块内的所有token选择专家，导致激活的专家数量成倍增加。

块扩散解码直接导致了本文要解决的问题——一个块内有多个token各自独立选择专家，导致唯一激活的专家数量过多（从70个降到16个），这是理解问题场景和技术方案的关键。

研究动机

现有MoE dLLM存在根本性架构不匹配问题。在块扩散解码中，dLLM需要在一次前向传播中同时处理多个token（如整个32个token的块），但现有的MoE路由仍然沿用自回归模型的token级专家选择范式，为每个token独立选择专家。这导致在一个前向传播中激活的唯一专家数量急剧增长，作者在LLaDA2.0-mini上的实验表明，平均激活的唯一专家数量达到69.5-71.5个（取决于数据集）。MoE延迟与唯一激活专家数量呈线性正相关（Pearson相关系数0.462），MoE部分占据了端到端延迟的主要部分（Transformers框架下MoE占比约60-70%）。这种内存访问瓶颈严重限制了MoE dLLM的实际部署效率，尤其是在资源受限的环境中。

本文的目标是本文的核心目标是解决MoE dLLM中块并行解码与token级专家选择之间的根本性不匹配，大幅减少推理过程中激活的唯一专家数量，同时不损失模型性能。具体而言，作者希望将平均唯一激活专家数量从约70个降低到接近每token选择的专家数量（即LLaDA2.0-mini中的8个），从而减少内存开销并提升推理速度。作者设定的性能目标是在如此激进的压缩下仍保留原始模型99%以上的性能，实现性能与效率的最佳平衡。

与已有工作不同的是，本文的独特切入角度在于提出了可学习的块级专家路由框架，这与现有方法形成鲜明对比。现有高效MoE策略主要针对自回归模型设计，分为两类：预执行压缩（如专家剪枝和专家合并）和运行时自适应执行（如专家跳过和专家复用）。针对dLLM的工作很少，EC-DLM将token-choice路由改为expert-choice路由，TEAM利用时空一致性复用专家，DES引入候选约束路由策略。然而，这些方法要么没有学习过程（如DES），要么没有充分利用块级信息。本文的创新之处在于首次提出了完全可学习的块级路由方案，通过自蒸馏训练让模型学会如何根据token级专家分数动态选择最优的块级专家子集，这是之前工作未曾探索的方向。

核心方法

dMoE采用粗到细的专家选择策略。首先，对于当前块内的所有token，计算每个token对各个专家的路由分数，其中ti是第i个token的表示。然后，通过聚合操作将所有token的专家分数汇总成块级专家分数Sblock，其中B是当前块的token集合。接着，对块级分数进行归一化，并应用top-p准则选择专家核心集C，其中p是累积概率阈值。最后，原始的token级路由被限制在这个核心集内进行，即每个token只从核心集C中的专家选择top-k个。训练时采用自蒸馏范式，使用与推理相同的前向过程和CSE损失。

dMoE的核心创新点在于通过两个关键观察实现了可学习的块级专家路由。第一个观察是token级专家分数能够有效指示专家重要性（Pearson相关系数0.462），这为将token级信息聚合为块级决策提供了理论基础。第二个观察是专家集中度在不同去噪步骤和块之间存在显著变化（变异系数CV达到15.93%-17.80%），这启发使用自适应的top-p准则而不是固定的核心集大小。与DES等非学习方法相比，dMoE通过自蒸馏训练让模型学会如何根据当前块的特征动态调整专家选择策略；与自回归模型的MoE优化方法相比，dMoE专门针对dLLM的并行解码特性设计，充分利用了块级信息。这种可学习性使得dMoE能够在极端压缩下仍保持近乎无损的性能。

方法步骤详情

dMoE的完整流程包含三个主要步骤。第一步是token级专家分数计算，对块内的每个tokenti，路由器输出对|E|个专家的分数向量si = [sE1, ..., sE|E|]。第二步是块级专家核心集选择，首先将所有token的专家分数聚合Sblock = sum_{i in B} si（也可以用其他聚合算子），然后归一化，最后应用top-p准则C = Top-P(Sblock, p)选择累积概率达到阈值p的专家作为核心集。第三步是受限的token级路由，对于每个token，将其专家分数中不属于核心集的部分屏蔽掉sC,i = MaskOutside(si, C)，然后从核心集中选择top-k个专家Ri = Top-K(sC,i, k)。训练时使用自蒸馏，前向过程与推理完全一致，损失函数确保模型学会正确的块级路由策略。

技术新颖性

dMoE的技术新颖性体现在多个方面。首先，这是首个针对MoE dLLM的可学习块级路由方案，填补了该领域的研究空白。其次，dMoE引入了自适应的top-p核心集选择机制，根据专家分数分布的集中程度动态调整核心集大小，这与固定大小的方法（如简单的top-K）相比更符合dLLM去噪过程的特点。第三，训练范式创新性地采用了自蒸馏，使得模型能够学习到如何从token级信号推断最优的块级专家选择策略。第四，dMoE与现有MoE优化方法正交，可以与专家剪枝、专家复用等技术结合使用，为MoE dLLM的效率优化提供了新的技术维度。最后，dMoE实现了从69.5个专家压缩到16.5个（76.15%压缩率）仍保留99.55%性能的突破性结果，证明了块级路由在MoE dLLM中的巨大潜力。

(a) We demonstrate the correlation between the router weights (token-level expert scores) and loss increase, drawn from GSM8K. (b) We present the distribution of the unique expert count during the inference process, drawn from GSM8K.
Figure 3: (a) We demonstrate the correlation between the router weights (token-level expert scores) and loss increase, drawn from GSM8K. (b) We present the distribution of the unique expert count during the inference process, drawn from GSM8K.
Overview of our proposed dMoE. For each noisy block, we aggregate token-level router scores into block-level expert scores and apply a top-p criterion to select an adaptive expert coreset. The final token-level routing is then restricted to this coreset. Training follows a self-distillation paradigm using CSE loss, and inference uses the same forward process.
Figure 4: Overview of our proposed dMoE. For each noisy block, we aggregate token-level router scores into block-level expert scores and apply a top-p criterion to select an adaptive expert coreset. The final token-level routing is then restricted to this coreset. Training follows a self-distillation paradigm using CSE loss, and inference uses the same forward process.

实验结果

dMoE在四个基准测试上取得了突破性的结果。在LLaDA2.0-mini上，dMoE（ptrain=0.6, ptest=0.6）将平均唯一激活专家数量从69.5个降低到16.5个，减少了76.15%，同时仅损失0.38%的性能（从83.95%降至83.575%），保留了99.55%的原始性能。具体来看，在MATH500上，专家数量从70.0降至16.3（-76.71%），精度从72.0%降至71.2%；在GSM8K上，从71.5降至16.6（-76.78%），精度从90.8%降至88.6%；在ARC-C上，从65.9降至16.9（-74.36%），精度从84.4%升至87.1%（竟然有提升）；在MMLU上，从70.6降至16.5（-76.63%），精度从88.6%降至87.4%。效率方面，dMoE将内存占用降低了74.36%-76.78%，实现了1.36x-2.04x的端到端延迟加速。在极端压缩设置下（ptrain=0.5, ptest=0.5），dMoE仍能将专家数量压缩到13.7个左右，同时保留97.5%的性能。消融实验表明，dMoE在不同块大小（8、16、24、32）下都保持稳定的性能-效率权衡，且top-p阈值可调节，能够适应不同的硬件特性和应用需求。

Results of the coefficient of variation (CV) on various benchmarks.
Table 1: Results of the coefficient of variation (CV) on various benchmarks.
Results on various benchmarks with LLaDA2.0-Mini. For MMLU, we use the math part for evaluation. Since the official DES code is not publicly available, we implement DES for comparison. Expert count denotes the average number of unique activated experts per layer.
Table 2: Results on various benchmarks with LLaDA2.0-Mini. For MMLU, we use the math part for evaluation. Since the official DES code is not publicly available, we implement DES for comparison. Expert count denotes the average number of unique activated experts per layer.
Ablation on the cumulative probability threshold p at training and test stages. We evaluate cumulative probability thresholds of 0.4, 0.5, 0.6, 0.7, and 0.8 while keeping other settings fixed.
Table 3: Ablation on the cumulative probability threshold p at training and test stages. We evaluate cumulative probability thresholds of 0.4, 0.5, 0.6, 0.7, and 0.8 while keeping other settings fixed.
(a) We report the average memory footprint of uniquely activated MoE parameters across all layers for both the baselines and our dMoE. (b) We present the average end-to-end latency and MoE latency for the original model and our dMoE.
Figure 5: (a) We report the average memory footprint of uniquely activated MoE parameters across all layers for both the baselines and our dMoE. (b) We present the average end-to-end latency and MoE latency for the original model and our dMoE.
Comparison of the performance-efficiency trade-off between our method and baselines. We report the results of GSM8K (a), MATH500 (b), and ARC-C (c).
Figure 6: Comparison of the performance-efficiency trade-off between our method and baselines. We report the results of GSM8K (a), MATH500 (b), and ARC-C (c).
查看结构化数据
任务指标本文基线提升
数学推理（MATH500） 准确率 71.2% 72.0% (LLaDA2.0-mini) 专家数量从70.0降至16.3（-76.71%），性能仅损失1.1%
数学推理（GSM8K） 准确率 88.6% 90.8% (LLaDA2.0-mini) 专家数量从71.5降至16.6（-76.78%），性能损失2.2%
科学推理（ARC-C） 准确率 87.1% 84.4% (LLaDA2.0-mini) 专家数量从65.9降至16.9（-74.36%），性能反而提升3.2%
一般推理（MMLU） 准确率 87.4% 88.6% (LLaDA2.0-mini) 专家数量从70.6降至16.5（-76.63%），性能损失1.2%
内存占用（所有数据集平均） 唯一激活MoE参数内存 降低74.36%-76.78% 原始模型 与DES-S相比额外降低59.18%-67.64%
推理延迟（所有数据集平均） 端到端加速比 1.36x-2.04x 原始模型 MoE延迟部分加速达到1.61x-2.04x

局限与改进

作者承认的局限性主要集中在评估范围和方法依赖性上。首先，本文只评估了语言模态，虽然作者提到方法原理上可以扩展到图像、视频等其他模态，但缺乏实证验证。其次，dMoE的训练依赖于自蒸馏，需要约700K训练样本（从GSM8K训练集、PRM12K、Numina-Math和OpenThoughts的子集生成），这在一定程度上增加了方法的训练成本。第三，dMoE的性能对top-p阈值较为敏感，需要根据不同数据集和任务特点进行调优。从我的观察来看，还有一些潜在的局限性：论文只在LLaDA2.0-mini这一个基座模型上验证了方法，缺乏在其他架构dLLM（如不同层数、不同专家数量）上的泛化性验证；论文主要关注推理效率，但训练效率的提升未被详细讨论；dMoE需要修改模型架构和训练流程，对于已经部署的模型可能需要重新训练才能应用。

独立分析的弱点

dMoE存在几个值得改进的弱点。首先，训练成本问题：论文中使用4块H100 GPU训练144小时（6天），这对于资源有限的研究团队来说是一个不小的负担。改进方向可以包括探索更高效的训练策略，如部分微调、LoRA等参数高效微调方法，或使用更小的合成数据集。其次，阈值敏感性：top-p阈值p对性能和压缩率有显著影响，需要针对不同应用场景手动调优。改进方向可以是设计自适应的阈值选择机制，根据输入特征或硬件约束动态调整p值，或者引入强化学习来学习最优的阈值策略。第三，方法适用性：dMoE需要在训练时就确定使用块级路由，这意味着无法直接应用于已经训练好的MoE dLLM。改进方向可以是探索无训练的蒸馏或适配方法，或者设计兼容自回归模型和dLLM的统一路由框架。第四，极端压缩的性能边界：当将专家数量压缩到接近每token选择的专家数量（如8个）时，性能仍有一定损失。改进方向可以是结合专家剪枝、专家合并等预执行压缩技术，在块级路由之前进一步减少候选专家数量。

未来方向

作者提出的未来工作方向包括将dMoE扩展到其他模态和多模态任务。具体而言，作者提到方法可以自然扩展到图像和视频等模态，以及视觉问答（VQA）和视觉推理等下游任务。基于本文的成果，我认为有几个值得深入探索的方向。首先是极端压缩：可以研究如何让块内的所有token共享同一个专家组，这将把唯一专家数量压缩到理论下限（即每token选择的专家数量），同时需要探索如何保持性能。其次是计算成本降低：可以联合优化专家选择策略，不仅减少唯一专家数量，还减少每token选择的专家数量（从k=4降到k=2或k=1），实现更激进的压缩。第三是动态自适应路由：可以探索更复杂的路由机制，根据去噪步骤、块位置、输入难度等因素动态调整路由策略，实现更精细化的效率优化。第四是与系统级优化的结合：可以将dMoE与专家通信优化、内存管理等系统级技术结合，进一步提升端到端效率。最后是跨模态验证：在实际的多模态任务上验证dMoE的有效性，探索不同模态之间的专家复用可能性。

复现评估

本文的复现性评估较为理想。代码已开源（https://github.com/fscdc/dMoE），这为其他研究者复现和扩展工作提供了便利。训练数据采用自蒸馏方式生成，约700K样本，从公共数据集（GSM8K、PRM12K、Numina-Math、OpenThoughts）构建，不依赖外部标注的高质量回答，确保了数据获取的可重复性。算力要求方面，论文明确指出使用了4块NVIDIA H100 GPU（每块80GB VRAM）进行训练，单次训练耗时约144小时（6天），评估在配备双AMD EPYC 9654处理器和1.5TiB内存的服务器上进行，每次评估约12小时。这些硬件配置虽然要求较高，但对于大模型研究来说是标准配置。基线对比中，由于DES的官方代码未公开，作者自行实现了DES-S和DES-V，这可能引入实现差异，但论文详细描述了实现细节。整体而言，本文提供了足够的技术细节和开源代码，复现难度中等偏高，主要障碍是算力需求而非技术不透明性。

论文图表

Figure 1: Comparison between the original LLaDA2.0-mini and our proposed dMoE. Unlike the original token-level expert routing in MoE dLLMs, our dMoE replaces token-level routing with block-level routing, substantially reducing the unique expert count while preserving performance.

Figure 1对比了原始LLaDA2.0-mini和dMoE的性能。左侧展示原始模型，在MATH500上达到72.0%准确率，但需要激活70.0个唯一专家；右侧展示dMoE，在MATH500上达到71.2%准确率，仅激活16.3个唯一专家。图中央展示了8个专家（E1-E8）和专家池的概念，直观说明了dMoE如何通过块级路由大幅减少激活专家数量。

这张图对理解论文至关重要，因为它以最直观的方式展示了dMoE的核心价值——在几乎不损失性能的前提下，将唯一激活专家数量减少了约76%，这是本文最重要的实验结果，也是论文的贡献所在。

Empirical studies on LLaDA2.0-mini. (a) & (b) We report the latency breakdown of three components: MoE (e.g., routing and execution), attention (e.g., attention computation and related projections), and others (e.g., LM head and embeddings). (c) We present the relationship between the number of unique activated experts and MoE kernel latency.
Figure 2: Empirical studies on LLaDA2.0-mini. (a) & (b) We report the latency breakdown of three components: MoE (e.g., routing and execution), attention (e.g., attention computation and related projections), and others (e.g., LM head and embeddings). (c) We present the relationship between the number of unique activated experts and MoE kernel latency.

Figure 2展示了LLaDA2.0-mini的实证研究结果。子图(a)和(b)分别展示了在Transformers和SGLang框架下，MoE、Attention和其他组件的延迟分解。结果表明MoE部分占据了端到端延迟的60-70%。子图©展示了唯一激活专家数量与MoE内核延迟的关系，呈现明显的线性正相关关系，相关系数r=0.46，支持了论文的核心假设。

这张图对理解论文的动机和方法设计至关重要。它通过实证数据证明了MoE延迟确实是推理瓶颈，且与唯一激活专家数量线性相关，这为dMoE的设计目标（减少唯一专家数量）提供了坚实的实验基础。

Table 4: Ablation on the block size. The evaluated model is trained with a block size of 32 with ptrain = 0.6. We set ptest = 0.6 here. We evaluate block sizes of 8, 16, 24, and 32 while keeping all other hyperparameters fixed.

Table 4展示了块大小的消融实验结果。表格分为四部分，分别对应块大小32、24、16、8。在每个块大小下，对比了原始LLaDA2.0-Mini和dMoE的专家数量和准确率。结果显示：块大小32时，dMoE将专家从70.0降至16.3（-76.71%），保留99.55%性能；块大小24时，从60.7降至16.2（-73.31%），保留98.62%性能；块大小16时，从48.6降至14.8（-69.55%），保留99.79%性能；块大小8时，从31.5降至14.1（-55.24%），保留98.81%性能。

这个表格对理解dMoE的泛化性至关重要。它展示了dMoE在不同块大小下都能保持稳定的性能-效率权衡，证明了方法不依赖于特定的块大小设置，增强了方法的适用性和鲁棒性。