Research map / Image generation / Fast diffusion sampling

Search-Augmented Masked Diffusion Models for Constrained Generation

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Discrete diffusion solvers · 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 SearchDiff, a training-free method that embeds candidate search and local refinement into each reverse step of masked diffusion to enforce constraints with black-box violation functions.

Problem

Standard discrete diffusion models optimize a likelihood objective and cannot natively enforce hard constraints or optimize non-differentiable properties at inference time.

Approach

At each reverse denoising step, draw best-of-M candidates from the denoiser's predicted distribution (CSS), refine them with greedy Hamming-1 local search minimizing a weighted constraint-violation function, and use the refined tokens to drive stochastic unmasking.

Key results

• On QED maximization, SearchDiff reaches 497.6 admissible molecules and mean QED 0.77 versus 0.62 for the best baseline (UDLM CFG).
• On tRNA generation, SearchDiff reaches 719.0 molecules satisfying all constraints, about 27x the best baseline (TreeG-SC 25.8).
• On Boolean SAT (7 variables), SearchDiff raises accuracy from 9.6% (vanilla MDM) to 76.0% with a 6M-parameter model.

Surprising findings

• Applying search only at the final denoising step gives little or no gain (e.g., Boolean SAT 11.4%), indicating search must be distributed across the trajectory.
• Local search alone generates over 200 molecules with QED above the training maximum, while standard decoding and CBG generate none.

Limitations

• Inference is slower than vanilla MDLM (e.g., 3.66-10.68 s vs 0.67 s per molecular batch), and cost grows with CSS budget.
• Results on tRNA have high variance (719.0 ± 105.1), and the search operator relies on hand-designed violation functions per task.
• Baseline comparisons use mixed setups and some baseline numbers are copied from prior work (CDD).

Experimental setup

Datasets: QM9 GRAMPA GtRNAdb (mammal tRNA) Sudoku random 3-SAT

Benchmarks: QM9 molecule constraints (SA, QED) GRAMPA peptide constraints tRNA structural constraints Sudoku accuracy Boolean SAT accuracy

Metrics: admissible molecules mean QED SA<=4 count constraint satisfaction count accuracy

Baselines: MDLM MDLM CBG UDLM CBG MDLM CFG UDLM CFG AR FUDGE AR PPLM AR CFG CDD TreeG-SC GPT-2 (6M, 85M, 303M)

Builds on: masked diffusion models (MDLM) MDLM CBG CFG TreeG

Closest papers by meaning

• Primal-Dual Guided Decoding for Constrained Discrete Diffusion (NeurIPS 2026 · similarity 0.90)
• Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-based Decoding (NeurIPS 2025 · similarity 0.90)
• Constrained Discrete Diffusion (NeurIPS 2025 · similarity 0.90)
• Encoder-Decoder Diffusion Language Models for Efficient Training and Inference (NeurIPS 2025 · similarity 0.89)
• Fast constrained sampling in pre-trained diffusion models (NeurIPS 2025 · similarity 0.89)
• Simplified and Generalized Masked Diffusion for Discrete Data (NeurIPS 2024 · similarity 0.89)