扩散模型也能"边写边改"：让掩码模型学会自己回头改错

你有没有观察过自己改一份草稿的过程？

正常人不会把整段话删掉重写一遍。你会盯着某个用词不对的地方，划掉它，换一个，然后顺带把后面跟着错的半句也修了。改的是局部，留的是大部分。

但现在的大模型不是这么干的。自回归（AR）模型做推理靠思维链、靠反思，可一旦发现前面写错了，它的"修正"本质上还是从头顺序再生成一遍——哪怕你只想改一个词。更糟的是，那些错误的中间内容会一直赖在上下文里，占着容量，污染后面的推理。有篇论文标题说得很直接："LLMs get lost in multi-turn conversation"，模型在多轮长对话里会"迷路"，倾向于把之前的错误一路传下去，而不是回头改掉。

这篇 arXiv 2606.16700 让我挺感兴趣的地方在于：它换了个生成范式来解这个问题。掩码扩散模型（Mask Diffusion Models, MDM）天生就支持"局部编辑"——掩掉某几个不确定的 token 重新采样，其他地方原封不动。这跟人类"迭代式局部修订"的纠错方式几乎是一回事。作者把这件事提炼成一个新原语，叫 Reflective Masking（反思性掩码），再配一个零参数的 History Reference 机制，让 MDM 真正学会"回头看、主动改"。

先把这篇论文一句话讲清楚

核心摘要：AR 模型修正错误必须顺序重写，局部编辑也得付全序列的代价；而 MDM 的掩码机制天然能做局部修订，却一直没被用来做"多轮反思"。这篇论文提出 Reflective Masking——把生成重新定义为"逐位置决策的迭代自修订过程"，每个位置每一步都在做三选一：保留、重新掩码、揭示。再用一套轻量级 post-training（2 张 H100 跑 5 小时）把这个能力激活出来，无需改动任何架构。配套的 History Reference 用零参数的旋转嵌入把整条去噪轨迹压成一个向量，让模型"记得"自己改过什么。结果很能打：图像编辑的编辑精度从基线的 65.81 拉到 99.73，MBPP 代码任务比 Vanilla SFT 涨了 8.8 个点，数独修订的精确解出率从 82.4 提到 93.4。

我的判断：这篇论文最值钱的不是某个 SOTA 数字，而是它给 MDM 找到了一个专属于扩散范式的 test-time scaling 方式。AR 模型靠"多生成 token"来 scaling，MDM 可以靠"多轮回头改"来 scaling——这个视角是真的漂亮。

论文信息

  • 标题：Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models
  • 作者：Yanming Zhang、Yihan Bian（共同一作）、Jingyuan Qi、Yuguang Yao、Lifu Huang、Tianyi Zhou
  • 机构：马里兰大学、弗吉尼亚理工、Intuit、加州大学戴维斯分校、MBZUAI
  • arXiv：2606.16700v1，2026 年 6 月 15 日
  • 链接：https://arxiv.org/abs/2606.16700

为什么是扩散模型，而不是继续卷 AR

先说点背景，免得后面看不懂为什么这条路走得通。

自回归模型生成文本是从左到右、一个 token 接一个 token。它的"反思"本质上是在序列末尾追加新内容——"我刚才说错了，应该是……"。问题在于，前面那段错的东西并没有被删，它还在上下文里待着。你想真正改掉它，只能从那个位置往后整段重新生成。一个局部的小错，付的是全序列的计算账。

掩码扩散模型走的是另一条路。它的生成过程是个吸收态马尔可夫去噪：一开始整个序列都是 MASK，模型一步步把 MASK 替换成确定的 token，一旦某个 token 被自信地"去噪"出来，就被固定住。这套机制有个天然的好处——你可以只掩掉某几个位置重新预测，其他地方完全不动。这就是"局部编辑"。

但现有的 MDM 没把这个能力用足。它们的 remasking 策略只盯着"低置信度 token"重新掩码，模型自己没法主动回头审视已经提交的预测。换句话说，MDM 有改局部的"手"，却没有反思的"脑"。这篇论文要补的就是这个脑子。

说实话，第一次看到这个 motivation 我是认可的。MDM 这两年（LLaDA、Lumina-DiMOO 这些）一直在追赶 AR，但很多工作都是在模仿 AR 的玩法。这篇反过来问：MDM 有什么是 AR 天生做不到的？局部修订就是答案。这个问题问得对。

方法核心：把生成变成"逐位置三选一"

Reflective Masking 的核心 idea 一句话就能讲清：在每个去噪步骤、每个位置上，模型都在做一个三选一的决策。

  • Keep（保留）：当前 token 看着没问题，留着
  • Re-mask（重新掩码）：当前 token 可能错了，把它掩掉重新预测
  • Reveal（揭示）：当前位置是 MASK，揭示出最可能的 token

关键在 Re-mask 这个动作——这是标准 MDM 没有的。标准 MDM 一旦揭示了 token 就锁死，而 RM 允许模型把已经写下的东西重新盖回去。这就是"反思"的物理实现。

那模型怎么知道该 keep 还是 re-mask？论文给了一条很干净的确定性规则（公式 1）。在时间步 \(t\)，模型对每个位置输出一个在扩展词表 \(\bar{\mathcal{V}} = \mathcal{V} \cup \{\text{MASK}\}\) 上的分布 \(p_\theta(\cdot \mid \tilde{x}^{(t)})_i\)，下一步状态这样定（\(M\) 表示 MASK）：

\[ \tilde{x}_i^{(t+1)} = \begin{cases} M & \tilde{x}_i^{(t)} \neq M \text{ 且 } p_\theta(M)_i > p_\theta(\tilde{x}_i^{(t)})_i \\ \tilde{x}_i^{(t)} & \tilde{x}_i^{(t)} \neq M \text{ 且 } p_\theta(M)_i \leq p_\theta(\tilde{x}_i^{(t)})_i \\ \arg\max_{v \in \mathcal{V}} p_\theta(v \mid \tilde{x}^{(t)})_i & \tilde{x}_i^{(t)} = M \end{cases} \]

直觉很有意思：如果模型给 MASK 的概率比给当前 token 的概率还高，说明模型在暗示这个位置当前的词是错的，于是把它重新掩掉。否则就保留。掩码位置则直接揭示最可能的词。

这个设计的妙处在于，它没有引入任何新的输出头或动作分类器。模型本来就要预测词表分布，现在只是把 MASK 也当成词表里一个可选项——MASK 的概率高低，自然就编码了"要不要回头改"的信号。简洁。

等等，这里有个坑。这条规则只看当前状态 \(\tilde{x}^{(t)}\)，那如果两个不同时间步产生了完全相同的状态，岂不是会陷入循环、来回 mask/reveal 同一个位置？论文确实意识到了这个问题——而这正是 History Reference 要解决的（后面讲），它通过改变模型的输入来打破循环，但保持公式 1 本身不变。

怎么训出这个能力：合成轨迹 + Oracle 监督

光有推理规则不够，得让模型真的学会"什么时候该 re-mask"。论文的训练范式很轻量，思路是造合成的修订轨迹，用 oracle 动作做监督。

先定义 oracle 修订规则（公式 3）。对每个可编辑位置，给定当前取值 \(z_i \in \{w_i, \text{MASK}, x_i^*\}\)（\(w_i\) 是错误 token，\(x_i^*\) 是目标 token），最优动作是确定的：

\[ \tau(z_i, x_i^*) = \begin{cases} \text{MASK} & z_i \in \mathcal{V} \setminus \{x_i^*\} \\ x_i^* & z_i = \text{MASK} \\ x_i^* & z_i = x_i^* \end{cases} \]

翻译一下：错的 token 就该掩掉，MASK 位置就该揭示成目标，已经对的就保留。

图1：使用 History Reference 的推理流程。所有状态被嵌入，历史状态进一步由 HER 处理；当前步的嵌入加到 history reference 上作为模型输入，预测下一步状态

图 1：带 History Reference 的推理流程。每一步的状态都被嵌入，历史状态经 HER（History Embedding Rotation）旋转处理后累积成 history reference，再和当前步嵌入相加喂给模型，预测下一步。这张图是整个推理时 pipeline 的总览。

训练数据怎么造？关键在于模拟出真实推理时会遇到的错误状态。从一条干净序列 \(x^*\) 出发：

  1. 采样一批要破坏的位置，采样时间步 \(t \sim \text{Uniform}\{0, \dots, T-1\}\)
  2. 把这些位置分成两组：一组替换成 MASK，另一组替换成错误 token \(w_i\)
  3. 错误 token 从一个破坏分布 \(\nu(\cdot \mid x_i^*)\) 里采。实践中很讲究——用冻结的 MDM backbone 的 top-k 预测（排除真值）来生成"以假乱真"的错误，让训练时的错误分布尽量贴近推理时模型自己会犯的错

然后用位置级转移规则把这些错误状态串成一条历史轨迹：正确的 token 一路保持；MASK 位置在某个时间步 \(t_1\) 转成正确 token；错误 token 先在 \(t_1\) 转成 MASK，再在更晚的 \(t_2\) 转成正确 token。这就模拟出了"先发现错→掩掉→重新填对"的完整修订过程。

图2：掩码扩散模型训练的合成历史数据构造。从干净序列通过 mask 和 wrong-token 两种破坏方式造出噪声序列，再用位置级转移规则采样合成历史并定义训练目标

图 2：合成历史数据的构造流程。左边是干净序列，通过 mask 破坏和 wrong-token 破坏制造噪声，右边用位置级转移规则把状态串成一条带"改错"过程的轨迹。这是整个训练范式的数据引擎，决定了模型能不能学到真实的修订模式。

训练目标就是逐位置的交叉熵（公式 4），用 oracle 动作 \(\tau(z_i, x_i^*)\) 做监督：

\[ \mathcal{L}_{\text{train}}(\theta) = \mathbb{E}\left[ \sum_{i \in E} -\log p_\theta(\tau(z_i, x_i^*) \mid a^{(t)})_i \right] \]

这个 loss 自然分解成三块：reveal loss（MASK 位置预测正确 token）、mask loss（错误位置预测 MASK）、keep loss（正确位置保持不变）。注意条件是 \(a^{(t)}\) 而不是裸状态 \(\tilde{x}^{(t)}\)——这个 \(a^{(t)}\) 就是下面要讲的 History Reference。论文还顺手证了一条性质：用更丰富的输入 \(a^{(t)}\) 做条件，不会增加最优训练风险，只会改善或持平。

History Reference：用零参数把整条轨迹"记"下来

回到前面那个循环的坑。如果模型只看当前状态，它没有记忆，可能在同一个位置反复横跳。人改稿子不会这样——你记得"这个词我刚才改过一次了"，所以会换个思路。History Reference（HR）就是给模型装上这个记忆。

核心做法：维护每个位置的累积嵌入，把整条前缀轨迹 \(\tilde{x}_i^{(0:t)}\) 压成一个向量喂给模型。最关键的是——零额外可学习参数，每步 O(1) 更新。

具体怎么压？论文借了 RoPE 的思路。设 \(e_i^{(k)} = \text{wte}(\tilde{x}_i^{(k)})\) 是位置 \(i\) 在第 \(k\) 步的 token 嵌入，定义一个 History Embedding Rotation（HER） \(R_\Delta\)，由历史步与当前步的距离 \(\Delta = k - t\) 索引。它满足旋转组合规则 \(R_0 = I\)、\(R_a R_b = R_{a+b}\)，用标准旋转编码的二维正弦块来实例化。然后 history embedding 这样累积（公式 2）：

\[ a_i^{(t)} = \sum_{k=0}^{t} \gamma^{t-k} R_{k-t} \, e_i^{(k)} \]

其中 \(\gamma \in (0, 1]\) 是历史衰减因子。展开看更直观：

\[ a_i^{(t)} = e_i^{(t)} + \gamma R_{-1} e_i^{(t-1)} + \gamma^2 R_{-2} e_i^{(t-2)} + \cdots + \gamma^t R_{-t} e_i^{(0)} \]

当前状态不旋转，每个过去状态按"滞后多少步"旋转一个角度、再乘个衰减。不同步的状态因此被旋到了不同的"相位"上，模型能区分"这是第几步写的"。

工程上更香的是它能写成 O(1) 递推（附录 B）：

\[ a_i^{(t)} = e_i^{(t)} + \gamma R_{-1} a_i^{(t-1)} \]

只需维护一个 \(N \times d\) 的运行张量，每步一次块对角旋转加逐元素相加。完整轨迹本来要 \(O(T \times N \times d)\) 的存储，递推直接压成单个 \(O(N \times d)\) 的缓冲区。而且 \(R_\Delta\) 是正交变换保范数，配合 \(\gamma \in (0,1]\)，运行和的范数随 \(t\) 最多线性增长，数值上稳。

说实话这个设计我挺喜欢的。它没有引入任何新参数，纯靠旋转嵌入的几何性质把"时间顺序"编码进历史信息里。零参数还能 O(1) 更新，这种"白嫖"式的机制设计，工程上太友好了。

图4：RM 在推理时基于演化的全局上下文主动重新掩码并修正 token，随着前置上下文被精炼，修正传播到最终答案

图 4：文本推理中 RM 的工作实况。模型在推理时主动 re-mask 并修正 token，随着前面的上下文被逐步精炼，修正会传播到最终答案。这张图直观展示了"边写边改、改了还会连锁修正"的过程——这正是 AR 模型做不到的局部+全局联动修订。

实验：三个任务，三种"指导强度"

实验设计有个我很认可的巧思：三个任务对应三种不同的"外部指导强度"，从强到无，层层递进验证 RM 的泛化性。整个训练在 2 张 NVIDIA H100 80GB 上约 5 小时就能跑完，对比那些要改架构、训几天几周的方法，这个成本优势很实在。

图像编辑（强指导）

Base 模型是 Lumina-DiMOO，从 ImgEdit 采 85,000 个样本训练，留 1,700 个测试，跟 vanilla SFT 在同样设置下对比。

方法 Edit Prec.↑ Coverage↑ MAE-RGB↓ PSNR(dB)↑ SSIM↑ VQAScore↑ User Study↑
Lumina 65.81 41.68 12.497 23.09 0.6256 71.95 41.8
Lumina+SFT 71.84 48.42 11.035 23.900 0.6570 81.61 53.3
RM（本文） 99.73 73.02 3.613 34.759 0.9744 85.17 68.2

编辑精度从基线的 65.81 干到 99.73，PSNR 从 23 拉到 34.8，提升幅度大得有点夸张。原因其实很直白：RM 能精确定位该编辑的区域并掩掉，只重新生成掩码区域，没掩的地方跟原图保持得极好。这正是局部编辑的天然优势——背景保留指标（MAE-RGB、PSNR、SSIM）的碾压式领先就是证据。

图3：图像编辑任务的定性结果。RM 预测的掩码用红色高亮，底部热图可视化编辑图像与原图的逐像素差异

图 3：图像编辑的定性结果。红色高亮是 RM 预测要修改的掩码区域，底部热图是编辑前后的逐像素差异。可以看到差异高度集中在该改的地方，背景几乎纹丝不动——这就是"局部修订"的视觉证据。

数独修订（结构化推理，有限指导 + 消融实验）

这块用的是个极轻量的 MDM——4 层 Transformer，0.81M 参数。测试时从已解的 9×9 棋盘随机破坏 4 到 20 个单元格，换成错误数字，让模型迭代修订。这个任务我觉得选得很聪明：数独有明确的对错，能干净地衡量"模型到底改对了没、有没有反复犯同一个错"。

数独这张表同时也是核心的消融实验：

变体 Exact Acc.(%)↑ Valid Rate(%)↑ Replay Mistake(%)↓ Conflict Cells↓
RM w/o HR 82.4 86.6 0.57 0.578
RM + HR 91.4 91.8 0.07 0.300
RM + HR + decay 89.4 89.6 0.07 0.362
RM + HR + decay + HER 93.4 93.6 0.03 0.236

这张表信息量很大，值得细看。加入 HR 后，精确解出率从 82.4 跳到 91.4，更关键的是 Replay Mistake（重复犯同一个错）从 0.57 暴降到 0.07——这直接验证了"记忆"的作用，模型不再傻乎乎地重访同一个错误预测。

但有个反直觉的细节：单独加 decay 反而比只用 HR 还差（89.4 vs 91.4）。这说明简单地"削弱历史信号"是不够的——你不能只是让旧信息衰减掉。只有当 HER 把历史信息显式解耦（旋转到不同相位）之后，再叠加 decay 才变得有益，最终四件套（RM+HR+decay+HER）拿到 93.4 的最优。

论文给出的结论我很认同：恰当地"结构化"历史信息，比单纯"衰减"它重要得多。这个消融做得诚实，没藏着 decay 单独用会掉点这件事，反而把它讲清楚了。

文本推理（完全自主，无指导）

Base 模型换成 LLaDA，在 MATH500、MBPP、ARC-Challenge 上测：

基准 类别 LLaDA Vanilla SFT Ours Δ
MATH500 Math 19.4 22.4 24.8 ↑2.4
MBPP Code 28.0 30.6 39.4 ↑8.8
ARC-Challenge MCQA 73.7 81.3 86.1 ↑4.8

这里有个很有意思的现象：MBPP（代码）涨了 8.8 个点，远超 MATH500 的 2.4 个点。论文的解释挺有说服力——代码的正确性依赖很多 token 都对，而 MATH500 主要看最终答案那一个数。RM 的迭代修订是 token 级的，能逐个纠正错误 token，所以在"对的 token 越多越好"的代码任务上收益自然更大。这个分析让我觉得作者是真的理解自己方法的作用边界，而不是看到哪个数高就吹哪个。

Minerva MATH 的分学科细分（表 4）也佐证了这点：RM 在几乎所有学科上都超过 Vanilla SFT，总分从 22.62 提到 24.10，其中数论（Num Theory）涨了 4.26 个点最猛。

图6：文本生成的更多定性结果

图 6：文本生成任务的补充定性案例，展示 RM 在实际生成中如何回头修订 token。

我的判断：漂亮的范式，克制的论文

先说亮点。这篇论文最打动我的是视角——它没有继续让 MDM 模仿 AR，而是挖出了 MDM 独有的能力：局部修订。把生成重新定义成"逐位置三选一的迭代自修订"，再用 MASK 的概率天然编码"要不要改"的信号，这套设计干净到几乎没有多余的零件。History Reference 用零参数旋转嵌入做轨迹记忆，O(1) 更新，工程上极其友好。整套方法无需改架构、2 张卡 5 小时就能激活，落地门槛低得不像研究原型。

更难得的是它诚实。消融实验里 decay 单独用会掉点这件事没有被藏起来，反而被拿出来讲清楚了背后的逻辑；MBPP 涨得多、MATH 涨得少的差异也给了合理解释，而不是含糊带过。这种"知道自己方法边界在哪"的论文，比那种到处吹 SOTA 的可信多了。

再说我的保留意见。作者自己在局限性里也承认了：评估的这三个任务——图像编辑、数独、文本生成——说到底还是比 AR 模型里最难的长程推理要简单。这个天花板很大程度上是被当前 base MDM 自身的推理能力卡住的。数独修订用的是 0.81M 参数的玩具模型，文本用的 LLaDA 也不算最强的 MDM。所以"RM 能不能在大规模、强 base 模型上同样 work"这个最关键的问题，论文其实没回答。

还有一个更本质的隐忧，作者也点到了：未来想用到 block diffusion（比如 LLaDA2.0 那种扩到 100B 的），直接套只能改"当前生成块内"的 token，没法跨块做全局修订。而真正的长程推理恰恰需要跨块的全局回头改。这个跨块修订机制怎么设计，是 RM 想从"漂亮的原语"走向"实用的推理引擎"必须迈过的坎。

所以我的定位是：这是一篇范式贡献大于刷榜贡献的论文。它不一定立刻能在你的产线上替换掉什么，但如果你在做 MDM、做扩散语言模型，或者在思考"除了堆 token 还能怎么做 test-time scaling"，这篇值得认真读。它至少证明了一件事——扩散模型的"边写边改"不是噱头，是能训出来、能涨点的真能力。

如果你也在折腾 MDM 的推理增强，RM 这套轻量 post-training + 零参数历史记忆的组合拳，是个值得直接上手试的 baseline。成本足够低，risk 足够小。

觉得有启发的话，欢迎点赞、在看、转发。跟进最新 AI 前沿，关注我