VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing
2026-06-26 • Machine Learning
Machine LearningData Structures and Algorithms
AI summaryⓘ
The authors study a way to improve generative models by changing their output during the generation process based on rewards or constraints. They focus on a model called Masked Diffusion Model (MDM) and create MDM-VGB, which allows the model to both reveal and hide parts of its output to get better results. Inspired by a known method called the Jerrum-Sinclair backtracking Markov chain, their approach lets the model explore different possibilities more flexibly to fix mistakes and increase the quality of the output. They prove that their method is efficient and performs well on tasks like Sudoku and molecular generation, even when verification is noisy.
generative modelsdiffusion modelsmasked diffusionreward-guided samplingbacktracking Markov chainJerrum-Sinclair algorithmconstraint satisfactioninference-time scalingSudokuQM9 dataset
Authors
Kijung Jeon, Thuy-Duong Vuong, Molei Tao
Abstract
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.