DMax:面向扩散语言模型的激进并行解码

用on-policy训练+软嵌入解码让dLLM并行度翻倍且不丢精度。

前置知识

扩散语言模型 (Diffusion Language Model, dLLM)

一种把连续域扩散思想离散化后用于文本生成的语言模型,代表工作有 LLaDA、Dream 等。它不采用自回归(AR)的「逐 token 生成」,而是把整段文本视作联合去噪对象,通过多轮并行预测填充 [MASK] 来生成完整序列,理论上有更高的并行推理效率。

DMax 正是为 dLLM 设计的解码算法,理解 dLLM 的「mask-to-token」迭代过程、它与 AR 模型在并行性上的差异,以及目前 dLLM 并行度受限于误差累积的现状,是读懂本文 motivation 的前提。

掩码扩散语言模型 (MDLM) 与均匀扩散语言模型 (UDLM)

MDLM(如 LLaDA)训练时把 token 替换为 [MASK],推理从全 [MASK] 序列开始逐步替换;UDLM 则把 token 替换为「从词表均匀采样的随机 token」,对所有位置施加去噪损失,天然具自修正能力,但起点噪声过大致生成不稳定。

论文把 MDLM 的「稳定起点」与 UDLM 的「自修正」视为互补,提出 On-Policy Uniform Training 把两者融合;不理解这两种范式在「训练目标、推理起点、可修正性」上的差异,就抓不住 OPUT 的核心动机。

误差累积与并行解码 trade-off

dLLM 的解码是把 [MASK] 替换为 token 的「一次性承诺」:一旦某位置解码完成,它在后续步里就作为不可修改的上下文传给其他位置。当一次性并行解码多个 mask 时,若其中有错,错误会以「错误上下文」形式污染后续所有预测,触发连锁崩坏,导致语义崩塌。

DMax 的全部设计——OPUT 让模型学会修正自己的错误、SPD 用软嵌入保留不确定性——都是为了直接缓解误差累积;理解这条因果链是理解 DMax 一切设计选择的钥匙。

On-Policy 训练 (On-Policy Rollout)

在强化学习里指「用当前策略自身产生的轨迹作为训练数据」的训练方式。本文特指:在每次迭代中先用当前模型 M_θ 在带噪输入上做一次「无梯度预测」,得到一组 sampled noisy sequence,再把这组序列当作下一轮去噪训练的输入,使训练分布与推理分布严格匹配。

OPUT 是 DMax 与传统 UDLM 训练最关键的区别:传统 UDLM 用「词表均匀随机 token」构造噪声,与模型实际推理时遇到的「自身预测」分布严重错位,导致自修正能力训练不充分;on-policy rollout 直接消除这一 train–inference gap,是 DMax 训练侧的核心创新。

TPF / TPS / AUP Score

TPF (Tokens Per Forward) 指一次前向平均解码的 token 数,反映并行度;TPS (Tokens Per Second) 反映实际推理吞吐;AUP Score 是把不同并行度下的精度和速度综合的「曲线下面积」,数值越大越好。

本文所有对比都同时报告这三个指标,核心论点是「TPF 翻倍但精度不降」;理解 AUP 含义能帮助读懂 Table 1 和 Figure 4 中 DMax 相对基线全面领先的具体含义。

块扩散 (Block Diffusion) 与半自回归

把整段生成序列切成长度为 B 的若干块(block),块内采用并行 dLLM 方式迭代去噪,块与块之间按顺序处理,是一种「局部并行、整体自回归」的混合范式。本文 block size = 32,既保留长序列的稳定性又支持大规模并行。

DMax 的 SPD 算法就是按 block 设计的,包括「块内连续前缀策略」「块级收敛条件」等都依赖 block 概念;不理解 block diffusion 的话,Algorithm 1 与 Figure 3 都会显得很抽象。

研究动机

扩散语言模型(dLLM)被广泛认为是 AR-LLM 的有力替代,核心吸引力来自并行解码能力。然而现有基于 MDLM 范式的 dLLM(LLaDA、Dream)存在根本瓶颈:在激进并行解码下性能断崖式下跌。MDLM 的解码是单向的 mask-to-token 二值过程,一旦某 [MASK] 位置被解码为 token,就被当作不可修改的上下文,无法修正。在高并行度下,几乎必然出现错误预测,这些错误会向后续去噪步传播,造成连锁式误差累积,最终导致语义崩塌。具体数据:当 $\tau_{dec}=0$ 时,即便已过 OPUT 训练的 LLaDA-2.0-mini 在 GSM8K 上只取得 68% 精度,远低于 0.95 阈值下的 92.6%。标准 UDLM 训练(用均匀随机 token 作噪声)不仅没提升并行度,反而让 GSM8K 精度从 92.6% 暴跌到 68.7%,原因是「词表均匀随机 token」严重偏离模型实际推理时遇到的「自身预测」分布,造成 train–inference 严重错配。

本文的目标是本文目标是设计一种新型 dLLM 范式 DMax,使模型在「激进并行解码」(TPF 翻倍以上)下仍能保持甚至略微超过原模型的精度。具体包含三个子目标:(1)在训练侧提出 On-Policy Uniform Training(OPUT),把预训练 MDLM 扩展为兼具 mask 去噪与「自预测去噪」能力的统一模型,且保持训练成本可控(2 epoch, batch size 8, 8 H200);(2)在推理侧提出 Soft Parallel Decoding(SPD),用「token embedding 与 mask embedding 的插值」替代离散的硬 token,让模型在每一步都能感知到之前预测的不确定性,从而更好地做自我修正;(3)在 6 个标准 benchmark(GSM8K、MATH500、Minerva-Algebra、ASDIV、HumanEval、MBPP)上同时实现「并行度大幅提升 + 精度不降 + 推理吞吐过 1000 TPS」,并把 AUP Score 推到新的 SOTA。

与已有工作不同的是，已有工作大致从三类角度缓解并行度问题:改进解码策略(Hierarchical Decoding 等)、模型蒸馏(dParallel-SFT)、KV cache / token dropping 等系统优化。这些方法都没触及问题的根本——MDLM 的二值 mask-to-token 范式本身缺乏「自我修正」机制。UDLM 范式虽天然可修正,但「从全随机 token 出发」使起点噪声太大、训练不稳定。本文的核心切入角度是「在保留 MDLM 稳定起点的同时,赋予其 UDLM 式的自修正能力」,技术上又分为两步走:训练侧用 on-policy rollout 消除 train–inference 错配(而非用均匀随机 token 凑合),推理侧用混合软嵌入(而非离散 token)传递不确定先验。这条路线与 dParallel-SFT「用 SFT 强行拟合高置信解码」不同,也与 SM/EvoToken「把软嵌入用于其他目的」不同,是把「on-policy 训练 + 软嵌入解码」首次系统地组合起来用于 dLLM 并行加速。

核心方法

DMax 的方法论围绕「让模型学会修正自己」这一直觉,把 MDLM 的 mask-to-token 二值解码重新形式化为「mask 嵌入 → 混合软嵌入 → token 嵌入」的可微过程。训练阶段,在标准 MDLM 训练基础上,增加一条 on-policy rollout 分支——先用当前模型 $M_\theta$ 在 75% 掩码率的带噪序列上预测所有 mask 位置并采样得到「预测噪声序列」,再把这个「预测序列」当作输入做一次普通去噪,对两次前向都计算去噪损失 $\mathcal{L}_{mask} + \mathcal{L}_{pred}$,反向传播更新 $\theta$。这让模型既保留 mask 去噪能力,又学到「从自身错误预测回到正确 token」的能力。推理阶段,对每个长度为 32 的 block,从全 [MASK] 初始化开始,每一步用高阈值 $\tau_{dec}$ 从左到右扫描,只把连续高置信前缀里的 mask 位置提升为 token 位置,并用混合嵌入作为下一轮输入,迫使模型在下一轮重新评估不确定位置,从根本上打破「错误一旦产生就无法修正」的魔咒。

DMax 的核心创新是「on-policy 训练 + 软嵌入解码」这一组合,本质区别于已有工作体现在三点。第一,训练侧的标准 UDLM 用「词表均匀随机 token」构造噪声,而 OPUT 用「模型自身当前参数采样得到的预测 token」构造噪声,使训练分布严格匹配推理分布;第二,推理侧 SPD 不再用「one-hot token 嵌入」作为去噪起点,而是用「置信度加权插值」后的混合软嵌入 $h_j = \pi_j \cdot e_{\hat{y}_j} + (1-\pi_j) \cdot e_{mask}$,把不确定性显式编码到模型输入,让模型能区分「自信预测」与「待修正预测」;第三,块内解码采用「从左到右最长连续高置信前缀」策略,保证 mask 区域始终是「一个连续的尾巴」,防止右侧未预测的不可靠 token 干扰左侧去噪——这是 SM、EvoToken 等同样使用软嵌入的工作所没有的设计,也是 DMax 在 $\tau_{dec}=0$ 这种极端激进设置下仍能保持 90% 精度的关键。

方法步骤详情

DMax 流程分五阶段。数据准备:从 GSM8K、PRM12K、Numina-Math、OpenThoughts 抽 0.7M 数学 prompt,从 OpenCodeInstruct 抽 1.0M 代码 prompt,用 LLaDA-2.0-mini 在 τ=0.95、block 32、max 2048 下自蒸馏。On-policy 训练:对 $x^0$ 采样 t,以概率 t 替换为 [MASK] 得 $x^{(m)}_t$;用 M_θ(no_grad) 在 $x^{(m)}_t$ 前向,在 mask 位采样得 $x^{(p)}_t$;再对两者各做有梯度前向,$L_{on-policy}=L_{mask}+L_{pred}$;在 8 H200、batch 8、lr 2e-6、cosine 下微调 2 epoch 得 DMax-Math/Coder。Block 初始化所有 h_j=e_mask。SPD 迭代:找最长连续高置信前缀 P 提升为 token,构造 h_j=π_j·e_yj+(1-π_j)·e_mask 范数归一化,直到收敛。Block 提交后进入下一块。

技术新颖性

技术新颖性集中在四点。第一,首次明确指出「词表均匀随机 token」是 UDLM 训练在并行解码下表现崩塌的根因(GSM8K 68.7% vs DMax 92.1%,23.4 个百分点差距),并用 on-policy rollout 给出可落地修复;第二,SPD 虽形式上像 SM、EvoToken 的软嵌入,但 DMax 把它定位为「保留不确定性先验」载体,直接与 on-policy 训练联动,并设计「连续前缀策略 + 归一化」两个工程细节,使其在 $\tau_{dec}=0$ 下仍保持 90% 精度;第三,Algorithm 1 引入块级收敛条件(连续一致性 + 置信度阈值),Table 4 证明 confidence 准则可让 GSM8K TPF 从 5.13→5.48、MBPP 从 5.16→5.86;第四,同时报告 TPF/TPS/Acc/AUP 四指标并给出 Figure 4 的 Acc-TPF Pareto 曲线,在 dLLM 加速领域是较系统的「精度-效率」综合评估框架。

**Figure 2**: Overview of the proposed On-Policy Uniform Training.

**Figure 3**: Overview of the Soft Parallel Decoding procedure in DMax.

实验结果

在 6 个 benchmark 上对比 4 类基线,核心发现分四层。第一层「并行度大幅提升 + 精度不降」:GSM8K TPF 2.04→5.48(↑168%)、精度 92.6%→92.1%;MATH500 TPF 2.58→5.94(↑130%)、精度 75.8%→75.4%;HumanEval TPF 4.38→7.36(↑68%)、精度 84.2%→83.5%;MBPP TPF 2.71→5.86(↑116%)、精度 80.6%→79.2%。AUP 全部领先(GSM8K +64%,MBPP +75%)。第二层 on-policy 是关键:Uniform Diffusion Training 全部断崖下降(GSM8K 68.7%、HumanEval 15.2%)。第三层 SPD 与 OPUT 协同:无 OPUT 直接 SPD 模型完全坍缩(0.0%);仅 OPUT 在 $\tau_{dec}=0$ 时只有 68.2%;OPUT+SPD 联合后 90.4%。第四层低并行度也涨点:精度普遍提升 0.7~3.0 个百分点,自修正对推理路径本身也有改善。

**Table 1**: Comparison with the original model and different baselines.

**Table 2**: Our proposed new paradigm also improves the model’s accuracy at low parallelism.

**Table 3**: Ablation on different training and inference strategies with different decoding parallelism.

**Table 4**: Ablation study on block-level convergence criteria.

**Figure 4**: Comparison of accuracy-TPF trade-off curves between original LLaDA-2.0-mini model and our method.

查看结构化数据

任务	指标	本文	基线	提升
GSM8K 数学推理	TPF (Tokens Per Forward, ↑) / Acc (↑) / AUP (↑)	TPF 5.48, TPS 1258, Acc 92.1%, AUP 557	LLaDA-2.0-mini TPF 2.04 / TPS 512 / Acc 92.6% / AUP 340;dParallel-SFT TPF 2.79 / Acc 92.3% / AUP 395;Uniform Diffusion Training TPF 2.26 / Acc 68.7% / AUP 0	TPF 提升 +168%,AUP 提升 +64%,Acc 仅下降 -0.5%;同时显著优于 dParallel-SFT (AUP +41%) 和 UDLM (Acc +23.4%)
MATH500 数学推理	TPF / TPS / Acc / AUP	TPF 5.94, TPS 1286, Acc 75.4%, AUP 507	LLaDA-2.0-mini 2.58 / 626 / 75.8% / 257;dParallel-SFT 3.42 / 823 / 75.8% / 310;UDLM 2.43 / 530 / 33.6% / 0	TPF +130%,AUP +97%,UDLM 在此任务上几乎完全崩溃(Acc 33.6%),证明 on-policy 必要性
Minerva-Algebra 数学推理	TPF / TPS / Acc / AUP	TPF 7.03, TPS 1492, Acc 91.5%, AUP 658	LLaDA-2.0-mini 3.01 / 755 / 91.4% / 363;dParallel-SFT 3.91 / 943 / 91.4% / 430;UDLM 2.55 / 551 / 42.7% / 0	TPF +134% (为 6 个任务中最大相对提升),AUP +81%,Acc 持平甚至略升
ASDIV 数学推理	TPF / TPS / Acc / AUP	TPF 5.62, TPS 1172, Acc 92.5%, AUP 556	LLaDA-2.0-mini 2.03 / 512 / 92.8% / 354;dParallel-SFT 2.72 / 663 / 93.0% / 459;UDLM 2.51 / 515 / 80.8% / 0	TPF +177%,AUP +57%,在 dParallel-SFT 已经把 Acc 推高到 93% 的情况下,DMax 仍能以 -0.3% Acc 换取 TPF 翻倍
HumanEval-Instruct 代码生成	TPF / TPS / Acc / AUP	TPF 7.36, TPS 1557, Acc 83.5%, AUP 637	LLaDA-2.0-mini 4.38 / 1044 / 84.2% / 369;dParallel-SFT 5.12 / 1229 / 76.8% / 394;UDLM 2.93 / 628 / 15.2% / 0	TPF +68%,TPS +49%,AUP +73%;UDLM 在代码任务上崩塌最严重(Acc 15.2%),凸显 on-policy 在代码生成上的关键作用
MBPP-Instruct 代码生成	TPF / TPS / Acc / AUP	TPF 5.86, TPS 1264, Acc 79.2%, AUP 482	LLaDA-2.0-mini 2.71 / 662 / 80.6% / 276;dParallel-SFT 3.66 / 880 / 74.7% / 273;UDLM 2.84 / 608 / 23.4% / 0	TPF +116%,AUP +75%;DMax 在低 TPF(2.71)时 Acc 80.6%,在 TPF 5.86 时仍维持 79.2%,几乎实现「精度无损加速 2 倍以上」
低并行度 Acc 提升 (Table 2)	Acc 提升幅度	GSM8K 93.4% (+0.8%),MATH500 78.0% (+2.2%),Minerva-Algebra 93.6% (+2.2%),ASDIV 93.5% (+0.7%),HumanEval 87.2% (+3.0%),MBPP 83.4% (+2.8%)	LLaDA-2.0-mini 92.6% / 75.8% / 91.4% / 92.8% / 84.2% / 80.6%	DMax 在 TPF 较低的设置下精度全面提升 0.7~3.0 个百分点,显示自修正机制对推理路径本身也有改善
推理吞吐 (Hardware Utilization)	TPS (Tokens Per Second)	在 2 张 H200 上 batch size 1,数学任务平均 TPS 1258~1492,代码任务 1264~1557	LLaDA-2.0-mini 512~755 (数学) / 662~1044 (代码)	实际推理吞吐 2 倍以上提升(GSM8K 512→1258 = +146%,MBPP 662→1264 = +91%),首次在 batch size 1 下让 dLLM 推理突破 1000 TPS

局限与改进

作者未单独设 Limitations 小节,通过仔细阅读可总结出六点。其一,DMax 精度虽接近原模型但仍略低 0.4~1.4 个百分点,在 LLaDA-2.0-mini 原本就接近 92% 的高基准上特别敏感;其二,DMax-Coder 在 MBPP 上精度 -1.4% 是 6 个任务中最大,提示自修正对代码长程逻辑链(变量名、签名、依赖)修正能力有限;其三,所有结果都是 zero-shot、batch size 1,block size 固定 32,未在更大 batch 或更长 block 下做评测;其四,训练数据 100% 来自 LLaDA-2.0-mini 自蒸馏,精度上限被 base 模型天花板锁死(MATH500 75.4%、HumanEval 83.5% 都未超过 base);其五,on-policy 训练在每次迭代都额外做一次前向,8 H200 上才能 batch 8,中小团队难复现;其六,评测集中在 6 个标准 benchmark,缺乏长文本、agentic reasoning、tool use、multilingual 等场景迁移性验证。

独立分析的弱点

独立审视后有以下弱点值得改进。第一,自修正只发生在单 block 内部,跨 block 按自回归顺序提交,前一块提交 token 错则后一块无法回头,改进方向是引入跨 block 二次校验或相邻 block 软参考;第二,on-policy 训练每次额外一次完整前向,实际成本是普通 SFT 的 1.5-2 倍,改进方向是用 REINFORCE-style 截断或仅在高 mask ratio 时执行 rollout;第三,SPD 混合嵌入依赖上一轮 top-1 置信度,若早期对正确 token 给出 <0.5 置信度,会延长收敛时间,改进方向是引入置信度衰减或熵正则;第四,代码任务精度仍下降 0.7-1.4 个百分点,改进方向是混入少量外部高质量代码做 cold-start,或把 OPUT 与 RLHF/DPO 结合;第五,作者只报告 batch size 1 的 TPS,缺 batch 4/8/16 下的延迟、显存、吞吐曲线;第六,Figure 4 显示极高 TPF(>7)时 DMax 也开始掉点,自修正存在次数上限,改进方向是研究多轮与少轮修正的收敛 trade-off。

未来方向

基于论文隐含方向有四条值得研究的方向。其一,长度泛化:本文固定生成长度 2048,未来可探索 block size 32 在 8K/32K/128K 长文本上是否依然有效,长上下文错误传播模式可能完全不同;其二,多模态扩展:把 OPUT+SPD 推广到图像/视频/音频多模态 dLLM(参考 LLaDA-V、Dream-V),研究软嵌入插值在跨模态 token 上的可行性;其三,与 RLHF/DPO 联训:on-policy rollout 用「模型自身采样」,若换成「人类偏好排序的采样」可在 OPUT 损失中加入偏好正则,让自修正不仅学「正确」还学「人类偏好」的正确;其四,理论分析:目前纯实验驱动,缺乏「为什么 on-policy 能消除 train-inference 错配」的收敛性理论分析,未来可借鉴 score matching 给出形式化证明。

复现评估

复现评估分四方面。开源情况:论文给出代码仓库 https://github.com/czg1225/DMax,作者声明代码会发布,但 DMax-Math/Coder 权重与训练脚本是否已公开需确认;基线 LLaDA-2.0-mini 已开源便于对比。数据可获得性:训练数据完全来自 GSM8K、PRM12K、Numina-Math、OpenThoughts、OpenCodeInstruct 公开子集,无需额外申请,按 0.95 阈值自蒸馏即可复现 0.7M+1.0M 样本。算力需求:训练需 8 H200、batch 8、2 epoch,数学约 175k 步、代码约 250k 步,单卡 H200 约 8-10+12-15 天;推理只需 2 H200 在 batch 1 下跑 1000+ TPS。复现难度:中等偏高,主要门槛在 on-policy rollout、SPD 混合嵌入 + 归一化 + 连续前缀扫描、块级收敛条件、dInFer 框架对接,具备 LLaDA 训练经验的团队 1-2 人月可完成核心复现。

DMax:面向扩散语言模型的激进并行解码 DMax: Aggressive Parallel Decoding for dLLMs