Search-Augmented Masked Diffusion Models for Constrained Generation
This addresses the challenge of generating sequences under constraints for applications like biological design and symbolic reasoning, representing a novel method for a known bottleneck rather than an incremental improvement.
The paper tackled the problem of discrete diffusion models lacking mechanisms for enforcing hard constraints or optimizing non-differentiable properties during inference, and introduced SearchDiff, a training-free neurosymbolic framework that integrates informed search into denoising to improve constraint satisfaction and property adherence, outperforming baselines in biological design and symbolic reasoning tasks.
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.