dMoE: dLLMs with Learnable Block Experts¶

会议: NeurIPS 2026
arXiv: 2605.30876
代码: https://github.com/fscdc/dMoE
领域: 模型压缩
关键词: 扩散大语言模型, 混合专家模型, 块级路由, 内存瓶颈, 并行解码

一句话总结¶

针对扩散大语言模型在块并行解码时因逐 token 独立路由导致的激活专家数暴涨与显存受限瓶颈，dMoE 通过汇聚块内 token 路由分布并利用自适应 Top-\(p\) 截断选拔专家核心集，在保持单 token 激活专家数不变的前提下将去重激活专家数从 69.5 骤降至 16.5，显存占用降低超 74%，实现最高 2.04 倍的端到端推理加速且几乎无损精度。

研究背景与动机¶

扩散大语言模型（dLLMs，如 LLaDA 系列）以其非自回归迭代去噪的特性，打破了传统自回归模型严格的从左到右单向依赖，在自然支持块级并行解码（Block Diffusion Decoding）方面展现了显著的吞吐优势。为了进一步扩展模型容量并控制推理激活计算量，将稀疏混合专家（MoE）架构引入 dLLM 成为当下模型规模化的主流范式。然而，在基于块并行解码的扩散推理过程中，传统 MoE 层依然沿袭了自回归模型的逐 token 独立选择专家（Token-level Routing）机制。

这种解码范式与路由机制的错配引发了严重的效率灾难：在 dLLM 的单次前向去噪中，模型需要同时预测当前块（如块大小为 32）内的多个 token；由于双向上下文依赖与多位置的独立分散路由，同一层内所有 token 激活的去重专家总数（Unique Activated Experts）急剧膨胀（LLaDA2.0-mini 单层平均激活近 70 个独特专家）。此时推理性能不再受算力限制，而是被频繁的专家权重显存搬运与内核加载彻底拖垮，使 MoE 内核耗时成为端到端推理延迟的主导瓶颈（耗时占比高达 80% 以上，且与去重专家数呈强线性正相关）。

尽管已有工作尝试通过序列级路由或投票机制进行无训练后处理截断，但在极端压缩下性能会发生灾难性雪崩，且无法动态应对去噪步间显著波动的专家集中度。本文切入点在于利用 token 级打分对专家重要性的保序性，在块级别对路由偏好做整体一致性收敛。核心 idea：将块内 token 级路由打分聚合为块级专家分布，通过自适应累积概率（Top-\(p\)）截断动态选拔块级专家核心集，并在同构前向自蒸馏训练下将逐 token 路由严格限定在该核心集内，实现深度的显存带宽解脱。

方法详解¶

整体框架¶

dMoE 旨在保持每个 token 激活的专家数量（如 top-\(k=8\)）完全不变以保障模型容量，同时大幅压低整个块在当前层所激活的独特专家并集大小 \(|\mathcal{C}|\)。整个推理前向过程划分为三个核心阶段：首先在块级别并行计算所有 token 对完整专家库的路由打分，随后将 token 级分数在块内进行汇聚归一化并执行 Top-\(p\) 动态筛选得到当前块的专家核心集（Coreset），最后在屏蔽核心集之外所有专家后重新对每个 token 独立执行 Top-\(k\) 选择与特征融合加权。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入噪声块 Tokens<br/>t_i, i ∈ B (块大小 32)"] --> B["Token 级初步路由打分<br/>计算 s_i = Router(t_i)"]
    B --> C["块级专家分数汇聚<br/>累加各 token 掩码得分并归一化"]
    C --> D["自适应 Top-p 核心集筛选<br/>动态确定块级专家集合 C"]
    D --> E["核心集内局部 Top-k 路由<br/>屏蔽外部专家，选取各 token 激活专家"]
    E --> F["受限 MoE 前向计算与加权输出<br/>仅加载 |C| 个专家显存权重"]

关键设计¶

1. 块级专家分数汇聚：保留 token 级重要性信号的全局偏好投影

以往启发式方法往往直接根据多数投票或固定启发规则粗暴裁剪专家，容易丢失对特定关键 token 具有决定性作用的专家。作者经验性验证表明，token 级路由器打分与掩蔽该专家后的重构损失增量呈正相关（Pearson 相关系数达到 0.462），证明局部权重是专家重要性的可靠先验。为此，dMoE 提出显式块级打分汇聚机制：

\[S_{block} = \sum_{i \in \mathcal{B}} \operatorname{TopKMask}(\mathbf{s}_i, k)\]

其中 \(\mathbf{s}_i = \operatorname{Router}(t_i) \in \mathbb{R}^{|\mathcal{E}|}\) 为块 \(\mathcal{B}\) 内第 \(i\) 个 token 的原始路由向量，\(\operatorname{TopKMask}\) 仅保留每个 token 得分前 \(k\) 的专家打分（其余置零），避免低分非活跃专家的背景噪声累积。随后将累加向量 \(S_{block}\) 经过 Softmax 或和归一化处理得到反映当前块整体需求的块级专家概率分布 \(\tilde{S}_{block}\)。

2. 自适应 Top-\(p\) 核心集筛选：拟合去噪步间与块间专家集中度动态波动

实证分析发现，扩散模型在不同去噪时间步 \(t\) 及不同层之间，专家激活的集中程度存在极大的动态跨度。如果在整个去噪过程中强行指定固定的专家候选池大小，在专家需求离散的步骤会引发严重容量瓶颈，而在高度集中的步骤则存在冗余。统计表明，GSM8K、MATH500 等基准上的去重专家数变异系数（CV）高达 15.93% 至 17.80%。

因此，dMoE 摒弃了固定大小的硬截断，引入自适应累积概率（Top-\(p\)）截断策略：

\[\mathcal{C} = \operatorname{Top-P}(\tilde{S}_{block}, p)\]

按排序后的块级得分从大到小依次纳入专家，直至累积概率达到超参数阈值 \(p\)。当上下文语义明确、专家偏好高度一致时，该准则自动收敛为极小核心集（如 12~14 个专家）；当去噪初期不确定性极高或块内语义丰富时，核心集自适应扩展以容纳更多专家多样性，完美契合了扩散模型由粗到细去噪阶段的物理特性。

3. 核心集受限 Token 路由：零微观结构破坏的约束解码与自蒸馏对齐

在获得当前块的专家核心集 \(\mathcal{C}\) 后，为了保证每个 token 的计算浮点数（FLOPs）和模型表达深度完全不退化，每个 token 依然需要在核心集内分配 \(k\) 个专家。算法构造受限打分向量 \(\mathbf{s}_i^\mathcal{C}\)：

\[\mathbf{s}_{i}^{\mathcal{C}}[e] = \begin{cases} \mathbf{s}_i[e], & \text{if } e \in \mathcal{C} \\ -\infty, & \text{otherwise} \end{cases}\]

随后每个 token 独立执行局部选择 \(R_i = \operatorname{TopK}(\mathbf{s}_i^\mathcal{C}, k)\) 并完成常规 MoE 加权。这种设计使得 GPU 硬件在执行当前块的 MoE 内核时，仅需将显存中的 \(|\mathcal{C}|\) 个专家权重拉取到高速缓存或片上 SRAM 中，彻底消除大范围散布权重加载。在训练阶段，模型无需外置额外复杂的教师网络，而是直接采用同构自蒸馏（Self-Distillation），以前向传播中施加核心集截断后的输出去拟合全量未截断的基座模型生成目标，使得路由器与 FFN 权重主动学会协作与信息代偿。

损失函数 / 训练策略¶

dMoE 采用基于自蒸馏的全参数微调方案。训练数据完全基于基座模型（LLaDA2.0-mini）自身在 GSM8K、PRM12K、Numina-Math 和 OpenThoughts 上由置信度阈值 0.95 生成的无外源监督响应（约 70 万条）。在训练期间，掩码率在 \([0.3, 0.8]\) 随机采样，块大小固定为 32，训练 Top-\(p\) 阈值设为 \(p^{train}=0.6\)。优化器采用 Cosine 学习率调度，峰值学习率为 \(2.0 \times 10^{-6}\)，在 4 张 NVIDIA H100 GPU 上迭代 2 个 epoch，使得模型在有限微调开销下快速适应块级路由约束。

实验关键数据¶

主实验¶

在 LLaDA2.0-Mini 基座模型上，dMoE 与原版模型以及现有代表性 MoE 压缩基线进行了全面对比。测试涵盖数学推理（MATH500、GSM8K）、科学推理（ARC-C）和高阶通用推理（MMLU-Math）。

方法	MATH500 专家数 / 准确率	GSM8K 专家数 / 准确率	ARC-C 专家数 / 准确率	MMLU-Math 专家数 / 准确率	平均激活专家数	平均准确率
LLaDA2.0-Mini (原版)	70.0 / 72.0%	71.5 / 90.8%	65.9 / 84.4%	70.6 / 88.6%	69.50	83.95%
Top-4 (降 token 专家数)	43.5 / 52.6%	44.3 / 72.4%	42.2 / 74.0%	43.8 / 61.5%	43.45	65.13%
DES-Seq (序列级压缩)	41.0 / 70.8%	41.9 / 90.7%	39.9 / 83.3%	41.4 / 87.1%	41.05	82.98%
DES-Vote (投票级压缩)	36.9 / 71.8%	37.0 / 90.5%	36.8 / 82.3%	37.0 / 85.6%	36.93	82.55%
DES-Seq* (极端压缩)	12.9 / 53.8%	13.2 / 74.0%	13.0 / 62.9%	13.0 / 70.7%	13.03	65.35%
DES-Vote* (极端压缩)	13.0 / 53.0%	12.9 / 73.5%	13.0 / 49.1%	13.0 / 72.6%	12.98	62.05%
dMoE (\(p^{train}=0.6\))	16.3 / 71.2%	16.6 / 88.6%	16.9 / 87.1%	16.5 / 87.4%	16.58	83.58%
dMoE (\(p^{train}=0.5\))	13.7 / 69.8%	13.8 / 87.1%	13.9 / 84.7%	13.6 / 86.7%	13.75	82.08%

消融实验¶

以 \(p^{train}=0.6\) 训练的 dMoE 在不同块大小（Block Size）下的鲁棒性消融（固定 \(p^{test}=0.6\)）：

块大小配置	MATH500 专家数 / 准确率	GSM8K 专家数 / 准确率	ARC-C 专家数 / 准确率	MMLU 专家数 / 准确率	平均专家数缩减率	相对基线性能保留率
Block size = 32 (原版)	70.0 / 72.0%	71.5 / 90.8%	65.9 / 84.4%	70.6 / 88.6%	基线 (69.5)	100.0% (83.95%)
Block size = 32 (dMoE)	16.3 / 71.2%	16.6 / 88.6%	16.9 / 87.1%	16.5 / 87.4%	-76.15% (16.58)	99.55% (83.58%)
Block size = 24 (原版)	60.7 / 73.6%	62.2 / 91.0%	57.8 / 82.3%	61.6 / 86.9%	基线 (60.58)	100.0% (83.45%)
Block size = 24 (dMoE)	16.2 / 68.6%	17.4 / 87.4%	16.1 / 87.9%	15.8 / 85.3%	-72.97% (16.38)	98.62% (82.30%)
Block size = 16 (原版)	48.6 / 70.2%	49.8 / 90.5%	47.0 / 86.9%	49.6 / 88.5%	基线 (48.75)	100.0% (84.03%)
Block size = 16 (dMoE)	14.8 / 70.0%	16.7 / 89.9%	15.1 / 88.1%	14.9 / 87.4%	-68.46% (15.38)	99.79% (83.85%)
Block size = 8 (原版)	31.5 / 71.0%	32.2 / 90.1%	31.1 / 87.7%	32.3 / 87.4%	基线 (31.78)	100.0% (84.05%)
Block size = 8 (dMoE)	14.1 / 70.2%	16.1 / 88.1%	14.2 / 87.2%	14.1 / 86.7%	-53.97% (14.63)	98.81% (83.05%)

关键发现¶

极端压缩极限逼近理论下界：LLaDA2.0-mini 每个 token 必须选 8 个专家。dMoE 在 \(p^{train}=0.5\) 且 \(p^{test}=0.4\) 时，能将全块独特专家数压到 13.0 个左右，几乎接近单 token 的理论下限 8，同时依然保持 80% 以上的高准确率；而强行将 DES 等基线压缩至 13 个专家时，准确率发生毁灭性下跌（DES-Vote* 平均准确率从 82.55% 暴跌至 62.05%）。
显存与时延实现突破性优化：dMoE 将激活 MoE 权重平均显存开销降低了 74.36% ~ 76.78%（相比基线方法进一步减少 59.18% ~ 67.64%），在端到端推理上录得 1.36× 至 2.04× 的实际物理加速比（主要归因于 MoE 内核加载时间减少 1.14× 至 1.66×）。
推断阶段阈值具有连续可调性：测试期 \(p^{test}\) 可以在 0.4 到 0.8 之间连续调节。增大 \(p^{test}\) 可渐进提升模型精度，甚至在 \(p^{test}=0.8\) 时在 GSM8K 和 ARC-C 上超越原版基座性能，展现出面向多级算力设备极高的部署弹性。

亮点与洞察¶

抓住扩散并行解码核心矛盾：敏锐地指出 dLLM 与 MoE 结合时，块并行使得微观上分散的 token 级路由在宏观上叠加为显存带宽爆炸，直击 MoE 扩散模型真正的推理瓶颈。
自适应 Top-\(p\) 替代硬性截断：紧扣去噪过程专家需求集中度随时间步大幅波动的物理规律，通过无参动态概率截断使得核心集大小随上下文确定度弹性伸缩。
简单优雅的零破坏设计：不改变单 token 的激活专家数与 FFN 算力结构，仅通过前置打分掩码实现跨 token 协同路由，易于集成到各类开源部署引擎（如 vLLM/SGLang）。

局限与展望¶

实验受限于单文本模态与单基座模型：论文目前仅在语言任务与 LLaDA2.0-mini 上进行了详尽测试，尚未在多模态扩散大模型（如 DiT、视觉扩散模型）上展开规模化验证。
未能联合压缩计算量：目前方法保持了每个 token 的 Top-\(k=8\) 计算量，未来若能将块级专家路由与 token 级计算跳过（Dynamic Expert Skipping）相结合，有望实现显存搬运与计算 FLOPs 的双重极致压缩。

评分¶

新颖性: ⭐⭐⭐⭐ [针对 dLLM MoE 显存墙提出块级动态路由与自蒸馏机制，思路清晰且直击痛点]
实验充分度: ⭐⭐⭐⭐⭐ [对比全面，消融实验涵盖阈值动态范围、不同块大小以及端到端真实显存与时延评测]
写作质量: ⭐⭐⭐⭐⭐ [逻辑层层递进，动机论证与实证观察互为表里，结构严谨规范]
价值: ⭐⭐⭐⭐⭐ [为端侧或显存敏感场景部署大规模并行生成 dLLM 提供了高度落地的基准方案]