Search-Augmented Masked Diffusion Models for Constrained Generation

Discrete diffusion models generate sequences by iteratively denoising samples corrupted by categorical noise, offering an appealing alternative to autoregressive decoding for structured and symbolic generation. However, standard training targets a likelihood-based objective that primarily matches the data distribution and provides no native mechanism for enforcing hard constraints or optimizing non-differentiable properties at inference time. This work addresses this limitation and introduces Search-Augmented Masked Diffusion (SearchDiff), a training-free neurosymbolic inference framework that integrates informed search directly into the reverse denoising process. At each denoising step, the model predictions define a proposal set that is optimized under a user-specified property satisfaction, yielding a modified reverse transition that steers sampling toward probable and feasible solutions. Experiments in biological design and symbolic reasoning illustrate that SearchDiff substantially improves constraint satisfaction and property adherence, while consistently outperforming discrete diffusion and autoregressive baselines.

💡 Research Summary

The paper tackles a fundamental limitation of discrete diffusion models: while they excel at globally denoising sequences and capturing complex dependencies, their training objective is purely likelihood‑based and offers no native way to enforce hard constraints or optimize non‑differentiable properties during generation. Existing controllable diffusion approaches either require additional fine‑tuning, rely on gradient‑based guidance that assumes differentiable objectives, or introduce projection steps that can bias the distribution. Moreover, autoregressive search‑based methods suffer from left‑to‑right bias and error accumulation when constraints depend on the full sequence.

To address these issues, the authors propose Search‑Augmented Masked Diffusion (SearchDiff), a training‑free inference framework that embeds a discrete search step into each reverse denoising iteration of a masked diffusion model. The algorithm proceeds as follows. Starting from the fully masked latent state x_T, the model iterates backward from t = T down to 1. At each step:

1. Clean‑state proposal – The denoiser x_θ produces a per‑position categorical distribution \hat{x}^{0}_{(t)} over the vocabulary, interpreted as a proposal prior for the clean sequence.
2. Constraint‑aware refinement – Using \hat{x}^{0}_{(t)} and the current diffusion context (x_t, t), a search operator generates a discrete candidate \bar{x}t. The search is formulated as a step‑conditioned optimization problem S_t = ⟨X_t, A_t, Succ_t, V⟩ where V(x) = Σ_k λ_k ν_k(x) aggregates weighted violations of K user‑defined constraints. The operator first performs Candidate Search Sampling (CSS), which samples high‑probability token configurations guided by \hat{x}^{0}{(t)} and selects the one with lowest V(x). It then applies a local refinement that explores single‑token edits (insert, delete, replace) to further reduce V(x) without leaving the admissible set X_t. Importantly, the search relies only on black‑box evaluations of ν_k, so non‑differentiable simulators or symbolic solvers can be incorporated.
3. Modified reverse transition – The refined candidate \bar{x}_t feeds into a modified reverse kernel \bar{p}_θ. Tokens already unmasked are set deterministically to their values in \bar{x}_t; for each currently masked position, a binary unmasking decision is sampled. If unmasked, the token is set to the value proposed by \bar{x}_t; otherwise it remains masked for the next step.

By interleaving search with denoising, SearchDiff preserves the inductive biases of the pretrained diffusion model while actively steering the trajectory toward regions that satisfy user constraints. The method is training‑free: no additional parameters are learned, and no gradient information from the constraints is required.

The authors evaluate SearchDiff on five tasks spanning molecular design (small molecules, peptides, tRNA), and symbolic reasoning (Sudoku, Boolean SAT). In molecular generation, the approach dramatically improves synthetic accessibility—up to a four‑fold increase over standard diffusion—while achieving the highest average QED (drug‑likeness) scores. For Boolean SAT, accuracy jumps from 9.6 % (baseline) to 76.0 %, an almost eight‑fold gain, and it outperforms strong autoregressive baselines. Across all domains, the method consistently yields higher constraint‑satisfaction rates and better property adherence without sacrificing sample quality.

Key contributions are: (1) introducing a training‑free, inference‑time search augmentation that enforces arbitrary hard constraints in discrete diffusion; (2) formalizing constrained diffusion via a weighted violation function and defining search‑guided reverse transitions that respect both the model’s proposal distribution and the constraints; (3) demonstrating broad applicability to non‑differentiable scientific objectives and showing substantial empirical gains over existing diffusion and autoregressive methods.

The paper opens several avenues for future work, including more sophisticated search heuristics (e.g., Monte‑Carlo tree search), multi‑objective optimization, and scaling to large‑scale scientific design problems such as protein engineering or materials discovery. Overall, SearchDiff represents a significant step toward practical, constraint‑aware generative modeling in domains where feasibility and target properties are as critical as data fidelity.