VGB for Masked Diffusion Model: Efficient Test-time Scaling for Reward Satisfaction and Sample Editing
Kijung Jeon, Thuy-Duong Vuong, Molei Tao
cs.LG, cs.DS, math.NA, math.PR, stat.ML
Submitted: 2026-06-26
Comments: 72 pages
Code: https://github.com/KraitGit/MDM-VGB
License: http://creativecommons.org/licenses/by/4.0/
The gist: Inference-time scaling is a promising paradigm to improve generative models, especially when outputs must satisfy structural constraints or optimize downstream rewards.
Terminology
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.
Sources
- Large Language Monkeys: Scaling Inference Compute with Repeated Sampling
- Training Verifiers to Solve Math Word Problems
- Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment
- Fine-Tuning Masked Diffusion for Provable Self-Correction
- Test-Time Scaling in Diffusion LLMs via Hidden Semi-Autoregressive Experts
- Effective Test-Time Scaling of Discrete Diffusion through Iterative Refinement
- UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching
- Inference-Time Scaling of Discrete Diffusion Models via Importance Weighting and Optimal Proposal Design
- Taming Imperfect Process Verifiers: A Sampling Perspective on Backtracking
- Qwen3 Technical Report
- D3LM: A Discrete DNA Diffusion Language Model for Bidirectional DNA Understanding and Generation
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks