MDM-VGB: Efficient Test-time Scaling for Masked Diffusion Models via Reward-Guided Remasking
Field: Machine Learning Authors: Kijung Jeon, Thuy-Duong Vuong, Molei Tao arXiv: 2606.28301
Abstract (original)
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 suffer exponential complexity due to error accumulation. Our theoretical findings are corroborated by strong empirical performance, particularly on Sudoku and QM9.
Key Points
- Introduces MDM-VGB, a discrete diffusion sampler for masked diffusion models combining unmasking with reward-guided remasking at inference time.
- Extends the classical Jerrum-Sinclair backtracking Markov chain from a fixed prefix tree to a masked-state graph, allowing tokens to be unmasked/remasked at arbitrary positions.
- Favors moves toward higher-value partial configurations, supporting both high-reward generation and efficient repair of low-reward samples.
- Theoretical guarantees: robustness to process-verifier noise and quadratic complexity, versus potentially exponential complexity of best-of-N style heuristics due to error accumulation.
- Strong empirical validation, notably on Sudoku and QM9.