扩散语言模型中习得式去掩码策略的受控研究 What Makes a Good Unmasking Policy? A Controlled Study of Learned Unmasking in Diffusion Language Models
受控研究揭示去掩码策略设计规律，提出轨迹收缩监督，856K小策略显著加速dLLM解码
前置知识
掩码扩散语言模型（dLLM）
一类以离散扩散方式生成文本的大语言模型：从全掩码序列出发，每步对当前序列做一次骨干前向，得到每个掩码位置的预测分布，再把部分掩码替换为预测 token，迭代若干步直到全部解出。与自回归模型从左到右逐 token 生成不同，dLLM 可在一步内并行提交多个位置，代表模型有 LLaDA-8B、Dream-7B。
本文的全部问题都发生在 dLLM 的迭代解码循环中，策略、NFE、参考轨迹等核心概念都建立在这个解码过程之上。
去掩码策略（unmasking policy）
决定每个解码步提交哪些掩码位置的机制。训练无关版本直接用规则（如置信度超过阈值就提交）；学习版本是一个小神经网络，输入当前解码状态的特征，输出每个候选位置的承诺概率，再按阈值提交。由于策略在每个解码步都要评估一次，其自身开销会抵消省下的骨干前向。
这是本文的研究对象；论文提出的四个设计轴——教师标签、训练状态、策略输入、策略容量——全部围绕它展开。
NFE（number of function evaluations）
生成一条序列所需的骨干模型前向评估次数。一次前向可并行解出多个 token，所以 NFE 越小解码越快；逐个解码（one-by-one）的 NFE 等于生成长度 L（本文为 256）。NFE 是效率的代理指标，但不包含策略自身的计算开销。
精度—NFE 前沿是本文的主要评价方式，而盈亏平衡分析（约10%规则）回答的是多少 NFE 下降才真正转化为端到端加速。
Fast-dLLM 与置信度并行解码
Fast-dLLM 是一种训练无关的 dLLM 加速方法：给每个位置的 top-1 预测设一个置信度阈值（如 0.9），一步内把所有超过阈值的位置并行提交，并配合 KV 前缀缓存近似。它既是本文生成参考轨迹的教师，也是最重要的对比基线之一。
本文的训练标签完全建立在其解码路径之上，理解它的承诺规则才能明白终端一致与轨迹收缩两种标签的差异。
策略蒸馏（模仿学习）
让小策略模仿一个固定教师决策过程：先在训练数据上运行教师，得到状态—决策对，再让策略在这些状态上用交叉熵损失拟合教师的标签。与强化学习不同，它不需要在线交互，训练稳定且便宜，但策略的上限受教师行为约束。
本文把所有监督式去掩码策略统一视为对固定解码过程的蒸馏，受控研究正是围绕蒸馏什么、在哪些状态上蒸馏展开的。
前缀缓存（prefix caching）
一种 KV 缓存近似：解码中固定不变的前缀部分的注意力键值被缓存复用，后续每步只对变化部分做完整计算，从而大幅降低每步延迟（本文测得 2048 提示下降低 73%）。代价是注意力计算变为近似，可能轻微改变模型输出。
本文发现学习型策略对这种近似明显更敏感（精度最多降约 2 点），是重要局限；缓存同时改变延迟结构，影响策略开销的盈亏平衡点。
研究动机
掩码扩散语言模型(dLLM)通过迭代精化部分掩码序列、每步并行提交多个 token 来生成文本,效率高度依赖去掩码策略:保守策略需要很多次去噪前向(NFE),激进策略则因各位置在相同状态下独立预测而可能提交互相矛盾的 token,且已提交的 token 后续不再修正。近年出现多种学习型策略:Learn2PD 用约 2.1K 参数的过滤器从 oracle 正确性信号学习;TraceLock 用约 7.5M 参数的网络处理骨干隐状态;LUP 系列用强化学习训练 300K(置信度输入)或 300M(隐状态输入)的策略。问题在于这些方法同时在多个设计轴上不同——教师如何标注可提交、在哪些状态上训练、策略看什么输入、策略多大——无法归因哪个选择有效。例如 LUP-hidden 参数约为 LUP-conf 的 1000 倍,效果反而更差;是否该用隐状态无法单独判断,因为使用它的策略同时也改变了容量与监督。此外,策略每步都要前向,策略规模、KV 缓存与提示长度如何影响多少 NFE 下降才值得,始终没有被量化。
本文的目标是本文的目标是把学习型去掩码策略的设计空间拆开，逐一受控地回答三个问题。第一，监督信号应当如何构造：什么算一次安全提交（承诺目标），以及在哪些解码状态上向教师学习（训练状态）。第二，策略应当观察什么输入：置信度、top-2 概率间隔、骨干隐状态，还是它们的组合。第三，策略带来的 NFE 下降中，有多少会被策略自身的每步开销吃掉，即不同参数规模、不同缓存与提示长度下的盈亏平衡点在哪里。在找到答案后，作者把结论组合成一个单一策略 OnTrack，验证它在精度—NFE 前沿和端到端吞吐（tokens/s）上是否确实优于训练无关方法（Fast-dLLM、EB-Sampler）和已发表的学习型策略（Learn2PD），并考察跨任务（数学到代码）、跨骨干（LLaDA 到 Dream）和跨生成长度的泛化能力。
与已有工作不同的是，本文的独特切入是用蒸馏视角统一所有监督式策略：解码过程在训练提示上跑出参考轨迹，教师从中标注每个位置可否提交，策略模仿教师。这一视角把设计拆成教师标什么、在哪些状态模仿、策略看什么、策略多大四个正交维度，从而可以做单因素受控实验。在此框架下作者提出新候选：轨迹收缩（trajectory contraction）标签——只有当一次提交能直接到达参考轨迹上的某个后续状态时才算有效，替代与最终输出一致即可的终端一致标签；以及置信度+隐状态的组合输入。此外，此前工作几乎只报告 NFE，本文额外对每步延迟做了系统 profile（H200、batch 1、128–2048 提示长度、有无前缀缓存），首次给出策略参数不超过 7.5M 时约 10% 的 NFE 下降即可抵消开销的可操作阈值，把精度—NFE 前沿与真实加速连接起来。
核心方法
直觉上,某个位置的预测等于最终答案并不代表现在提交它就安全:提交会改变后续所有去噪步骤的输入,可能把模型带离原本的解码路径。作者据此把监督式去掩码策略视为对固定解码过程(教师)的蒸馏,逐因素受控比较。技术路线:以冻结的 LLaDA-8B-Instruct 为骨干,用 Fast-dLLM(置信度阈值 0.9、$L=256$、块长 32、温度 0)在 GSM8K+MATH 训练集上生成参考轨迹,仅保留答案正确的 9,169 条(共 686,052 个状态)。监督研究为 2×2 设计:承诺目标(终端一致 vs 轨迹收缩)× 训练状态(参考轨迹 vs 教师诱导);输入研究比较 $c_{t,i}$、$h_{t,i}$、$[c_{t,i},\Delta_{t,i}]$、$[c_{t,i},h_{t,i}]$,架构统一只改输入投影;开销研究在 H200 上对 331K/856K/7.5M/300M 策略 profile 每步延迟。所有比较用同一阈值网格扫描精度—NFE 前沿,以逐个置信度解码为参考精度(GSM8K 78.0%、MATH-500 38.8%),设 98%/95%/90% 目标线。
核心创新是轨迹收缩监督目标。已有教师(Learn2PD、TraceLock 共用)用终端一致标签 $y^{tok}_{t,i}=\mathbb{1}[\hat{x}_{t,i}=x^\star_i]$:预测与参考轨迹最终 token 相同即可提交。这在低 NFE 区域过于宽松:正确 token 若在有错预测时被提前提交,会带偏离参考轨迹。轨迹收缩改为:设 $A_{t\to u}=M(s_t)\setminus M(s_u)$ 为从 $s_t$ 到未来状态 $s_u$ 间被提交的位置集合,仅当 $C(s_t, A_{t\to u})=s_u$(提交后恰好落在参考轨迹上)时有效;取最大可行 $u^\star_t$,目标为 $y^{tc}_{t,i}=\mathbb{1}[i\in A_{t\to u^\star_t}]$;有效性对 $u$ 单调,故不需额外骨干前向。本质区别:标签编码了提交对后续状态的影响,即路径一致性而非仅终点一致性。配套发现同样关键:隐状态单独弱于置信度但与之互补;策略不超过 7.5M 参数时每步仅增约 1.5ms(不超过 10%),约 10% NFE 下降即可盈利。
方法步骤详情
训练分四步。1) 轨迹采集:用 Fast-dLLM(阈值 0.9)在 GSM8K+MATH 训练集上解码,仅保留答案正确的轨迹,得 9,169 条、686,052 个状态。2) 目标构造:对每个状态 $s_t$ 取最大有效收缩 $A^\star_t$,标签 $y_{t,i}=\mathbb{1}[i\in A^\star_t]$。3) 策略训练:输入置信度 $c_{t,i}=\max_v p_\theta(x_i=v\mid s_t)$ 与末层隐状态 $h_{t,i}$,投影到 $d=128$ 维相加为 $z_{t,i}=W_c c_{t,i}+W_h h_{t,i}$;单层 DiT 以 $m_t,t$ 为条件输出 $q_{t,i}=\sigma(f_\phi(z_t;m_t,t))$,优化交叉熵 $\mathcal{L}=-\sum[y\log q+(1-y)\log(1-q)]$;AdamW,1 epoch,单卡 H200 约 3.25 小时,总参数 855,713。4) 推理:块内提交 $q_{t,i}>\gamma$ 的位置并填骨干 top-1 预测,若无则提交策略分最高者。
技术新颖性
技术新颖性体现在三点。其一，方法学上，这是首个对学习型去掩码策略做单因素受控归因的工作：统一架构（同一 DiT 主干，只改输入投影）、统一数据与训练日程、统一阈值网格，把现有方法纠缠的四个设计轴（教师标签、训练状态、输入、容量）拆开，得出终端一致标签在低 NFE 区域过宽（28.7% 的正标签违反轨迹收缩，56.0% 的参考状态含至少一处违反）这类此前无法得到的结论。其二，标签设计上，轨迹收缩把哪些提交是安全的形式化为能否一步到达参考轨迹上的后续状态，既保留了并行提交的跳跃能力（允许跳过中间状态），又施加了路径一致性约束；利用有效性对 $u$ 的单调性，标签构造零额外骨干开销，可与任意参考轨迹即插即用。其三，系统层面，首次量化策略自身开销的盈亏平衡：对 331K/856K/7.5M/300M 四种规模、128–2048 提示长度、有无前缀缓存做每步延迟 profile，给出不超过 7.5M 参数约 10% NFE 下降即回本、300M 约需 28% 的经验法则，弥补了该领域只看 NFE 的评估盲区。
实验结果
受控研究给出三个发现。发现 1(监督):轨迹收缩在两种训练状态下都改进前沿,GSM8K 上贡献大部分增益;域外 MBPP 上参考轨迹带来额外增益。证据:28.7% 的终端一致正标签违反收缩约束;教师诱导仅保留参考状态的 34.0%/25.4%。发现 2(输入):隐状态单独弱于置信度(105.3 vs 52.9 NFE),组合以最少 NFE 达 98%/95% 目标(42.1);加间隔无一一致增益。发现 3(开销):不超过 7.5M 的策略每步开销至多 10%,约 10% NFE 下降即回本;300M 最高 +39%,需约 28%。主实验:GSM8K 全目标最优,98% 目标对三基线加速 1.56×/1.45×/1.59×(仅一处 CI 含 1);MATH-500 95% 目标最少 NFE 但 98% 饱和;HumanEval 90% 目标最少 NFE;MBPP 95%/98% 最少 NFE、90% 落后。迁移:Dream-7B 上优势大体保留(HumanEval 59.8% vs 55.5%);长度迁移成功,但域+长度双重偏移下落后;缓存使精度最多降 2 点。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| GSM8K 数学推理（LLaDA-8B-Instruct，98% 精度目标） | 达到目标的最少平均 NFE | 42.1 | Fast-dLLM 65.6 / Learn2PD 61.2 / EB-Sampler 66.8 | NFE 加速 1.56×/1.45×/1.59×（paired bootstrap 95% CI，除对 98% 目标的 EB-Sampler 外均排除 1） |
| GSM8K 98% 目标，零样本、无前缀缓存 | 吞吐 TPS（tokens/s，单卡 H200） | 331 | Fast-dLLM 232 / Learn2PD 242 / EB-Sampler 221 | 约 +43%/+37%/+50%，所有精度目标下均为最高 TPS |
| MBPP 代码生成（域外迁移，98% 精度目标） | 达到目标的最少平均 NFE | 48.6 | Fast-dLLM 60.8 / EB-Sampler 65.1 | NFE 加速约 1.25×/1.34× |
| GSM8K（策略输入受控研究，98% 目标） | 最少平均 NFE | 42.1（置信度+隐状态） | 仅置信度 52.9 / 仅隐状态 105.3 / 置信度+间隔约 43.6 以下无一一致更优 | 较仅置信度省约 20% NFE；隐状态单独使用需翻倍 NFE |
| MBPP（监督设计 2×2 受控研究，98% 目标） | 最少平均 NFE | 48.6（参考轨迹+轨迹收缩） | 教师诱导+收缩 64.8 / 参考轨迹+终端一致 192.4 | 较次优监督设计省约 25% NFE，验证轨迹收缩与参考轨迹的互补 |
| GSM8K 盈亏平衡分析 | 抵消策略每步开销所需的 NFE 下降 | 约 10%（331K–7.5M 策略，每步开销 +2%~+10%） | 300M 策略开销最高 +39%，需约 28% NFE 下降 | 确立轻量策略（≤7.5M）的现实可用性阈值 |
| Dream-7B-Instruct，HumanEval（跨骨干迁移） | 精度 / NFE | 59.8% @ 84.1 NFE，超过 one-by-one 参考精度 55.5% | Fast-dLLM 53.7% @ 83.7 / EB-Sampler 53.7% @ 85.7 | 唯一在最少 NFE 下超过参考精度的方法 |
局限与改进
作者承认的局限有三。一是训练范围:策略只在数学数据、单一生成长度 $L=256$、固定块长 32 上训练,可迁移到代码(同长度)或其他长度(同数学域),但域+长度双重偏移下优势消失,落后于 Fast-dLLM 和 EB-Sampler。二是高精度区域:MATH-500(LLaDA)和 GSM8K(Dream)上即使很保守的阈值,精度也饱和于参考之下,而训练无关方法可用更多 NFE 超过参考,故本方法最适合中等精度目标。三是对骨干特征的敏感性:前缀缓存使其精度最多降约 2 个百分点,置信度基线几乎不受影响;即使每步只解一个 token 也有下降,说明问题在于提交哪些 token,推测是隐状态被缓存扰动。我的补充观察:教师天花板明显,标签来自 Fast-dLLM 阈值 0.9 的轨迹且只保留答案正确的约 61%,错误轨迹信息被丢弃,策略难以超越教师;LUP 与 TraceLock 无公开 checkpoint 未被直接比较;统计功效有限,HumanEval 仅 164 题、MBPP 257 题,bootstrap 区间很宽,部分优势声明并不显著。
独立分析的弱点
独立分析可见四个弱点。其一,监督天花板:轨迹收缩标签继承自单一教师(Fast-dLLM 阈值 0.9),策略最多模仿教师的高效行为,无法学会更聪明或更保守的提交,98% 目标的精度饱和正是教师轨迹缺乏更激进安全跳跃的证据。改进:混合多阈值教师构造目标,或用最终答案奖励做 RL 微调。其二,隐状态依赖带来工程脆弱性:任何改变隐状态分布的近似推理(KV 缓存、量化、投机解码)都可能使策略失效,前缀缓存已造成约 2 个百分点损失,对追求吞吐的部署是实际障碍。改进:缓存感知训练(对隐状态加扰动)、归一化与鲁棒化,或自动回退纯置信度模式。其三,泛化维度单一:换骨干(LLaDA 到 Dream)需重新生成轨迹并重训策略(790K vs 856K),域与长度需混合训练才能覆盖。改进:骨干无关特征(相对置信度、熵)、多骨干联合训练、长度条件化。其四,结论的外部效度:约 10% NFE 回本是在 H200、batch 1、bfloat16 的 8B/7B 模型上测得,更大 batch、更小骨干或不同硬件会改变骨干与策略的延迟比,阈值需重测;小数据集上的前沿差异部分落在噪声内,应普遍报告置信区间。
未来方向
作者明确提出两个方向：一是跨域、跨生成长度的混合训练，以消除域+长度双重偏移下的失效；二是让学习型去掩码策略对近似推理（前缀缓存等）鲁棒；并指出高精度区域的饱和可能需要覆盖更保守承诺行为的监督。基于本文成果还可延伸：将轨迹收缩标签与强化学习结合，用最终答案奖励微调蒸馏策略以突破教师天花板；把标签构造扩展到温度采样的随机轨迹（本文全部温度 0），研究随机解码下的提交安全性；把延迟盈亏平衡分析发展为通用评估协议，联合优化策略容量与 NFE；应用于更大骨干（如 30B 级 dLLM）和多模态扩散模型，检验隐状态投影约 52 万参数的开销在更大模型上是否仍可忽略；探索自适应块长与策略的协同设计，以及在服务端 batch 推理场景下重新 profile 开销与吞吐。
复现评估
复现条件较好。数据与模型全部公开:LLaDA-8B-Instruct、Dream-7B-Instruct 为公开权重,GSM8K/MATH-500/HumanEval/MBPP 均公开,划分与提示模板在附录 C.1 说明;论文给出代码链接,实现基于 LUP(Jazbec et al., 2026)公开代码库修改。算力门槛低:策略仅 856K 参数,训练为单张 H200 约 3.25 小时(AdamW、lr $3\times10^{-5}$、1 epoch、batch 16),骨干冻结无需训练;主要一次性成本是用 Fast-dLLM 对 14,958 条训练提示解码生成参考轨迹,8B 模型单卡可完成。实现细节充分:架构(单层 DiT、$d=128$、RoPE)、超参(含 seed 123)、阈值网格、吞吐测量协议均有记录,附录 D 给出全部操作点数值表(Tables 13–26)便于逐点核对。需约 40GB 以上显存的 GPU 做 8B 推理。综合:复现难度中低,主要不确定性在参考轨迹生成与硬件差异导致的吞吐漂移,精度—NFE 前沿相对稳定。
论文图表
前缀缓存实验图：零样本与四样本提示下，无缓存（实线）与有缓存（虚线）解码的精度—吞吐曲线。OnTrack 对缓存更敏感（精度最多降约 2 个百分点），零样本 98% 目标被 Fast-dLLM 反超，但四样本 98% 目标 OnTrack 仍最快（423 vs 349 tokens/s）。
揭示了本文最重要的实用局限：学习型策略消费隐状态，因而对近似推理更脆弱；同时验证了在常用四样本+缓存配置下吞吐优势依然成立。