MDM-VGB:掩码扩散模型的推理时缩放与奖励引导生成

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

精选理由

这篇论文给掩码扩散模型加了重新掩码机制,像有个纠错开关,Sudoku和QM9上效果很好,值得做推理优化的读者看看。

AI 摘要

MDM-VGB是一种面向掩码扩散模型的高效采样器,在推理时通过奖励引导的重新掩码机制来增强生成质量。理论证明MDM-VGB对过程验证器噪声具有鲁棒性,计算复杂度为二次方,而best-of-N方法可能因误差累积达到指数复杂度。在Sudoku和QM9等约束满足与科学基准上,MDM-VGB实现了强实证性能,同时支持高奖励生成与低奖励样本的有效修复。

原文 · arXiv cs.LG

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

Inference-time scaling is a promising paradigm to improve generative models, especially when outputs must satisfy structural constraints or optimize downstream rewards. We consider Masked Diffusion Model (MDM) and introduce MDM-VGB, a discrete diffusion sampler that augments unmasking generation with theoretically principled reward-guided remasking. Inspired by the recent success of the classical Jerrum-Sinclair backtracking Markov chain in reward-tilted generation, MDM-VGB extends the backtracking random walk from a fixed prefix tree to a masked-state graph, allowing tokens to be unmasked and remasked at arbitrary positions. The resulting sampler favors unmasking and remasking moves that lead to higher-value partial configurations, enabling both effective high-reward generation and efficient repair of low-reward samples. We prove that MDM-VGB is robust to process-verifier noise and achieves quadratic complexity, while popular test-time heuristics such as best-of-$N$ can incur exponential complexity due to error accumulation. Our theoretical findings are corroborated by strong empirical performance, particularly on popular constraint-satisfaction and scientific benchmarks such as Sudoku and QM9.