Search-Augmented Masked Diffusion Models for Constrained Generation

2026-02-02

AI Summary

This work addresses the challenge of enforcing hard constraints or optimizing non-differentiable properties during inference in discrete diffusion models. To this end, the authors propose SearchDiff, a neuro-symbolic inference framework that requires no additional training and, for the first time, seamlessly integrates heuristic search into the reverse denoising process of discrete diffusion. At each denoising step, the method constructs a candidate set based on the model’s predictions and refines it through user-specified constraints or desired attributes, yielding sequences that are both high-probability and compliant with the given criteria. Evaluated on biological sequence design and symbolic reasoning tasks, SearchDiff consistently outperforms existing discrete diffusion and autoregressive baselines, achieving substantially higher constraint satisfaction rates and attribute alignment.

Technology Category

• Search and Optimization: Mixed Discrete/Continuous Search
• Constraint Satisfaction and Optimization: Search
• Computer Vision: Diffusion Models for Vision

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: ML for personalized search and recommendations

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.

Problem

• constrained generation
• discrete diffusion models
• hard constraints
• non-differentiable properties
• structured generation

Innovation

• search-augmented generation
• masked diffusion models
• constrained generation
• neurosymbolic inference
• discrete diffusion