Uncertainty Estimation for Molecular Diffusion Models

기술 요약: 분자 확산 모델을 위한 불확실성 추정

문제 정의
확산 모델(Diffusion models)은 3D 분자 생성 분야의 지배적인 패러다임이 되었습니다. 그러나 이러한 모델에는 생성된 분자가 화학적으로 유효하지 않거나 불안정할 때 이를 식별할 수 있는 원칙적인 신호가 부족하다는 결정적인 실무적 병목 현상이 존재합니다. 예측 모델은 일반적으로 불확실성 추정을 사용하여 신뢰할 수 없는 출력을 표시하지만, 기존의 방법들은 주로 이미지 도메인에 집중되어 왔습니다. 다운스트림 평가(예: 도킹 또는 실험실 검증) 비용이 높은 분자 도메인에서는, 조잡한 규칙 기반 필터나 모델 자체의 가능도(likelihood)—종종 보정이 잘 이루어지지 않는—에 의존하는 것이 여전히 표준으로 남아 있습니다. 분자 확산 모델을 위한 샘플 품질 신호를 제공하는 유사한 연구가 부족한 실정입니다.

방법론
저자들은 생성기를 재학습시키지 않고도 사전 학습된 분자 확산 모델에 대해 사후(post-hoc) 방식으로 샘플별 불확실성을 추정하는 방법을 제안합니다. 이 접근 방식은 모델이 내부적으로 불확실함을 느끼는 샘플들이 더 가변적인 디노이징 궤적(denoising trajectories)을 보일 것이라는 직관에 기반합니다.

1. 라플라스 근사(Laplace Approximation): 어떤 사전 학습된 모델과도 호환성을 유지하기 위해, 이 방법은 디노이징 네트워크의 파라미터에 라플라스 근사를 적용합니다. 구체적으로, 사전 학습된 파라미터 $\hat{\theta}$ 주변에 $q(\theta) \approx \mathcal{N}(\theta | \hat{\theta}, \Sigma)$ 형태의 사후 분포를 구축합니다. 공분산 $\Sigma$는 확장성을 보장하기 위해 최종 레이어에만 적용된 경험적 피셔 정보 행렬(empirical Fisher information matrix)과 사전 정밀도(prior precision) 항을 사용하여 근사합니다.
2. 노이즈 예측 가변성(Noise Prediction Variability): 생성 궤적 동안 선택된 타임스텝 $\mathcal{T}$의 부분 집합에 대해, 방법론은 사후 분포 $q(\theta)$로부터 $M$개의 가중치 샘플 $\theta_m$을 추출합니다. 각 샘플에 대해 디노이저는 노이즈 $\epsilon^m_t = f_{\theta_m}(\hat{x}_t, t)$를 예측합니다.
3. 불확실성 점수 집계: 각 타임스텝에서의 중간 불확실성 $u_t$는 $M$개의 샘플에 대한 노이즈 예측값들의 요소별(element-wise) 샘플 분산으로 정의됩니다. 이 단계별 불확실성은 선택된 타임스텝, 모든 원자($N_a$), 그리고 모든 특징 차원($D$)에 걸쳐 집계되어 생성된 분자에 대한 단일 스칼라 불확실성 점수 $u(\hat{x})$를 산출합니다:
   $$u(\hat{x}) := \frac{1}{|\mathcal{T}| \cdot N_a \cdot D} \sum_{t \in \mathcal{T}} \sum_{n=1}^{N_a} \sum_{d=1}^{D} [u_t]_{n,d}$$
   경험적으로 저자들은 생성 과정의 마지막 10%(즉, $t$가 작은 시점)에서만 불확실성을 측정하는 것이 성능 측면에서 충분하며, 이를 통해 계산 오버헤드를 제한할 수 있음을 발견했습니다.

주요 기여

• 분자 확산에 대한 최초 적용: 본 연구는 확산 불확실성 접근법을 이미지 도메인에서 분자 생성으로 확장하여, 실험실 검증과 같이 비용이 많이 드는 과학적 도메인에서의 샘플 품질 신호 필요성을 해결했습니다.
• 사후(Post-hoc) 프레임워크: 이 방법은 기존의 EDM이나 GeoLDM과 같은 최첨단 아키텍처에 적용 가능하도록 기초가 되는 확산 모델을 재학습시킬 필요가 없습니다.
• 테스트 타임 스케일링(Test-Time Scaling): 본 논문은 불확실성 점수를 오버샘플링 및 필터링을 통한 테스트 타임 스케일링 설정에서 어떻게 품질이 높은 분자를 선별하는 기준으로 사용할 수 있는지 보여줍니다.

결과
저자들은 EDM 및 GeoLDM 모델을 사용하여 QM9 및 GEOM-Drugs 데이터셋에 대해 접근 방식을 평가했습니다.

• 품질와의 상관관계: QM9 데이터셋에서 제안된 불확실성 점수는 확립된 품질 지표(분자 안정성, 원자 안정성 및 유효성)와 통계적으로 유의미한 음의 상관관계를 보였습니다. 이는 불확실성 점수가 높을수록 품질이 낮은 샘플임을 나타냅니다.
• 가능도(Likelihood)와의 비교: 불확실성 점수는 샘플 품질의 예측 인자로서 확산 모델의 음의 로그 가능도(NLL)보다 일관되게 우수한 성능을 보였습니다. 예를 들어, EDM/QM9에서 불확실성 점수는 분자 안정성과 $-0.284$의 스피어먼 상관계수를 기록한 반면, NLL은 $-0.150$을 기록했습니다.
• 테스트 타임 스케일링 성능: 2만 개의 분자를 생성하고 불확실성이 가장 낮은 1만 개를 남기도록 필터링함으로써, 저자들은 EDM 모델에서 분자 안정성($\sim10\%$), 원자 안정성($\sim1\%$), 유효성($\sim5\%$)의 상당한 개선을 관찰했습니다. 이러한 이득은 기본 모델 아키텍처를 EDM에서 GeoLDM으로 교체했을 때 얻은 개선 효과와 대등한 수준이었습니다.
• 절제 연구(Ablation Studies):
  • 피셔 근사(Fisher Approximation): 피셔 근사를 제거하고(공분산을 등방성 가우시안으로 대체) 수행했을 때 거의 동일한 성능을 보였으며, 이는 해당 점수가 엄격한 베이지안 인식론적 불확실성이라기보다는 작은 섭동에 대한 민감도 측정치로서 작동함을 시사합니다.
  • 타임스텝 선택: 신호는 생성 궤적의 "깨끗한" 끝부분(작은 $t$)에서 가장 강력하게 나타났으며, 이는 불확실성 측정을 마지막 단계에 집중하는 선택을 정당화합니다.
• 한계점: 본 방법론은 더 크고 복잡한 GEOM-Drugs 데이터셋에서는 의미 있는 개선을 이끌어내지 못했으며, 이는 해당 접근 방식이 현재 높은 구조적 복잡성을 가진 데이터셋에는 전이되지 못함을 나타냅니다.

의의 및 주장
본 논문은 확산 불확실성 접근법을 분자 도메인으로 확장하는 "첫걸음"을 뗐다고 주장합니다. 주요 의의는 생성된 분자를 필터링하기 위한 실용적인 사후 도구(post-hoc tool)를 제공함으로써, 다운스트림 평가 파이프라인의 효율성을 높이는 데 있습니다. 저자들은 이 불확실성 점수를 적응형 테스트 타임 스케일링 방법에서 NLL 기반 검증기를 대체할 자연스러운 대안으로 위치시킵니다.

저자들은 자신의 주장에 대해 겸허한 태도를 유지하며, 이 접근 방식이 현재 더 크고 복잡한 데이터셋(GEOM-Drugs)에서는 실패한다는 점과, 근본적인 메커니즘이 진정한 베이지안 인식론적 불확실성이라기보다는 민감도 측정치에 가깝다는 점을 명시적으로 언급했습니다. 또한 향후 연구에서 고차원 이미지 생성에서의 최근 성공 사례와 유사하게, 사전 학습된 분자 인코더를 사용하여 "의미론적(semantic)" 특징 공간에서 불확실성을 계산하는 방향을 탐구해야 한다고 제안합니다.