기술 요약: 대리 가능도 추정기를 이용한 확장 가능한 추론 시간 어닐링 (SITA)

문제 정의

볼츠만 분포 $\pi(x) \propto \exp(-E(x)/k_B T)$로 정의되는 분자 구성의 평형 앙상블을 샘플링하는 것은 통계 물리학 및 계산 화학에서 기초적이면서도 계산적으로 다루기 어려운 과제입니다. 분자 역학(MD) 및 마르코프 연쇄 몬테카를로(MCMC)와 같은 전통적인 방법들은 펨토초 단위의 타임스텝과 국소 에너지 최솟값에 갇히는 경향성으로 인해 높은 계산 비용 문제를 겪습니다.

딥 생성 모델은 아마티제이션(amortized)된 빠른 샘플링을 위한 경로를 제시하지만, 이들은 학습 데이터로서 평형 앙상블을 필요로 한다는 "순환적"인 학습 한계에 직면합니다. 이는 정확히 생성하기 어려운 대상인 평형 앙상블을 사용하여 모델을 반복적으로 재학습시키는 "부트스트래핑(bootstrapping)" 접근법을 통해 해결하려고 시도합니다. 그러나 확산 모델(예: PITA)에 기반한 최근의 방법들은 전체 역시간 적분 경로를 따라 스코어 필드(score field)의 발산(divergence)을 계산해야 하는 자기 정규화 중요도 샘플링(SNIS)에 의존하며, 이 발산 계산은 시스템 차원에 따라 기하급급하게 증가하여 대규모 분자 시스템에서는 실행 불가능해집니다.

방법론: SITA

저자들은 발산 계산의 계산적 병목 현상을 우회하기 위해 연속 흐름 매칭(continuous flow-matching) 모델과 대리 가능도 추정기(surrogate likelihood estimator)를 결합한 SITA (Scalable Inference-Time Annealing) 프레임워크를 제안합니다.

핵심 구성 요소

1. 연속 흐름 매칭 모델 (Continuous Flow-Matching Models): SITA는 기저 가우시안 분포로부터 타겟 분포로 샘플을 운송하는 연속 시간 프로세스를 정의하기 위해 확률적 보간법(stochastic interpolants)을 활용합니다. 확산 모델과 달리, 흐름 매칭은 회귀를 통해 속도장 $v_t(x)$를 학습하며, 명시적인 노이즈 스케줄을 피합니다.
2. 대리 가능도 추정기 (BoltzNCE): 흐름 기반 재가중화에 필요한 다루기 힘든 발산 항($\nabla \cdot v_t$)을 계산하는 것을 피하기 위해, SITA는 대리 가능도로서 매개변수화된 에너지 기반 모델(EBM)을 사용하는 BoltzNCE를 채택합니다. 이 EBM은 흐름의 출력 분포에 대한 에너지 함수 $U_\phi$를 학습하여 다룰 수 있는 밀도 추정치 $q_\phi(x)$를 제공합니다.
3. 온도 스티어링 (Temperature Steering): SITA는 흐름 모델에 내재된 온도 인코딩을 활용합니다. 추론 시 기저 가우시안 분포의 분산을 재조정함으로써 ($\kappa = T_{high}/T_{low}$인 $x_0 \sim \mathcal{N}(0, \kappa^{-1}I$), 아키텍처 변경이나 명시적인 온도 조건화 없이도 저온의 샘플을 생성할 수 있습니다.

SITA 알고리즘

이 방법은 감소하는 온도 사다리 $\{T_k\}$에 대해 반복적인 부트스트랩 루프를 통해 작동합니다:

1. 흐름 어닐링 (Anneal the Flow): 재조정된 기저 분포를 사용하여 온도 $T_{k+1}$에서의 샘플을 생성하기 위해 흐름 ODE를 적분합니다.
2. EBM 미세 조정 (Finetune the EBM): 생성된 흐름 샘플을 사용하여 BoltzNCE 목적 함수를 통해 EBM 파라미터 $\phi$를 업데이트하고, 대리 가능도 $q_\phi$를 정교화합니다.
3. 중요도 샘플링 (Importance Sampling): 중요도 가중치 $w(x) \propto \exp(-E(x)/k_B T_{k+1}) / q_\phi(x)$를 계산합니다. 흐름 출력을 이 가중치로 재샘플링하여 $T_{k+1}$을 목표로 하는 데이터셋을 만듭니다.
4. 흐름 미세 조정 (Finetune the Flow): $T_{k+1}$에서의 타겟 분포를 더 잘 근사하도록 재가중된 데이터셋을 사용하여 흐름 모델 $\theta$를 다시 학습시킵니다.

이 과정은 타겟 온도(예: 300K)에 도달할 때까지 반복됩니다. 이 방법은 EBM의 대리 가능도를 활용하여 가중치를 계산함으로써 값비싼 자코비안(Jacobian) 계산을 피합니다.

주요 기여

• 흐름을 위한 최초의 추론 시간 어닐링: SITA는 연속 시간 흐름 모델에 추론 시간 어닐링을 적용한 최초의 방법으로, 사전 정의된 사다리를 가로지르는 큰 온도 변화를 가능하게 합니다.
• 대리 기반 중요도 샘플링: BoltzNCE 스타일의 대리물을 통합함으로써 고온 분포에서 상온 타겟으로의 중요도 가중치 기반 운송을 가능하게 하며, 이는 기존의 확산 기반 부트스트래핑 방법에서 사용되는 경로 적분 기반 SNIS 추정기의 과도한 비용 문제를 우회합니다.
• 확장성: 발산 계산을 다룰 수 있는 EBM으로 대체함으로써, SITA는 기존의 어닐링 방식이 실패하는 많은 자유도를 가진 시스템으로 확장 가능합니다.

실험 결과

저자들은 알라닌 디펩타이드 (ADP) 및 **알라닌 트리펩타이드 (ATP)**라는 두 가지 벤치마크 분자 시스템에서 SITA를 평가하였으며, PITA, 온도 어닐링 볼츠만 생성기 (TA-BG), 그리고 300K MD 데이터로 직접 학습된 모델 (MD-Diff, MD-NF)과 비교하였습니다.

• 알라닌 디펩타이드 (Alanine Dipeptide): SITA는 라마카란드라 KL 발산 (Rama-KL) 및 에너지 Wasserstein-2 지표에서 최고 수준의 성능을 달성하였으며, PITA를 능가했습니다. MD-NF가 더 낮은 Energy-W1를 보였으나 모드 붕괴(mode collapse, 낮은 Rama-KL)의 징후를 보인 반면, SITA는 모든 주요 구조적 베이신(conformational basins)을 성공적으로 포착했습니다.
• 알라닌 트리펩타이드 (Alanine Tripeptide): SITA는 사후 MD 정밀화(post-hoc MD refinement) 없이도 거의 모든 지표(Rama-KL, Energy-W1, Energy-W2)에서 모든 베이스라인을 능가했습니다. 반면, PITA와 TA-BG는 경쟁력을 유지하기 위해 짧은 MD 정밀화 단계가 필요했습니다.
• 효율성: 초기 학습 데이터 생성 비용을 제외하면, 부트스트래핑 단계에서 베이스라인 방법들에 비해 1~2 자릿수 적은 에너지 평가를 소모했습니다.
• 정밀화 (Refinement): 저자들은 학습된 대리물을 사용하여 독립 메트로폴리스-헤이스팅스(IMH)를 적용하는 것이 샘플 다양성을 유지하면서도 (특히 Rama-KL 및 T-W2 측면에서) 지표를 더욱 개선함을 입증했습니다. 다만, 이는 추가적인 에너지 평가를 수반합니다.

의의 및 주장

본 논문은 SITA가 흐름 기반 어닐링의 발산 항 계산의 불가능성을 해결함으로써 분자 시스템의 볼츠만 분포 샘플링을 위한 새로운 최첨단(state-of-the-art)을 구축한다고 주장합니다. 저자들은 대리 가능도 추정기가 많은 자유도를 가진 분자 앙상블을 모델링하는 데 있어 실용적인 경로를 제공하며, 이는 기존의 확산 기반 부트스트래핑 방법들이 모드 붕괴나 계산 비용으로 인해 어려움을 겪는 영역임을 강조합니다.

이 연구는 대리 가능도 사용이 이론적 편향(정확한 타겟이 아닌 기울어진 분포로 수렴함)을 도입할 수 있지만, 벤치마크 시스템에 대한 실험 결과는 이러한 편향이 평형 앙상블을 포착하는 데 있어 우수한 성능을 저해하지 않음을 보여준다고 밝힙니다. 저자들은 SITA가 연속 시간 흐름의 아키텍처적 이점을 유지하면서 저온 타겟을 향한 효율적이고 반복적인 정밀화를 가능하게 하는 확장 가능한 대안임을 제시합니다.