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

2026-06-26

AI Summary

This work addresses the challenge of efficiently generating high-quality samples during inference that satisfy structural constraints or optimize downstream rewards, while also enabling effective editing of low-quality outputs. The authors propose MDM-VGB, a discrete diffusion sampler tailored for masked diffusion models, which introduces a theoretically grounded reward-guided remasking mechanism. This mechanism dynamically performs unmasking and remasking at arbitrary positions to enhance sample rewards and correct flawed generations. By extending the Jerrum–Sinclair backtracking Markov chain from fixed prefix trees to mask state graphs, MDM-VGB enables flexible remasking strategies that maintain robustness to verifier noise and retain quadratic time complexity. Experiments demonstrate that MDM-VGB significantly outperforms heuristic baselines such as best-of-N on constraint-satisfaction and scientific benchmarks like Sudoku and QM9, achieving substantial reductions in computational cost while avoiding error accumulation.

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.

Problem

• Masked Diffusion Model
• Inference-time scaling
• Reward-guided generation
• Constraint satisfaction
• Sample editing

Innovation

• Masked Diffusion Model
• reward-guided remasking
• backtracking Markov chain
• test-time scaling
• quadratic complexity