Search-Augmented Masked Diffusion Models for Constrained Generation

TL;DR

This paper introduces SearchDiff, a novel inference framework for discrete diffusion models that integrates informed search to enforce hard constraints and optimize non-differentiable properties during sequence generation.

Contribution

SearchDiff is a training-free neurosymbolic inference method that enhances discrete diffusion models by incorporating search into the denoising process for constrained generation.

Findings

Significantly improves constraint satisfaction in biological design and symbolic reasoning tasks.

Outperforms discrete diffusion and autoregressive baselines in property adherence.

Demonstrates effectiveness without additional training or modifications to the model.

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.