💬 NLP

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

本文介绍了 D3IM，一种能够实现掩码扩散语言模型中直接 Token 修订的无参数采样器，以及 SCOPE，一种用于减轻由此产生的保留偏差（preservation bias）的后训练方法，两者共同在推理和编程基准测试上实现了显著的性能提升。

原作者： Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu, Yue Dong, Greg Ver Steeg

发布于 2026-06-02

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu, Yue Dong, Greg Ver Steeg

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

核心问题：“冻结的错误”

想象你正在和一个非常聪明但有点紧张的机器人助手一起写故事。每当机器人猜出一个词时，它就会把这个词写在纸上。

在这些机器人工作的标准方式中（被称为掩码扩散语言模型，Masked Diffusion Language Models），一旦机器人写下一个词并觉得已经足够“自信”了，它就会把这个词冻结在原处。它会在词上面贴上一层胶带。即使机器人后来意识到：“等等，这个词放在这个句子里并不通顺，”它也无法更改，因为胶带在那里。它只能填充那些尚未被写下的“空白”部分。

对于数学或编程这类棘手的任务，这是一个巨大的问题。如果机器人在早期写错了一个数字（比如把“2”写成了“5”），它就会陷入困境。它会试图围绕这个错误的数字构建剩余的答案，从而导致最终结果完全错误。

解决方案：两部分的改进

作者 Longxuan Yu 及其同事意识到，机器人其实具备改变主意的能力，只是游戏规则阻止了它这样做。他们提出了一个由两部分组成的解决方案来解决这个问题：一种新的游戏方式（D3IM）和一种专门针对机器人的训练练习（SCOPE）。

第一部分：新的游戏规则 (D3IM)

类比： 想象一场大家都在猜句子中下一个词的“烫手山芋”游戏。

旧规则： 一旦你喊出一个词，你就必须一直拿着它直到结束。如果你喊错了词，你就只能认栽。
D3IM 规则（“白纸重来”）： 在游戏的每一个步骤中，机器人都会观察句子中的每一个词，甚至是那些它已经喊出来的词。它会问自己：“我仍然认为这是最好的词吗？”
- 如果机器人依然很有信心，这个词就会保留。
- 如果机器人觉得：“实际上，这不对，”它会擦掉这个词（把它变回空白）或者立即用一个更好的词替换它。
- 它不需要请求许可或使用特殊工具；它只是根据当前的信心水平重新评估一切。

难点在于： 机器人玩这个游戏表现得很好，前提是它知道自己错了。但在最初的训练中，机器人从未被教导去承认自己的错误。它被训练去信任自己的第一次猜测，即使那个猜测是错误的。所以，当你让它玩这个新的“白纸重来”游戏时，它会不断冻结那些错误的词，因为它认为它们是正确的。

第二部分：专门的训练 (SCOPE)

类比： 这就像是一场“带有错误的排练”。
为了纠正机器人过度信任自身错误的坏习惯，作者创建了一种名为 SCOPE（基于预测误差的自我调节训练，Self-Conditioned On Prediction Errors）的训练方法。

它是如何运作的： 在训练期间，机器人被要求写一个句子，但随后它会被迫观察自己的作品，并假装自己犯了错。
机器人写下一个句子，然后训练员说：“好了，我要隐藏掉你的一些词。现在，请再次尝试猜出它们。”
转折点： 训练员会专门挑选那些机器人猜错、但却表现出极高置信度的词。机器人必须看着自己错误的答案，意识到自己错了，然后用正确的词替换它。
这教会了机器人一个至关重要的教训：“仅仅因为我说得很有信心，并不意味着我是对的。我需要愿意改变主意。”

结果：一个协作的团队

当你将新的游戏规则（D3IM）与专门的训练（SCOPE）结合在一起时，机器人就变成了一个具有自我修正能力的天才。

没有经过训练： 如果你只给机器人新的规则而不进行专门训练，它的表现反而会变差。它试图改变词汇，但却不断重复同样的错误，因为它并不信任自己发现错误的能力。
经过训练后： 机器人学会了识别自己的“冻结错误”。它现在可以观察到一个数学题中的错误数字，说：“不，这不对，”然后将其替换为正确的数字，而这一切都是在句子构建的过程中完成的。

现实世界的结果

论文在名为 LLaDA-8B（一种智能语言模型）的模型上测试了这些任务，涵盖了困难的任务：

数学 (GSM8K & MATH)： 分数显著提升。例如，在一场高难度的数学测试中，准确率从 55.3% 跳升至 68.3%。
编程 (HumanEval & MBPP)： 编写可用代码的能力大幅提升，在其中一项测试中从 14.0% 跃升至 29.3%。

核心启示

主要的发现是：拥有改变主意的能力是不够的；你必须经过训练，才知道何时该改变主意。

作者展示了通过教导模型识别其自信的错误（SCOPE），并赋予其能够立即替换这些错误的机制（D3IM），该模型处理复杂推理问题的能力比以前强得多。这不仅仅是让机器人变得更聪明，而是教会它不要对早期的错误如此固执。

技术摘要：修订而非冻结：用于自纠错掩码扩散语言模型的采样器匹配训练

1. 问题陈述

掩码扩散语言模型（MDLM），如 LLaDA，通过迭代地对掩码序列进行去噪来生成文本，并在每一步重新预测所有位置。这种架构本质上具备一种“修订能力”：即使在某个 token 变得可见之后，去噪器在后续步骤中理论上仍可以为该位置分配不同的预测。

然而，标准的 MDLM 解码策略采用的是单调去掩码（monotonic unmasking），一旦提交 token 就会将其“冻结”。一旦 token 被揭示，它就被视为干净的 token，随后的步骤仅更新掩码位置。这导致了两个关键问题：

永久性早期错误： 生成过程早期犯下的错误（例如数学题中的错误数字或代码中的错误运算符）会变成永久性的，严重影响在重推理任务上的表现。
修订差距（The Revision Gap）： 现有的修订方法要么依赖启发式机制（置信度阈值、验证步骤），要么需要将已提交的 token 重新掩码回 [MASK] 再进行重新预测（例如 ReMDM）。目前仍缺乏一种能够实现直接“可见到可见”（visible-to-visible）修订，且无需辅助模块或重新掩码步骤的有原则的采样器。

此外，作者识别了一个被称为**保持偏差（preservation bias）**的模型端障碍：即使训练目标并不强制执行 token 保留，使用标准掩码目标训练的模型也倾向于复现其自身生成的错误已提交 token，而不是在暴露于自我生成上下文时对其进行纠正。

2. 方法论

本文提出了一个由两部分组成的解决方案：D3IM（一种新采样器）和 SCOPE（一种训练程序），旨在协同工作。

2.1 D3IM：全清重掩码（Clean-Slate Remasking）

D3IM 是一种作为校正器风格反向更新衍生的无参数采样器。不同于保留已提交 token 的标准采样器（设置保留权重 $c_t > 0$ ），D3IM 将 $c_t$ 设置为 0。

机制： 在每个去噪步骤中，模型预测所有活跃位置。D3IM 根据置信度对这些预测进行排序，并仅保留前 $K$ 个位置。
修订通道：
- Token-to-Token（Token 到 Token）： 原先已提交的 token 可以被新的、更高置信度的预测直接覆盖，而无需经过 [MASK] 状态。
- Token-to-Mask（Token 到掩码）： 置信度较低的已提交 token 会被降级回 [MASK]，以便在后续步骤中重新预测。
约束： 下一个状态完全由当前的预测决定，而不是通过复制之前的可见 token 来确定。这要求模型的置信度排序在自我生成的上下文中是可靠的。

2.2 SCOPE：基于预测误差的自我调节

直接将 D3IM 应用于基础 MDLM 会导致性能下降，因为模型缺乏纠正自身高置信度错误的能力（保持偏差）。SCOPE 是一种轻量级的后训练程序，旨在使模型的置信度行为与 D3IM 的采样动态相一致。

流程：
1. 模型在部分掩码的输入上执行前向传播（自我调节/自我条件化）。
2. 它对掩码位置进行采样预测。
3. 这些预测中置信度最高的 $\rho$ 比例被“提交”为可见 token，其余则保持掩码状态。至关重要的是，这一过程会让模型暴露在其自身的高置信度错误之下。
4. 然后，模型被训练以恢复包括已提交（可能错误的）token 和剩余掩码位置在内的真实序列。
目标： 训练模型将可见的自我生成 token 视为假设而非固定上下文，从而减少保持偏差，并使置信度排序能够为 D3IM 的前 $K$ 个选择提供可靠依据。

3. 核心贡献

D3IM 采样器： 一种无参数的直接修订采样器，实现了无需辅助模块或重新掩码的直接 visible-to-visible token 更新，它是作为预训练 MDLM 的有效校正更新而衍生的。
识别保持偏差： 作者证明了标准 MDLM 训练会导致一种行为倾向，即模型即使在训练目标允许纠正的情况下，也会复现其自身错误的可见预测。
SCOPE 训练： 一种模拟 D3IM 采样过程的轻量级后训练程序。它有效地降低了保持偏差，使 D3IM 能够按预期发挥作用。
采样器-模型对齐： 本研究表明，修订能力并非通用的；训练目标必须与推理采样器的决策规则（例如 D3IM 的全局 top- $K$ 规则）专门匹配才能生效。

4. 实验结果

实验在 LLaDA-8B-Instruct 上针对四个基准测试进行：GSM8K（数学）、MATH-500（数学）、HumanEval（代码）和 MBPP（代码）。

性能提升： 在 64 步去噪时，SCOPE + D3IM 的组合显著优于使用标准去掩码的原始 LLaDA：
- GSM8K： +13.0% (68.3% vs. 55.3%)
- MATH-500： +4.8% (23.6% vs. 18.8%)
- HumanEval： +15.3% (29.3% vs. 14.0%)
- MBPP： +10.4% (30.8% vs. 20.4%)
缩放效应： 在数学和代码任务上，性能随去噪步数的增加而提升（例如，GSM8K 在 256 步时达到 73.8%）。
对齐的必要性：
- 未使用 SCOPE 的 D3IM： 性能显著下降（例如，GSM8K 降至 48.0%），证实了未经适配的模型无法可靠地决定哪些 token 需要修订。
- 未使用 D3IM 的 SCOPE： 使用 SCOPE 训练的模型配合标准单调去掩码，结果褒贬不一，表明训练收益是针对 D3IM 采样器的修订机制特有的。
机制诊断：
- 解决保持偏差： 在一个模型提交错误 token 的压力测试中，原始模型 70.1% 的时间会保留该错误。SCOPE 将此降低至 17.0%，并能以 50.2% 的概率恢复真实序列。
- 轨迹分析： 未使用 SCOPE 时，D3IM 的修正经常导致“内容-到-空白符振荡”（在 token 与换行符之间切换）。使用 SCOPE 后，修正会针对特定的语义错误（如单位转换、数字修正）并快速稳定下来。

5. 重要性与主张

论文认为，MDLM 自纠错推理的瓶颈不在于能否重新访问 token，而在于在自我生成上下文中置信度排序的可靠性。

有原则的修订： D3IM 提供了一种有理论依据的方法进行直接 token 修订，超越了启发式阈值或间接重新掩码。
训练-推理对齐： 本研究强调，先进的推理策略（如全局 top- $K$ 重排序）需要特定的训练目标（SCOPE）才能发挥作用。标准的持续预训练或通用的自我纠错方法并不会自动迁移到这些新的采样动态中。
效率： 所提出的解决方案计算效率很高，仅需要轻量级的 LoRA 微调（在单张 A100 上约 20 分钟，500 步），且在推理过程中不需要架构更改或辅助模块。

作者总结道，虽然 D3IM+SCOPE 显著改善了推理和代码生成，但它仅限于 token 级别的纠正，无法在初始轨迹从根本上存在缺陷时重构整个推理路径。未来的工作可能会探索利用丰富的修订轨迹进行过程奖励建模或强化学习。