DMax: Aggressive Parallel Decoding for dLLMs

会议: NeurIPS 2026
arXiv: 2604.08302
代码: https://github.com/czg1225/DMax
领域: LLM 其他
关键词: 扩散语言模型, 并行解码, 策略内均匀训练, 软并行解码, 自我纠错

一句话总结

DMax 针对掩码扩散语言模型在激进并行解码下的级联误差累积瓶颈，提出结合策略内均匀训练与软并行解码的新范式，将离散二值掩码解码重构为嵌入空间的自纠错过程，在 GSM8K 上将每步前向生成 Token 数量（TPF）从 2.04 翻倍跃升至 5.48 且无损精度。

研究背景与动机

扩散语言模型（dLLMs）因具备打破自回归逐词生成依赖、实现多位置高度并行解码的潜力，被视为新一代高效文本生成的强有力竞争者。然而在现有掩码扩散语言模型（MDLMs）中，提高并行度（如降低置信度阈值）会导致生成质量迅速崩溃。这一困境的深层瓶颈在于传统的二值掩码到离散 Token 解码机制：一旦某个掩码位置在当前步骤被解码为一个具体 Token，该预测便被永久固化并作为后续推理的固定上下文，即便发生错误也无法撤销或修正。

这种单向不可逆解码在激进并行解码时会引发灾难性后果。由于大量 Token 在单次前向传播中被同时预测，早期产生的局部幻觉或错误预测无法像投机解码那样得到验证与回滚，反而迅速污染后续步骤的条件上下文，引发级联式误差累积并最终导致语义崩溃。虽然均匀扩散语言模型（UDLMs）允许在任意词表 Token 间持续去噪重评，但若从纯粹的词表随机噪声起步，往往因噪声脱离自然语言流形而导致训练与生成极不稳定；且传统均匀加噪与模型推理时自身的预测误差存在严重的分布偏差。

为了破除不可逆固化与随机噪声不稳的两难困境，DMax 提出将二值离散翻转重构为嵌入空间中连续软状态的渐进式自纠错。核心 idea：通过策略内自采样构造预测噪声以消除分布偏差（OPUT），并在解码阶段引入置信度加权的预测与掩码混合软嵌入（SPD），使模型能够在保留全掩码稳定起步的同时，具备在推理过程中动态识别不确定性并迭代纠错的能力。

方法详解

整体框架

DMax 采用半自回归块级扩散（block-wise semi-autoregressive diffusion）生成范式。对于长度为 \(B\) 的当前块，模型首先将其初始化为全掩码状态，以继承掩码扩散模型稳定的去噪起步；在后续迭代去噪过程中，模型不仅预测剩余掩码，还会对当前块内已解码的 Token 进行持续重评与修正。整个框架由两项核心技术驱动：在后训练阶段通过策略内均匀训练（OPUT）让模型学会纠正自身犯下的错误，在推理阶段通过软并行解码（SPD）将中间状态建模为预测词嵌入与掩码嵌入的插值，并在连续前缀约束与双重收敛准则下实现激进并行加速。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入块初始化<br/>全块置为掩码嵌入 e_mask"] --> B["策略内均匀训练 OPUT<br/>模型自身分布采样构造错误预测"]
    B --> C["软并行解码 SPD<br/>置信度加权混合预测嵌入与掩码"]
    C --> D["最长连续置信前缀扫描<br/>阻止右侧高不确定性干扰左侧预测"]
    D --> E["双重块级收敛判定<br/>连续两步 Top-1 一致或置信度超标"]
    E -->|未收敛| C
    E -->|收敛| F["固化当前块输出并推进至下一块"]

关键设计

1. 策略内均匀训练（OPUT）：基于自身输出分布消除训练与推理偏差

传统 UDLM 直接使用词表均匀采样的随机噪声进行训练，但这些随机词严重脱离自然语言的语义流形，使模型浪费大量容量在无意义的词表映射上；更为致命的是，推理阶段模型遇到的错误上下文源于其自身的错误预测，而非均匀随机词，导致严重的分布失配。为解决这一痛点，OPUT 从模型当前的预测分布中按策略内（On-Policy）采样来构建噪声。

在每一次训练迭代中，给定干净序列 \(x_0 \sim \mathcal{D}\) 并随机采样加噪比例 \(t \sim \text{Uniform}(t_l, t_h)\)，首先生成掩码序列 \(x_t^{(m)}\)。将该序列输入预训练模型 \(M_\theta\)，对所有被掩码位置并行预测并按当前参数分布进行实时采样，得到包含模型自生成误差的预测噪声序列 \(x_t^{(p)}\)：

\[x_{t}^{(p),i}=\begin{cases}x_{t}^{(m),i},&\text{若 }x_{t}^{(m),i}\neq\texttt{[MASK]}\\\hat{x}^{i}\sim p_{\theta}(\cdot\mid x_{t}^{(m)}),&\text{若 }x_{t}^{(m),i}=\texttt{[MASK]}\end{cases}\]

随后，分别以掩码序列 \(x_t^{(m)}\) 和预测序列 \(x_t^{(p)}\) 作为输入执行两次前向传播，并在所有位置上监督其对真实干净 Token \(x_0\) 的交叉熵损失：

\[\mathcal{L}_{\mathrm{on\mbox{-}policy}} = \mathcal{L}_{\mathrm{mask}} + \mathcal{L}_{\mathrm{pred}} = -\sum_{i=1}^L \log p_\theta^{(m)}(x_0^i \mid x_t^{(m)}) - \sum_{i=1}^L \log p_\theta^{(p)}(x_0^i \mid x_t^{(p)})\]

通过将自身预测错误显式纳入监督目标，模型既完整保留了原本从掩码恢复 Token 的能力，又获得了从自身产生的错误 Token 中自愈和纠错的机制。

2. 软并行解码（SPD）：置信度加权的嵌入空间不确定性传递

单纯拥有纠错能力依然无法应对极端激进并行下的多点复合错误：当一整个块内所有 Token 被一次性粗暴解码时，强相关的错误信号会相互干扰导致震荡。SPD 的核心洞察在于，解码过程中的中间状态不应是二值的离散 Token，而应将前一步的预测不确定性平滑保留并传递至下一步。

对于块内已解码的位置 \(j \in \mathcal{T}^{(t)}\)，设上一步预测的 Top-1 词为 \(y_j^{(t-1)}\)，其对应置信度概率为 \(\pi_j^{(t-1)}\)。SPD 将剩余概率质量 \(1 - \pi_j^{(t-1)}\) 赋予代表最大不确定性的掩码嵌入 \(\mathbf{e}_{\mathrm{mask}}\)，构建未归一化的混合嵌入：

\[\tilde{\mathbf{h}}_j^{(t)} = \pi_j^{(t-1)} \mathbf{e}(y_j^{(t-1)}) + (1 - \pi_j^{(t-1)}) \mathbf{e}_{\mathrm{mask}}\]

为防止高维嵌入在向量叠加时发生模长漂移与范数崩溃，SPD 引入范数再归一化操作，使其模长精确对齐两者的概率加权范数：

\[\mathbf{h}_j^{(t)} = \frac{\tilde{\mathbf{h}}_j^{(t)}}{\|\tilde{\mathbf{h}}_j^{(t)}\|_2} \left( \pi_j^{(t-1)} \|\mathbf{e}(y_j^{(t-1)})\|_2 + (1 - \pi_j^{(t-1)}) \|\mathbf{e}_{\mathrm{mask}}\|_2 \right)\]

这种连续插值赋予模型清晰的置信度先验：高确信度的 Token 接近真实词嵌入，低确信度位置则向掩码嵌入靠近，提示注意力机制在自回归去噪中对低置信度位置保持审慎并优先修正。

3. 最长连续置信前缀与双重收敛准则：保障单向因果流与高效退出

为了防止序列右侧高不确定性的投机预测向左污染尚在形成的稳定前缀，SPD 在每次迭代时执行严格的从左至右扫描：只有在掩码区域中置信度超过解码阈值 \(\tau_{\mathrm{dec}}\) 的最长连续前缀才能晋升为 Token 状态；一旦遇到首个置信度不达标的位置，其右侧的所有位置均强制维持掩码。若前缀为空，则强制仅释放最左侧的一个 Token，确保解码单调推进。

在块收敛判定上，DMax 设计了双重准则：（1）一致性收敛，即整块内所有位置的 Top-1 预测在连续两次迭代中完全相同；（2）置信度收敛，即整块所有 Token 的预测概率均超过高验收阈值 \(\tau_{\mathrm{acc}}\)（默认 0.9）。一旦满足任一条件，即视该块达到稳定态并一次性提交输出，极大缩减了无效的前向迭代开销。

损失函数 / 训练策略

训练以开源基础模型 LLaDA-2.0-mini 为基底，在 8 张 H200 GPU 上采用全参数微调（2 个 epoch，余弦学习率衰减，初始学习率 \(2 \times 10^{-6}\)，批大小 8，固定加噪比例 0.75，块大小 32）。为了避免额外显存峰值开销，掩码序列前向计算与预测序列前向计算在同一 epoch 内被解耦为独立的微调步交替更新。训练数据完全基于自蒸馏构建（使用基座模型以 0.95 极高置信度自生成，筛选出 0.7M 数学样本与 1.0M 代码样本），不依赖任何外部教师模型的监督。

实验关键数据

主实验

在主流数学推理与代码生成基准上，DMax 展现出显著超越原始模型与其他并行加速基线的生成效率，并在极高并行度下维持了模型固有的高准确率（测试环境为 2 张 H200 GPU，Tensor Parallelism，Batch Size 1）：

基准任务	方法	TPF ↑	TPS (tokens/s) ↑	准确率 (Acc.) ↑	AUP Score ↑
GSM8K	LLaDA-2.0-mini	2.04	512	92.6%	340
GSM8K	Hierarchical Decoding	2.44	577	91.6%	357
GSM8K	dParallel SFT	2.79	721	92.3%	395
GSM8K	Uniform Diffusion Training	2.26	493	68.7%	0
GSM8K	DMax-Math (本文)	5.48	1258	92.1%	557
MATH500	LLaDA-2.0-mini	2.58	626	75.8%	257
MATH500	Hierarchical Decoding	3.01	669	73.0%	268
MATH500	dParallel SFT	3.42	823	75.8%	310
MATH500	Uniform Diffusion Training	2.43	530	33.6%	0
MATH500	DMax-Math (本文)	5.94	1286	75.4%	507
Minerva-Algebra	LLaDA-2.0-mini	3.01	755	91.4%	363
Minerva-Algebra	DMax-Math (本文)	7.03	1492	91.5%	658
HumanEval-Instruct	LLaDA-2.0-mini	4.38	1044	84.2%	369
HumanEval-Instruct	DMax-Coder (本文)	7.36	1557	83.5%	637
MBPP-Instruct	LLaDA-2.0-mini	2.71	662	80.6%	276
MBPP-Instruct	DMax-Coder (本文)	5.86	1264	79.2%	482

消融实验

表 3 验证了在 GSM8K 上各个核心训练机制与推理模块在不同激进解码阈值 \(\tau_{\mathrm{dec}}\) 下的作用（当 \(\tau_{\mathrm{dec}}=0.0\) 时代表每步将整个块的所有掩码激进释放）：

On-Policy Rollout	连续前缀约束	混合软嵌入 (SPD)	\(\tau_{\mathrm{dec}}=0.95\) (TPF / Acc)	\(\tau_{\mathrm{dec}}=0.50\) (TPF / Acc)	\(\tau_{\mathrm{dec}}=0.0\) (TPF / Acc)
-	-	-	2.04 / 92.6%	4.47 / 78.0%	7.86 / 0.9%
-	✓	✓	1.04 / 0.0%	1.73 / 0.0%	5.39 / 0.0%
✓	-	-	2.95 / 92.6%	5.14 / 90.1%	5.89 / 68.2%
✓	✓	-	2.85 / 93.0%	5.28 / 91.3%	5.89 / 68.2%
✓	-	✓	3.25 / 92.8%	5.64 / 91.4%	6.01 / 90.4%
✓	✓	✓	3.00 / 93.3%	5.48 / 92.1%	6.01 / 90.4%

关键发现

• OPUT 是软解码生效的前提：如果不在训练阶段赋予模型从自身错误中恢复的能力，直接对基座模型应用混合软嵌入推理（第二行配置），模型在所有阈值下全部崩溃（准确率归零）。只有经过 OPUT 训练，模型才具备理解插值状态并去噪的能力。
• 激进并行下的韧性飞跃：在极端全量并行设置（\(\tau_{\mathrm{dec}}=0.0\)）下，原始模型因误差累积彻底失效（仅 0.9%），仅依靠 OPUT 纠错也只能恢复至 68.2%；而结合软嵌入后准确率高达 90.4%，证明连续不确定性传播对于抑制突发多点误差至关重要。
• 低并行度下的逆向增益：如表 2 所示，在保守的低并行设置下，DMax 由于具备自纠错机制，甚至能在 GSM8K（+0.8%）、MATH500（+2.2%）和 HumanEval（+3.0%）上反超原模型的基线精度，纠正了原有单向路径上的偶发逻辑错误。

亮点与洞察

• 将离散硬决策松弛为连续软嵌入：通过置信度将预测 Token 与掩码 Token 线性插值，利用掩码嵌入天然表征最大不确定性的几何先验，在无需修改模型底层架构的前提下将不确定性显式向前传递。
• 策略内噪声采样优雅解决流形偏离问题：摒弃传统 UDLM 从词表均匀采样的低效方案，直接从模型当前步的自身输出中取样，完美拟合了推理时自生成错误的真实统计分布。
• 纯自蒸馏无额外标注依赖：全量训练数据均由基础模型在保守参数下离线自生成，证明了 dLLM 的自纠错与并行加速潜力完全蕴含于自身参数分布之中，具有极高的方法通用性。

局限与展望

• 块内局部视野受限：当前软并行解码基于分块（Block-wise）自回归框架展开，跨块的误差依然属于单向不可逆范式，长程逻辑层面的全局错误修正仍有待进一步探索。
• 固定加噪比与两步训练成本：训练期固定掩码比例为 0.75，且每个训练样本需要额外的 rollout 采样与两次前向传播计算，训练阶段的计算开销约为普通微调的 2 倍。
• 更细粒度的动态置信阈值：未来可探索根据文本序列熵自适应调整解码阈值 \(\tau_{\mathrm{dec}}\) 与验收阈值 \(\tau_{\mathrm{acc}}\)，在保持极速吞吐的同时进一步挖掘极限推理质量。

相关工作与启发

• vs 传统 MDLMs (LLaDA / Dream): 传统掩码扩散模型采用严格的单向掩码替换机制，一旦做出预测便永久固化；DMax 通过引入自纠错和混合嵌入机制，成功打破了并行度提高与生成质量崩溃之间的负相关。
• vs 传统 UDLMs (Uniform Diffusion): 传统均匀扩散训练从全词表随机替换出发，加噪样本脱离语言流形且与推理阶段的预测错误存在严重分布偏差；DMax 的策略内采样大幅缩小了该偏差，收敛稳定且快速。
• vs 投机解码 (Speculative Decoding): 投机解码依赖额外的小型草稿模型与严格的验证回滚流程，而 DMax 将草稿生成与纠错自愈统一在单个扩散模型自身的迭代去噪动力学中。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 巧妙重构掩码与均匀扩散关系，软嵌入与策略内采样设计精练且直击痛点。
• 实验充分度: ⭐⭐⭐⭐⭐ 覆盖主流数学与代码基准，消融实验直观展示极端并行阈值下的模型鲁棒性。
• 写作质量: ⭐⭐⭐⭐⭐ 动机分析深入透彻，图表传达清晰直观。
• 价值: ⭐⭐⭐⭐⭐ 为推动扩散语言模型走向工业级高吞吐落地树立了强有力的里程碑。

相关论文

• [NeurIPS 2026] Dependency-Guided Parallel Decoding in Discrete Diffusion Language Models
• [NeurIPS 2026] Depth Exploration for LLM Decoding
• [NeurIPS 2026] Confidence-Based Decoding is Provably Efficient for Diffusion Language Models
• [NeurIPS 2026] On the Nature of Attention Sink that Shapes Decoding Strategy in Omni-LLMs
• [NeurIPS 2026] Constrained Decoding for Diffusion Language Models via Efficient Inference over Finite Automata