Research map / Image generation / Fast diffusion sampling

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

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Distillation for few-step generation · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes MDM-VGB, a reward-guided masked diffusion sampler whose random walk on a masked-state graph allows arbitrary-position unmasking and remasking, with theory on targeting the reward-tilted distribution and quadratic mixing, validated on Sudoku, QM9, DNA, protein, letter avoidance and Dyck editing.

Problem

Sampling from reward-tilted distributions over structured outputs with masked diffusion models at inference time is sample-inefficient: best-of-N suffers exponential error accumulation with imperfect process verifiers, and fixed-order backtracking cannot revise early mistakes.

Approach

Generalize the Jerrum-Sinclair value-guided backtracking chain from fixed-order prefix trees to the any-order masked-state graph, so tokens can be revealed and remasked at arbitrary positions with value-weighted transitions and depth-dependent normalization; add a momentum lift to speed leaf hitting.

Key results

• On QM9 generation at N=16, MDM-VGB-Momentum reaches Pass@95 of 90.6 vs 41.0 for BoN.
• On Sudoku, MDM-VGB-Momentum solves 100.0% of puzzles with 0.000 violated constraints vs 99.3% and 0.087 for BoN.
• On Dyck grammar editing, MDM-VGB reaches 99.41% accuracy with 29.98 moves vs 25.64% and 127.63 moves for AR-VGB.

Surprising findings

• The direct any-order generalization of AR-VGB (no depth normalization, lambda=0) is reported to have exponentially slow convergence.
• Momentum improves generation but slightly lowers editing accuracy for MDM-VGB on Dyck (97.54% with momentum vs 99.41% without) while reducing moves.

Limitations

• Training high-quality process verifiers is costly for long sequences and large vocabularies.
• Heuristic verifiers (e.g., Sudoku, Letter) can be far from the true value V*.
• Theoretical guarantees require kappa-accurate verifiers, and the leaf mass of the stationary distribution is only order 1/n without modification.

Experimental setup

Datasets: QM9 DeepSTARR Sudoku Dyck grammar

Benchmarks: Sudoku QM9 DNA enhancer design (DeepSTARR) Protein motif scaffolding Letter avoidance Dyck grammar editing

Metrics: success rate Pass@95 Success@1Å motif RMSD validity violated constraints adjusted NFE accuracy number of moves

Baselines: Base BoN MDM-VGR AR-VGB AR-VGB-Momentum ReMDM

Builds on: AR-VGB (Jerrum-Sinclair backtracking chain) MDLM Best-of-N ReMDM

Closest papers by meaning

• Accelerated Sampling from Masked Diffusion Models via Entropy Bounded Unmasking (NeurIPS 2025 · similarity 0.89)
• Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models (NeurIPS 2026 · similarity 0.88)
• Remasking Discrete Diffusion Models with Inference-Time Scaling (NeurIPS 2025 · similarity 0.88)
• Hierarchical Variational Policies for Reward-Guided Diffusion (NeurIPS 2026 · similarity 0.88)
• Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking (NeurIPS 2025 · similarity 0.87)
• Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions (ICML 2025 · similarity 0.87)