🤖 machine learning

Fine-Tuning Masked Diffusion for Provable Self-Correction

本文介绍了 PRISM，这是一种轻量级且与模型无关的方法，通过在推理过程中计算每个 token 的质量分数来检测并修正低质量 token，从而为预训练的掩码扩散模型实现可证明的自修正，且无需进行架构更改、强化学习或依赖外部验证器。

原作者： Jaeyeon Kim, Seunggeun Kim, Taekyun Lee, David Z. Pan, Hyeji Kim, Sham Kakade, Sitan Chen

发布于 2026-05-26

📖 1 分钟阅读☕ 轻松阅读

原作者： Jaeyeon Kim, Seunggeun Kim, Taekyun Lee, David Z. Pan, Hyeji Kim, Sham Kakade, Sitan Chen

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

想象一下，你正在尝试解决一个谜题，比如数独或撰写故事，但你必须一次性填满所有空白，而不是一次填一个词。这就是**掩码扩散模型（Masked Diffusion Model, MDM）**的工作原理。它从一个完全空白的页面（或一个由空方格组成的网格）开始，并试图同时猜测每个位置的内容。

问题在于？因为它一次性猜测所有内容，所以有时会出错。它可能会在应该填"2"的地方填上"7"，或者写出一句不符合语法的句子。在过去，一旦 AI 犯错，它就会被困住。它无法回头说“等等，这看起来不对”，然后进行修正。

这篇论文介绍了一种名为PRISM（用于掩码扩散推理时自我修正的即插即用重掩码）的新方法。可以将 PRISM 想象成给 AI 配备了一双额外的眼睛和一支红笔，让它在工作时就能使用。

以下是其工作原理，分解为简单的概念：

1. 问题：“盲目猜测”

想象一下你正在填写填字游戏，但你必须在完全相同的时刻猜出每一个单词，而不查看其他单词的线索。你可能会为一个 3 个字母的单词猜出"CAT"，但后来发现它与交叉单词"DOG"不匹配。在标准的 MDM 中，一旦你写下"CAT"，它就会留在那里。模型不知道如何检查"CAT"在整个谜题的上下文中是否是一个好的猜测。

2. 解决方案：“质量评分”

PRISM 教导 AI 为其写下的每一个单词或数字分配一个质量评分。

高分：“这个词与句子的其余部分完美契合。”
低分：“这个词在这里感觉很奇怪。它可能是错的。”

论文声称，PRISM 能够以可证明的正确性学习给出这些评分。它不仅仅是猜测；它会根据周围的单词，通过数学计算来评估特定单词正确的可能性。

3. 魔法技巧：“重掩码”

一旦 AI 获得了质量评分，它就会使用一种称为重掩码的巧妙技巧。

如果一个单词的评分很低（它很可能是错的），AI 就会将其擦除（重新在该位置加上“掩码”）。
然后，它尝试为该位置猜测一个新单词。
如果一个单词的评分很高，它就保持原样。

这发生在生成过程中。AI 不仅仅是在写作；它是在实时地写作、检查、擦除和重写。这就像一位作家打出一句话，意识到某个词错了，将其删除，然后打出一个更好的词，然后再继续下一句话。

4. 为何这很特别（“即插即用”方面）

通常，要教会 AI 自行修正错误，你必须从头重建整个大脑，或者使用一个独立的、重型“批评家”模型来检查工作。

PRISM 轻量级：它就像给现有 AI 添加了一个小型“适配器”或插件。你不需要从头开始重新训练整个模型。
高效：论文表明，PRISM 可以使用非常少的额外数据学习自我修正（远少于训练原始 AI 所需的数据）。
通用性强：作者在以下领域测试了该方法：
- 数独：修正网格中的错误数字。
- 文本：写出更好、更自然的句子。
- 代码：修复 Python 编程错误（如缺少括号或逻辑错误）。

5. 结果

论文显示，当他们使用 PRISM 时：

AI 犯的错误更少。
它能够自行修正错误，从而在编程和谜题方面获得更好的最终结果。
即使对于非常大的模型（例如一个拥有 80 亿参数的编程助手），它也能发挥作用，而无需大量的额外计算能力。

简而言之

PRISM 是针对特定类型 AI 的一种简单而智能的升级。它赋予 AI 审视自身工作的能力，说出“这看起来不对”，擦除错误，然后重试——所有这些都在它生成答案的同时进行。它将“一次性猜测”转变为“起草与编辑”的过程，使 AI 在解决谜题、撰写文本和编程方面更加可靠。

技术摘要：用于可证明自校正的掩码扩散微调

问题陈述

掩码扩散模型（MDMs）已成为离散空间（如文本、代码、数独）生成建模中自回归模型的一种有前景的替代方案。与按顺序生成 token 的自回归模型不同，MDMs 从一个完全被掩码的序列开始，并以任意顺序逐步取消掩码 token。尽管这种并行、少步的推理提供了效率和灵活性，但它引入了一个关键限制：标准 MDM 缺乏自校正能力。

在标准 MDM 推理中，一旦一个 token 被取消掩码，它就会保持固定。这种刚性使得 MDM 容易受到跨位置的依赖错误的影响，特别是在使用少量推理步骤时。此前尝试解决此问题的方法面临显著缺点：

不精确的代理：无需训练的方法（如 ReMDM、DFM）依赖于随机重新掩码或源自未训练以表示 token 质量的 logits 的分数，导致校正效率低下或不准确。
架构彻底改造：那些确实定义了有效自校正目标的方法通常需要改变 MDM 架构（例如空心 Transformer）或预训练目标，限制了它们对现有预训练模型的适用性。

因此，需要一种方法，能够在不改变核心架构的情况下，为任何预训练 MDM赋予自校正能力，同时提供“token 质量”的理论基础定义。

方法：PRISM

作者提出了PRISM（用于掩码扩散推理时自校正的即插即用重新掩码），这是一个轻量级、与模型无关的微调框架。

1. 理论基础：逐 Token 质量

PRISM 为自校正定义了一个严格的目标：逐 token 质量分数。对于部分被掩码的序列 $y$ ，位置 $i$ 处的 token $y_i$ 的质量定义为：在 $y_i$ 被掩码的序列其余部分条件下，该 token 的似然度：
$g^*_i(y) := p(x_i = y_i \mid y \oplus m_i)$
此处， $y \oplus m_i$ 表示将序列 $y$ 中的第 $i$ 个 token 替换为掩码 token $m$ 后的序列。如果 $y_i$ 与上下文一致，则该概率较高；如果是错误，则该概率较低。

2. PRISM 损失函数

为了学习一个近似 $g^*$ 的模型 $g_\phi$ ，PRISM 引入了一种新颖的损失函数。挑战在于， $g_\phi$ 必须预测输入 $y$ 中已经取消掩码的 token $y_i$ 的质量，然而目标却是定义在被掩码上下文 $y \oplus m_i$ 上的后验概率。

PRISM 通过特定的数据构建和损失公式解决了这一问题：

数据构建：采样一个真实序列 $x$ 和一个掩码版本 $z$ 。使用预训练 MDM $f_\theta$ 对 $z$ 中的子集 token 进行取消掩码，以创建新序列 $y$ 。
损失函数：训练模型 $g_\phi$ 使用二元交叉熵（BCE）预测二元标签 $1[x_i = y_i]$ （即取消掩码的 token 是否与真实值匹配）：
$\mathcal{L}(\phi) := \mathbb{E}_{x,z,(y,i)} \left[ \text{BCE}(1[x_i = y_i], g_\phi^i(y)) \right]$
理论上，无论用于生成 $y$ 的预训练 MDM $f_\theta$ 的质量如何，该损失的唯一最小化器恰好就是所需的逐 token 质量 $p(x_i = y_i \mid y \oplus m_i)$ 。

3. 实现细节

即插即用：PRISM 将轻量级辅助适配器（“质量头”）附加到任何预训练 MDM 上。骨干网络保持共享，一个头输出取消掩码后验（原始 MDM 功能），新头输出标量质量分数。
正则化：为了保留 MDM 原始的取消掩码能力，在微调过程中将标准 MDM 训练损失作为正则化项加入。
效率：该框架允许通过同时取消掩码多个 token，从 $f_\theta$ 的单次前向传播中生成多个训练对 $(y, i)$ ，从而提高了数据效率。

4. 推理过程

在推理过程中，微调后的 MDM 在每一步同时执行两个操作：

取消掩码：基于标准 MDM 后验概率，选择一组被掩码的 token 进行取消掩码。
重新掩码：选择一组已取消掩码且质量分数最低（由 $g_\theta$ 预测）的 token 并重新掩码。
这些步骤允许模型在序列最终确定之前迭代地校正错误。

主要贡献

理论保证：PRISM 提供了首个针对 MDM 中逐 token 质量的可证明学习目标，无需生成完整干净序列或使用验证器。它证明了损失最小化器收敛于给定掩码上下文的 token 真实条件似然。
与模型无关的设计：与 prior 工作不同，PRISM 不需要架构更改或从头开始重新训练。它作为“即插即用”的微调层，适用于任何现有的预训练 MDM。
效率：该方法以显著更少的训练样本（比预训练少几个数量级）和计算资源学习自校正能力。

实验结果

作者在三个领域和规模上评估了 PRISM：

数独（30M 规模 MDM）：
- PRISM 在微调过程中取得了逐步提高的成功率，在几个 epoch 内超越了 ReMDM 和 ReMDM-conf 等基线。
- 可视化证实，PRISM 正确地为网格中的错误单元格分配了低质量分数，实现了有效的错误检测。
无条件文本生成（170M 规模 MDM）：
- 在 OpenWebText 上使用生成困惑度（Generative Perplexity）和 MAUVE 进行评估。
- PRISM 优于基线，特别是在采样步骤较少（例如 64–256 步）的情况下，此时依赖错误最为普遍。
- 校准研究表明，预测的质量分数与 token 的经验似然度紧密跟踪。
代码生成（8B 规模 MDM - LLaDA）：
- 应用于 LLaDA-8B-Instruct，微调时间少于 500 GPU 小时。
- 在 MBPP 和 HumanEval 基准测试中，PRISM 在各种采样步骤下均优于普通 MDM 和其他重新掩码基线。
- 机制分析显示，PRISM 成功校正了生成代码中的语法错误（例如不完整的 elif 语句）和逻辑错误（例如分隔符放置）。
- 可迁移性：在一个 MDM 骨干上训练的 PRISM 微调模型，可以作为不同预训练 MDM 骨干（在相同分布上运行）的自校正器，带来显著的性能提升（例如 HumanEval 上 +10.9%）。

意义与主张

该论文将 PRISM 定位为构建通过校正和重排序组合离散序列的生成模型的一步，模仿人类般的生成过程。

可证明的自校正：作者强调，PRISM 的独特之处在于定义和学习了精确、可测试的逐 token 质量目标，而无需依赖强化学习（需要完整序列生成）或外部验证器。
可扩展性：该方法证明了即使在大规模模型（80 亿参数）上，也可以通过微调高效地获得自校正能力，且计算需求极低。
局限性：作者谦逊地承认，由于质量分数基于逐位置后验边缘分布，PRISM 可能难以处理全局推理错误（例如跨越多个 token 的结构逻辑错误），这些错误未被局部 token 似然度所捕捉。然而，他们指出，这项工作为未来的框架在离散扩散的高效训练范式内解决全局错误打开了大门。