제약 조건 하의 생성에 대한 검색 기반 마스크 확산 모델

Search-Augmented Masked Diffusion Models for Constrained Generation

이산 확산 모델은 범주형 노이즈에 의해 손상된 샘플을 반복적으로 제거하여 시퀀스를 생성하며, 구조화되고 상징적인 생성에 대한 자기 회귀 디코딩의 매력적인 대안을 제공합니다. 그러나 표준 훈련은 주로 데이터 분포와 일치하는 likelihood 기반 객관 함수를 목표로 하며, 추론 시에 엄격한 제약 조건을 적용하거나 미분 불가능한 특성을 최적화할 수 있는 기본적인 메커니즘을 제공하지 않습니다. 본 연구에서는 이러한 제한점을 해결하고, 사용자가 지정한 속성 만족 하에 최적화된 제안 집합을 정의하여 역 노이즈 제거 과정에 정보를 기반으로 한 검색을 직접 통합하는, 훈련이 필요 없는 신경-기호 추론 프레임워크인 Search-Augmented Masked Diffusion (SearchDiff)을 소개합니다. 각 노이즈 제거 단계에서 모델의 예측은 사용자가 지정한 속성을 만족하도록 최적화되어 수정된 역 전환을 생성하며, 이를 통해 샘플링이 가능하고 타당한 솔루션으로 향하도록 유도합니다. 생물학적 설계 및 기호 추론 분야에서의 실험 결과는 SearchDiff가 제약 조건 만족 및 속성 준수 측면에서 크게 향상되었으며, 이산 확산 및 자기 회귀 기반 모델보다 일관되게 우수한 성능을 보인다는 것을 보여줍니다.

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.