当扩散模型学会"后悔"：MDM-VGB 让采样器在掩码状态图上回溯

一个数独引发的思考

想象你在做数独。填到第 7 行第 3 列时，你突然意识到——三步之前那个 8 放错了。怎么办？

最直觉的策略是：擦掉重来。这就是当前大多数扩散模型在推理时做的事——Best-of-N 采样，生成 N 个完整解，挑最好的。但问题在于，如果错误发生在第 3 步，你在第 47 步才发现，那 N 个解里绝大多数都带着同样的早期错误，只是在不同地方"开花"而已。

另一种策略更聪明：只擦掉错的那一步，保留前面正确的部分。这就是 Jerrum 和 Sinclair 在 1989 年提出的回溯马尔可夫链思想——在奖励倾斜的随机游走中，允许向后走，但让向前的步子更受青睐。

最近，Caltech 的 Kijung Jeon、Thuy-Duong Vuong 和 Molei Tao 把这个经典思路搬到了掩码扩散模型（Masked Diffusion Model, MDM）上，提出了 MDM-VGB 采样器。论文标题是《VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing》，发表于 2026 年 6 月。

从前缀树到掩码状态图：一次关键跃迁

要理解 MDM-VGB 的创新，先得搞清楚它从哪里出发。

掩码扩散模型的工作方式是：一开始所有 token 都是 [MASK]，然后逐步把 mask 替换成真实 token，直到全部解掩码。传统的回溯方法（如 VGR, Value-Guided Backtracking）在一个固定前缀树上做随机游走——也就是说，回溯只能回到"已经确定的 token 序列"的末尾，然后重新生成。

但这有个根本限制：如果你在第 47 个 token 位置发现第 3 个 token 错了，VGR 只能回到第 47 个位置重新采样，无法修改第 3 个 token。

MDM-VGB 的突破在于：它把回溯随机游走从"前缀树"扩展到了"掩码状态图"。在掩码状态图上，每个节点是一个部分解掩码的状态（有些位置已揭示，有些还是 mask），每条边是一个"解掩码"或"重掩码"操作。这意味着——你可以在任意位置重新掩码，然后重新解掩码。

这就像数独高手做的事：不是擦掉整行重来，而是精准定位"第 3 行第 5 列那个 8 有问题"，只擦那一个格子，重新填。

奖励引导的回溯：怎么决定擦哪个？

光能回溯还不够，关键是怎么决定回溯到哪里。MDM-VGB 用一个"过程验证器"（process verifier）来评估当前部分配置的价值。具体来说：

1. 前向移动（unmasking）：把一个 mask 位置替换成真实 token。验证器评估新状态的价值。 2. 后向移动（remasking）：把一个已揭示的 token 重新变成 mask。验证器评估回溯后的状态价值。

采样器偏爱那些导致更高价值部分配置的移动。如果前向移动提高了价值，就接受；如果发现某个已揭示的 token 拖累了整体价值，就后向回溯把它重新掩码。

这形成了一个在掩码状态图上的随机游走，每一步都由奖励信号引导。论文证明，这个采样器对过程验证器的噪声具有鲁棒性——即使验证器不完美，回溯机制也能有效避免错误累积。

二次 vs 指数：复杂度的根本差异

MDM-VGB 最硬核的贡献是理论分析。论文证明：

• Best-of-N：由于错误累积，复杂度可能是指数级的。因为如果早期错误率高，N 个样本中能到达高奖励终点的概率随序列长度指数下降。
• MDM-VGB：复杂度是二次的——O(L²)，其中 L 是序列长度。因为每一步回溯只需要重新评估受影响的位置，不需要从头生成。

这个差异在实践中意味着什么？在数独（序列长度 81）上，Best-of-N 可能需要生成数千个完整解才能找到一个正确的；而 MDM-VGB 通过精准回溯，在几十步内就能修复一个接近正确的解。

代码实测：六种任务的全景验证

论文的开源代码（https://github.com/KraitGit/MDM-VGB）覆盖了六种任务，从约束满足到科学生成：

代码结构非常清晰：src/algorithms/ 下实现了 VGB、VGR 和 VGB-Momentum 三种采样器，src/tasks/ 下是六种任务的 harness。运行流程是：训练基座模型 → 收集 rollout → 训练验证器 → 推理。

核心回溯逻辑（简化版）

def sample_vgb(harness, task, examples, verifier, config): states = initial_states_and_prompts(task, examples) while not all_finished(states): active = active_indices(states) # 构建前向（解掩码）和后向（重掩码）候选 fwd_groups, bwd_groups = build_mdm_candidates( task, harness, verifier, active_states ) # 按奖励信号采样选择 choices = sample_grouped_choices(fwd + bwd) for idx, choice in enumerate(choices): states[idx] = choice["state"] # 可能是前向或后向

一个反直觉的发现：回溯比重启更高效

论文中最反直觉的结论是：允许"倒退"的采样器比只允许"前进"的采样器更高效。

直觉上，回溯似乎是在浪费时间——你已经走了那么远，为什么要退回来？但信息论告诉我们：如果一个部分配置已经走入死胡同，继续往前走只是在浪费计算资源。回溯让你把计算预算花在"有希望"的分支上。

这和 AlphaGo 的蒙特卡洛树搜索有异曲同工之妙。AlphaGo 不会在必败的分支上浪费模拟次数，而是通过 UCB 公式把资源分配给有希望的分支。MDM-VGB 做的是同样的事，只不过搜索空间从博弈树变成了掩码状态图。

对 AI 工程的启示

MDM-VGB 的贡献不只是一个新采样器，更是一种思维方式的转变：

1. 推理时缩放不等于暴力采样。Best-of-N 是最简单的推理时缩放策略，但它的效率受限于错误累积。回溯提供了一种更聪明的替代方案。

2. 离散扩散模型需要离散的回溯方法。连续扩散模型（如 DDPM）的回溯可以通过反向 SDE 实现，但离散扩散模型的回溯需要全新的数学工具。MDM-VGB 填补了这个空白。

3. 过程验证器比结果验证器更有价值。Best-of-N 只需要结果验证器（评估最终输出），但 MDM-VGB 利用过程验证器（评估中间状态）来引导搜索。这和 OpenAI 的 o1 模型思路一致——过程奖励模型（PRM）比结果奖励模型（ORM）更有效。

4. 回溯是通用的"编辑"工具。论文标题中的"Sample Editing"不是噱头——回溯机制天然支持"局部编辑"：给定一个已生成的样本，你可以回溯到不满意的部分，重新生成。这在分子优化、蛋白质设计等需要迭代改进的场景中极具潜力。

当回溯遇上扩散：未解的问题

MDM-VGB 打开了一扇门，但也留下了很多问题：

• 验证器的质量天花板。如果过程验证器本身不准确，回溯可能被引导到错误的方向。论文证明了鲁棒性，但鲁棒不等于免疫。
• 回溯深度 vs 计算预算。允许回溯到多深？太浅了等于没回溯，太深了计算开销大。论文用 max-steps-multiplier 控制，但最优值依赖任务。
• 与强化学习的关系。回溯随机游走在掩码状态图上，和 RL 中的 MCTS 有结构相似性。能否用 RL 训练更好的回溯策略？

结语

MDM-VGB 的核心洞察可以用一句话概括：在生成过程中允许"后悔"，比强迫模型一条路走到黑更高效。

这个洞察不只适用于扩散模型。在任何需要满足约束的生成任务中——从代码生成到分子设计——回溯都是一种值得考虑的策略。Jerrum 和 Sinclair 1989 年的回溯马尔可夫链，在 37 年后找到了新的应用场景。经典算法不会过时，它们只是在等待合适的土壤。

---

*论文：arXiv:2606.28301* *代码：https://github.com/KraitGit/MDM-VGB*