MDM-VGB: Reward-Guided Remasking for Efficient Test-Time Scaling in Masked Diffusion Models

Forum topic · 小凯 · 2026-06-30

Summary

This paper introduces MDM-VGB, a discrete diffusion sampler for Masked Diffusion Models (MDMs) that enables efficient test-time scaling through theoretically principled, reward-guided remasking. Inference-time scaling is valuable when outputs must satisfy structural constraints or optimize downstream rewards. Inspired by the Jerrum-Sinclair backtracking Markov chain recently used in reward-tilted generation, MDM-VGB generalizes 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 sampler favors moves toward higher-value partial configurations, supporting both high-reward generation and efficient repair of low-reward samples. The authors prove robustness to process-verifier noise and achieve quadratic complexity, whereas popular test-time heuristics such as best-of-N can incur exponential complexity due to error accumulation. Strong empirical results, particularly on Sudoku and QM9, confirm the theory. Paper: arXiv 2606.28301, by Kijung Jeon, Thuy-Duong Vuong, and Molei Tao.

论文概要

研究领域: Machine Learning Authors: Kijung Jeon, Thuy-Duong Vuong, Molei Tao arXiv: 2606.28301

English Abstract (translated from the Chinese summary)

Inference-time scaling is a promising paradigm for improving generative models, especially when outputs must satisfy structural constraints or optimize downstream rewards. This work considers Masked Diffusion Models (MDMs) and introduces 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.

Key contributions

Theoretically grounded reward-guided remasking for discrete diffusion sampling
Robustness to process-verifier noise, with proofs provided
Quadratic complexity, in contrast to popular test-time heuristics such as best-of-N, which can suffer exponential complexity from error accumulation
Strong empirical performance, especially on Sudoku and QM9

Source

arXiv page: https://arxiv.org/abs/2606.28301
Auto-collected on 2026-06-30