VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing

推理时缩放（inference-time scaling）是一种极具前景的范式，可用于提升生成模型的性能，尤其适用于输出需满足结构化约束或需优化下游任务奖励的情形。本文以掩码扩散模型（Masked Diffusion Model, MDM）为基础，提出MDM-VGB——一种离散型扩散采样器；该方法在原有“解掩码生成”（unmasking generation）流程之上，引入了具有理论保障的、基于奖励引导的重掩码（reward-guided remasking）机制。受近期经典Jerrum-Sinclair回溯马尔可夫链在奖励偏置生成任务中取得成功之启发，MDM-VGB将原本局限于固定前缀树上的回溯随机游走，拓展至更为灵活的“掩码状态图”（masked-state graph）上，从而允许模型在任意位置对已解掩码的token进行重新掩码（remasking），亦可对已掩码位置执行解掩码操作。由此产生的采样器倾向于选择那些能导向更高局部价值（即更高奖励的部分配置）的解掩码与重掩码动作，既可高效生成高奖励样本，亦能快速修复低奖励样本。我们从理论上证明：MDM-VGB对过程验证器（process verifier）中的噪声具有鲁棒性，且整体计算复杂度为二次方级（quadratic complexity）；相比之下，当前广泛应用的测试时启发式方法（如“最优-N选一”（best-of-$N$））则可能因错误累积而陷入指数级复杂度。我们的理论分析得到了扎实实证结果的有力支持，尤其在Sudoku（数独）和QM9等广受关注的约束满足类与科学计算类基准任务上表现尤为突出。

如何在扩散模型（特别是Masked Diffusion Model, MDM）的推理阶段高效生成高质量、结构合规且高奖励的离散序列（如Sudoku解、分子图），克服传统采样方法（如best-of-N）在约束满足和奖励优化中面临的指数级复杂度与误差累积问题。这是一个新兴且重要的推理时优化问题，尤其针对离散扩散模型的可控生成。

提出MDM-VGB——一种基于理论保证的离散扩散采样器，将经典Jerrum-Sinclair回溯马尔可夫链从固定前缀树推广至动态掩码状态图，支持任意位置的‘解掩码→重掩码’双向操作，并以价值函数引导重掩码决策；首次在离散扩散框架中实现奖励导向的可证明鲁棒性（抗验证器噪声）与O(N²)二次复杂度，突破了启发式方法的理论瓶颈。

理论贡献突出：给出收敛性证明与噪声鲁棒性保证；实验验证充分：在Sudoku（强逻辑约束）、QM9（分子性质优化）等硬基准上显著超越MDM基线及best-of-100；代码已开源（论文附GitHub链接）；关键洞见是‘状态图上的价值引导回溯’可兼顾探索效率与修复能力，为后续研究提供新范式——如扩展至多模态掩码空间或与LLM验证器协同。

1. 'Masked Diffusion for Discrete Data' (ICML 2023); 2. 'Jerrum-Sinclair Markov Chains for Sampling from Exponential Families' (JACM 1996); 3. 'Test-Time Training with Reward Models' (NeurIPS 2022); 4. 'Best-of-N Sampling is Exponentially Inefficient' (ICLR 2024 Spotlight); 5. 'Diffusion-LM: Diffusion-based Language Modeling' (ACL 2023)