기술 요약: 스타인 확산 안내 (SDG)

문제 제기

확산 모델은 복잡하고 고차원적인 분포를 모델링하는 데 탁월한 성과를 보여주었습니다. 그러나 레이블이 없는 데이터로 훈련된 경우, 이러한 모델은 주로 데이터 매니폴드의 고밀도 영역에서 샘플링하며, 종종 새로운 약물 후보와 같은 고가치 희귀 샘플이 존재할 수 있는 저밀도 영역을 간과합니다.

이러한 한계를 해결하기 위한 기존 방법들은 일반적으로 두 가지 범주로 나뉘며, 둘 다 상당한 단점을 가지고 있습니다:

1. 분류기 기반 안내: 이러한 방법들은 여러 노이즈 수준에서 노이즈가 포함된 데이터로 훈련된 보조 분류기를 활용하여 샘플링을 안내합니다. 그러나 훈련 복잡성을 증가시키며, 분류기 기울기가 노이즈가 많을 경우 샘플을 오도하여 특히 저밀도 영역에서 생성 매니폴드에서 벗어나게 할 수 있습니다.
2. 확률적 최적 제어 (SOC): SOC 는 최적 제어기를 통해 확산 과정을 유도함으로써 사후 확률 샘플링을 위한 원칙적인 프레임워크를 제공합니다. 이론적으로 타당하며 깨끗한 데이터로 훈련된 오프더셸 분류기를 활용할 수 있지만, 최적 제어를 계산하려면 전체 신경 SDE 궤적을 통해 보상 신호를 역전파해야 합니다. 이는 효율적인 추론을 위해 SOC 를 계산적으로 불가능하게 만듭니다.

최근의 "훈련 없는" 접근법들은 SOC 의 계산 비용을 피하기 위해 트위디 공식을 사용하여 확산 사후 확률을 근사하려 시도합니다. 그러나 본 논문은 이 근사가 본질적으로 편향되어 있고 비최적이며, 특히 점수 기반 모델이 가장 정확도가 낮은 저밀도 영역에서 신뢰할 수 없는 안내를 초래한다고 주장합니다.

방법론: 스타인 확산 안내 (SDG)

저자들은 훈련 없는 방법의 효율성과 SOC 의 원칙적 성격을 조화시키는 **스타인 확산 안내 (SDG)**라는 훈련 없는 프레임워크를 제안합니다. SDG 는 대리 SOC 목적 함수에 기반하며, 근사 사후 확률을 정제하기 위한 스타인 보정 메커니즘을 도입합니다.

1. 대리 SOC 공식화

저자들은 새로운 비용 함수 $\tilde{J}(u, x, t)$를 가진 확률적 최적 제어 문제로 저밀도 확산 샘플링을 공식화합니다. 이 함수는 다음을 포함합니다:

• 고밀도 영역을 패널티로 부과하는 상태 비용 (점화 일정 $\alpha(t)$를 통해 저밀도 영역의 탐색을 장려).
• 오프더셸 분류기 또는 보상 모델 $r(x_T)$에 기반한 종단 비용 (강도 $\beta(t)$를 통해 샘플을 원하는 속성으로 안내).

그들은 SOC 가치 함수 $V(x,t)$에 대한 변분 상한을 유도합니다. 이 상한은 트위디 공식에 의존하는 표준 훈련 없는 안내를 최적화하는 것이 근사 사후 확률 $q$와 실제 사후 확률 $p$ 사이의 클루버 - 라이블러 (KL) 발산을 무시하면서 상한의 처음 두 항만을 최소화하는 것에 해당함을 보여줍니다. 이러한 무시는 비최적 안내로 이어집니다.

2. 왕복 스타인 보정

전체 SOC 의 계산 비용을 들이지 않고 KL 발산 항을 해결하기 위해, SDG 는 왕복 스타인 보정 메커니즘을 도입합니다:

• 역방향 매핑: 확산 매니폴드 $M_t$상의 노이즈 입자 $x_t$를 트위디 공식을 사용하여 데이터 매니폴드 $M_T$로 역방향 매핑하여 초기 근사 사후 확률 샘플 집합 $x_T$를 얻습니다.
• 스타인 보정: 그런 다음 이러한 입자들을 **스타인 변분 경사 하강법 (SVGD)**을 사용하여 $M_T$상에서 진화시킵니다. 입자들은 근사 사후 확률과 실제 사후 확률 사이의 KL 발산을 최소화하기 위해 가장 가파른 하강 방향을 따라 업데이트됩니다. 이는 모델 점수를 통해 근사된 조건부 사후 확률의 점수를 활용합니다: $\nabla_{x_T} \log p(x_T|x_t) \approx s_\theta(x_T) - \eta(t)s_\theta(x_t)$.
• 정방향 매핑: 보정된 입자들은 다음 확산 단계에 사용되기 위해 현재 매니폴드 $M_t$로 정방향 매핑됩니다.

이 메커니즘은 오프더셸 분류기를 활용하더라도 사후 확률 샘플이 실제 확산 역학 및 생성 매니폴드에 가깝게 유지되도록 효과적으로 정규화합니다.

3. 알고리즘

SDG 알고리즘은 이 보정을 표준 역방향 확산 과정에 통합합니다. 각 단계에서 다음과 같이 수행합니다:

1. 트위디 공식을 통해 초기 사후 확률 추정을 생성합니다.
2. 스타인 보정 단계를 적용합니다 (커널 상호작용 및 점수 차이를 기반으로 한 입자 업데이트).
3. 보정된 샘플을 사용하여 역방향 SDE 업데이트를 위한 안내 신호 (보상 기울기) 를 계산합니다.

주요 기여

1. 저밀도 프레임워크: 새로운 비용 - 투 - 고 함수를 통해 저밀도 영역을 명시적으로 목표로 하는 확산 안내를 위한 새로운 SOC 기반 공식화.
2. 이론적 통찰: 트위디 공식을 통한 확산 사후 확률 근사가 비최적이며, 근사 및 실제 사후 확률 간의 KL 발산을 고려하기 위해 보정 단계가 필요함을 보여주는 이론적 증명.
3. 스타인 보정 메커니즘: 근사 및 실제 사후 확률 간의 KL 발산을 반복적으로 최소화하여 저밀도 영역에서 견고한 안내를 가능하게 하는 계산 효율적이고 훈련 없는 스타인 보정의 도입.
4. 일반화 가능성: 모델 파인튜닝이나 증류 없이 다양한 확산 사후 확률 샘플링 문제에 통합될 수 있는 플러그 - 앤 - 플레이 프레임워크.

실험 결과

저자들은 이미지 안내 및 단백질 도킹을 위한 분자 샘플링 등 다양한 작업에서 SDG 를 평가했습니다.

• 이미지 안내 작업: 이미지 레이블 안내, 가우시안 디블러링, 초해상도, 텍스트 - 투 - 이미지 스타일 전이와 같은 작업에서 SDG(특히 저밀도 샘플링이 없는 변형인 SDG$\heartsuit$) 는 대부분의 지표에서 표준 훈련 없는 베이스라인 (예: DPS, LGD) 과 심지어 반복 전략을 가진 고급 방법 (UGD, MPGD) 보다 일관되게 우수한 성능을 발휘했습니다. 이는 저밀도가 아닌 설정에서도 근사 사후 확률을 보정하는 것의 중요성을 강조합니다.
• 분자 샘플링 (저밀도): 리간드 - 단백질 도킹 작업 (Fa7, 5ht1b, Jak2, Parp1) 에서 SDG 는 새로운 "히트"분자 (신규성, 합성 접근성, 약물 유사성, 도킹 점수를 엄격하게 충족) 를 생성하는 데 있어 베이스라인 (GDSS, MOOD, FREED 등) 보다 크게 우월한 성능을 보였습니다.
  • 애블레이션 연구: 스타인 보정을 제거하면 (SDG$\clubsuit$) 성능이 저하되어 샘플이 종종 히트 기준을 충족하지 못하거나 생성 매니폴드에서 벗어나게 됩니다. 저밀도 샘플링을 제거하면 (SDG$\heartsuit$) 역시 비최적 결과가 나와 신규 발견을 위해 저밀도 영역을 목표로 하는 것의 필요성을 확인시켜 줍니다.
  • 보상 과대평가: 본 논문은 스타인 보정 없이는 보상 모델이 저밀도 영역에서 보상을 과대평가하여 비현실적인 분자를 생성한다고 입증합니다. SDG 는 정규화 메커니즘으로 작용하여 이러한 과대평가를 완화하고 샘플 유효성을 유지합니다.
  • 효율성: SDG 는 런타임 오버헤드 (솔버 및 작업에 따라 약 36-85%) 를 발생시키지만, 전체 SOC 에 비해 최소한의 메모리 오버헤드로 대규모 입자 크기 (최대 3000) 까지 확장 가능하며 계산적으로 실행 가능합니다.

중요성 및 주장

본 논문은 SDG 가 사후 확률 샘플링을 위한 prohibitively 비싼 확률적 최적 제어 방법에 대한 계산 효율적이고 훈련 없는 대안을 제공한다고 주장합니다. 사후 확률 보정의 필요성을 이론적으로 확립하고 스타인 변분 추론을 통해 이를 구현함으로써, SDG 는 고가치 발견 (예: 새로운 약물 후보) 이 존재할 가능성이 높은 저밀도 영역에서의 효과적인 샘플링을 가능하게 합니다.

저자들은 그들의 방법이 다음과 같음을 강조합니다:

• 트위디 기반 근사의 편향을 극복합니다.
• 샘플이 생성 매니폴드에서 벗어나는 것을 방지합니다.
• 저밀도 영역에서의 보상 과대평가를 완화합니다.
• 고밀도 영역을 넘어 다양한 사후 확률 샘플링 문제에 적용 가능한 일반 프레임워크를 제공합니다.

이 연구는 근사 사후 확률을 보정하는 것이 훈련 없는 확산 안내에서 종종 간과되는 중요한 구성 요소이며, SDG 가 과학 분야에서 희귀하고 고가치인 샘플을 발견하기 위한 확산 모델의 잠재력을 unlocked 한다고 시사합니다.