Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Search-Augmented Masked Diffusion Models for Constrained Generation

Huu Binh Ta, Michael Cardei, Alvaro Velasquez, Ferdinando Fioretto

TL;DR

This work tackles constrained generation for discrete diffusion models by introducing SearchDiff, a training-free inference framework that embeds a constraint-aware search into each reverse denoising step. By generating a denoiser-based proposal, performing CSS and local refinement to minimize a constraint-violation objective, and applying a modified reverse kernel, the method steers samples toward feasible regions without retraining. The approach yields substantial gains in constraint satisfaction and non-differentiable property adherence across molecular design, peptide design, tRNA generation, and symbolic reasoning benchmarks like Sudoku and Boolean SAT, often outperforming both diffusion-based and autoregressive baselines. The results demonstrate the practicality and versatility of trajectory-level, inference-time optimization for reliable, constraint-driven discrete sequence generation in scientific and reasoning domains.

Abstract

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.

Search-Augmented Masked Diffusion Models for Constrained Generation

TL;DR

This work tackles constrained generation for discrete diffusion models by introducing SearchDiff, a training-free inference framework that embeds a constraint-aware search into each reverse denoising step. By generating a denoiser-based proposal, performing CSS and local refinement to minimize a constraint-violation objective, and applying a modified reverse kernel, the method steers samples toward feasible regions without retraining. The approach yields substantial gains in constraint satisfaction and non-differentiable property adherence across molecular design, peptide design, tRNA generation, and symbolic reasoning benchmarks like Sudoku and Boolean SAT, often outperforming both diffusion-based and autoregressive baselines. The results demonstrate the practicality and versatility of trajectory-level, inference-time optimization for reliable, constraint-driven discrete sequence generation in scientific and reasoning domains.

Abstract

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.
Paper Structure (29 sections, 16 equations, 8 figures, 10 tables, 1 algorithm)

This paper contains 29 sections, 16 equations, 8 figures, 10 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related work
  3. Preliminaries
  4. Problem Formulation
  5. SearchDiff
  6. Search-augmented denoising (Step 2)
  7. Search Initialization and local refinement
  8. Modified reverse kernel (Step 3)
  9. Experiment
  10. Biological Sequence Generation
  11. Small Molecules
  12. Peptides
  13. Transfer ribonucleic acid
  14. Symbolic Reasoning Tasks
  15. Conclusion
  16. ...and 14 more sections

Figures (8)

  • Figure 1: Illustration of the SearchDiff method. Task: QED-optimized molecular generation.
  • Figure 2: Illustration of a single SearchDiff denoising step: (1) The denoiser predicts a proposal distribution $\hat{\bm x}_0^{(t)}$. (2) The search operator refines this into a low-violation candidate $\bar{\bm x}_t$ via (2a) CSS and (2b) a local search refinement step. (3) The modified reverse kernel used $\bar{\bm x}_t$ to deterministically update unmasked tokens and stochastically unmask a subset of masked tokens, yielding the next latent $\bm x_{t-1}.$
  • Figure 3: Average number of generated molecules with higher QED compared to the maximum value of QED in the training dataset for different algorithms.
  • Figure 4: tRNAs generated by SearchDiff (left) and MDLM (right).
  • Figure 5: Average QED by LS and no LS approaches with different numbers of CSS candidates.
  • ...and 3 more figures