검색 강화 마스크 확산으로 제약 만족 생성

초록

SearchDiff는 훈련 없이 이산 마스크 확산 모델의 역확산 과정에 제약 기반 탐색을 삽입해, 각 단계에서 모델이 제시한 후보 집합을 사용자가 정의한 제약·속성에 맞게 최적화한다. 이를 통해 생물학적 설계와 논리 퍼즐 등에서 제약 충족률과 목표 속성 점수를 크게 향상시키면서도 기존 확산 모델의 생성 품질을 유지한다.

상세 분석

본 논문은 이산 마스크 확산 모델이 “데이터 분포와 일치”하는 확률적 목표만을 최적화하고, 추론 시 하드 제약이나 비미분 가능 속성을 직접 제어할 메커니즘이 없다는 근본적인 한계를 지적한다. 이를 해결하기 위해 제안된 SearchDiff는 훈련 단계에 전혀 개입하지 않고, 역확산 과정의 각 타임스텝에 “검색‑증강(검색‑augmented) 디노이징” 단계를 삽입한다. 구체적으로, 현재 마스크된 상태 xₜ 에 대해 디노이저 x_θ 가 각 위치별 클린 토큰에 대한 확률 분포 \hat{x}^{0}{(t)} 를 출력한다. 이 분포는 후보 생성의 사전(prior)으로 활용되며, 후보 집합은 두 단계로 정제된다. 첫 번째는 “Candidate Search Sampling”(CSS)이라 부르는 전역 탐색으로, \hat{x}^{0}{(t)} 를 기반으로 높은 확률을 가진 토큰 조합을 샘플링하고, 제약 위반 점수 V(x) 를 최소화하는 후보를 선택한다. 두 번째는 로컬 리파인먼트 단계로, 선택된 후보에 대해 단일 토큰 교체·삽입·삭제와 같은 이산 액션을 적용해 위반 점수를 추가 감소시킨다. 이때 제약 위반 함수 ν_k(x) 와 가중치 λ_k 를 이용해 다중 제약을 가중합 형태로 통합한다. 최종 후보 \bar{x}_t 는 수정된 역전이 커널 \bar{p}θ(x{t‑1}| \bar{x}_t, x_t) 에 입력되어, 마스크된 토큰을 선택적으로 언마스크하고, 언마스크된 토큰은 후보값으로 확정한다. 이렇게 하면 탐색 단계가 모델의 사전 확률을 크게 왜곡하지 않으면서도 제약 만족도를 크게 끌어올릴 수 있다.

핵심 기술적 기여는 다음과 같다. 첫째, 제약‑위반 함수 V(x) 를 명시적으로 정의하고, 이를 최소화하는 이산 탐색 문제를 역확산 단계와 결합함으로써 “훈련‑프리” 제어 메커니즘을 구현한다. 둘째, 탐색 공간 X_t 를 완전한 토큰 시퀀스로 제한하고, 마스크 토큰은 확산 과정에만 맡겨 병렬 디노이징의 장점을 유지한다. 셋째, 탐색이 비미분 가능 블랙박스 제약(예: 화학 시뮬레이터, 논리 솔버)에도 적용 가능하도록 설계돼, 기존의 gradient‑guided 또는 classifier‑guided 방법이 요구하는 미분 가능성 가정에서 자유롭다.

실험에서는 5가지 도메인(분자 설계, 펩타이드, tRNA, Sudoku, Boolean SAT)에서 기존 이산 확산 및 autoregressive 모델을 크게 앞선 성능을 보였다. 특히 분자 생성에서는 QED 점수를 최대화하면서 합성 가능성 비율을 4배까지 끌어올렸고, Boolean SAT에서는 정확도를 9.6 %에서 76 %로 상승시켰다. 이러한 결과는 SearchDiff가 “제약 만족도”와 “목표 속성 최적화” 사이의 트레이드오프를 효과적으로 해결함을 입증한다.