전이 가능한 정규화 흐름을 이용한 암묵적 샘플링

초록

PROSE는 285 백만 파라미터 규모의 정규화 흐름 모델로, 2~8잔류 펩타이드 MD 궤적을 학습해 시퀀스 길이와 아미노산 조합이 다른 새로운 펩타이드에 대해 제로샷으로 무상관 샘플을 생성한다. SNIS 기반 중요도 샘플링과 순차적 몬테카를로에 제안 분포로 활용해 기존 MD와 비교해 GPU 시간당 4 천 배 빠른 샘플링을 달성한다.

상세 분석

본 논문은 기존 Boltzmann Generator(BG)와 Transferable Boltzmann Generator(TBG)의 한계를 극복하기 위해, Transformer 기반의 Autoregressive Normalizing Flow인 Tarflo를 확장한 PROSE 아키텍처를 제안한다. 첫 번째 핵심은 가변 길이 펩타이드를 동시에 학습할 수 있도록 마스킹과 패딩 무시 로그‑디터미넌트 집계를 도입한 점이다. 이는 입력 차원이 시스템마다 달라지는 상황에서도 동일한 파라미터 집합을 재사용하게 해, 학습 효율과 전이성을 크게 향상시킨다. 두 번째로, 고정된 위치 임베딩 대신 주기적 특성을 가진 사인파 위치 임베딩을 사용해 시퀀스 길이 외삽에 강인한 특성을 부여한다. 세 번째는 시스템‑조건부 정보를 단순 덧셈이 아닌, Transformer 레이어 내부에서 다중‑헤드 어텐션을 통해 아키텍처 전반에 걸쳐 통합함으로써 아미노산 종류, 잔기 위치, 온도 등 다양한 물리적 변수에 대한 정밀한 조절이 가능하도록 설계했다.

학습 목표는 정규화 흐름의 로그‑우도 최대화이며, 이는 실제 MD 궤적에서 추출한 좌표 분포와의 KL‑다이버전스를 최소화한다. 학습 후에는 제안 분포 qθ(x)를 이용해 자기‑정규화 중요도 샘플링(SNIS)으로 목표 Boltzmann 분포 p(x)∝exp(−U(x)/kT)를 추정한다. SNIS는 샘플 수가 충분히 크면 편향이 사라지는 일관성을 제공하므로, 별도 MCMC 파라미터 튜닝 없이도 정확한 자유 에너지와 구조적 메트릭을 얻을 수 있다.

실험에서는 21 700개의 펩타이드 시퀀스를 포함하는 ManyPeptidesMD 데이터셋(총 4.3 ms MD 시뮬레이션)으로 사전 학습하고, 길이 48의 새로운 펩타이드 30종에 대해 제로샷 성능을 평가했다. Wasserstein‑2 거리 기반 에너지, 이디히얼 토러스, TICA 투영 메트릭에서 기존 1 µs MD와 비교해 동일 GPU 시간 내에 45배 낮은 오류를 기록했으며, 특히 메타안정 상태 전이를 포착하는 T‑W2와 TICA‑W2에서 현저히 우수했다. 또한, 동일 모델을 순차적 몬테카를로(SMC)와 결합해 샘플 효율을 더욱 향상시켰으며, 사전 학습된 PROSE를 새로운 시스템에 대해 몇 천 번의 SNIS 샘플만으로 역전파 기반 역 KL‑다이버전스 미세조정(fine‑tuning)할 수 있음을 보였다.

비교 대상인 TBG는 연속 정규화 흐름(CNF) 기반으로, 벡터 필드 적분과 발산 계산에 높은 비용이 소요돼 4 GPU‑day가 필요했지만, PROSE는 1 GPU‑hour 내에 3 × 10⁴ 샘플과 로그우도를 제공한다. 이는 제안·평가 단계 모두에서 10³~10⁴ 배의 속도 향상을 의미한다.

한계점으로는 현재 8잔류 이하의 짧은 펩타이드에 국한된 점, 그리고 전이성 검증이 주로 동일 화학 환경(수용액, 동일 온도)에서 이루어진 점을 들 수 있다. 향후에는 더 큰 단백질 도메인, 다양한 용매·조건에 대한 일반화, 그리고 역동적 적응형 샘플링 전략과의 결합이 기대된다.