[논문 리뷰] Implicit Transfer Operator Learning: Multiple Time-Resolution Surrogates for Molecular Dynamics

Mathias Schreiner, Ole Winther외 1명|arXiv (Cornell University)|2023. 05. 29.

Machine Learning in Materials Science인용 수 8

한 줄 요약

ITO는 SE(3)-등가 확산 모델을 사용하여 분자 동역학의 다중 시간 해상도 대리 모델을 학습하고, 자가 일관된 장시간 역학과 규모 간 효율적 샘플링을 가능하게 한다.

ABSTRACT

Computing properties of molecular systems rely on estimating expectations of the (unnormalized) Boltzmann distribution. Molecular dynamics (MD) is a broadly adopted technique to approximate such quantities. However, stable simulations rely on very small integration time-steps ($10^{-15}\,\mathrm{s}$), whereas convergence of some moments, e.g. binding free energy or rates, might rely on sampling processes on time-scales as long as $10^{-1}\, \mathrm{s}$, and these simulations must be repeated for every molecular system independently. Here, we present Implict Transfer Operator (ITO) Learning, a framework to learn surrogates of the simulation process with multiple time-resolutions. We implement ITO with denoising diffusion probabilistic models with a new SE(3) equivariant architecture and show the resulting models can generate self-consistent stochastic dynamics across multiple time-scales, even when the system is only partially observed. Finally, we present a coarse-grained CG-SE3-ITO model which can quantitatively model all-atom molecular dynamics using only coarse molecular representations. As such, ITO provides an important step towards multiple time- and space-resolution acceleration of MD. Code is available at \href{https://github.com/olsson-group/ito}{https://github.com/olsson-group/ito}.

연구 동기 및 목표

• 분자 시스템에서 볼츠만 분포의 기댓값을 추정할 필요성을 동기 부여하고, 미세한 MD 시간단계를 넘어서는 장시간 스케일 동역학의 도전 과제를 다룬다.
• MD 전이 밀도들의 다중 시간 해상도 대리모음을 학습하기 위한 암시적 전달 연산자(ITO) 프레임워크를 제시한다.
• 공간 대칭성에 걸쳐 물리적 일관성을 유지하기 위해 SE(3)-등가 생성 모델을 개발한다.
• ITO가 다중 시간 스케일에 걸쳐 자가 일관된 확률적 동역학을 생성하고, 거칠게 표현된 형태로도 모든 원자 MD를 거친 표현으로 근사할 수 있음을 입증한다.

제안 방법

• ITO를 조건부 노이즈 제거 확산 확률 모델(cDDPM)을 통해 조건부 전이 밀도 p_{Nτ}(x_{Nτ}|x_0)를 학습하는 것으로 형식화한다.
• 3D 회전 및 평행이동에 대한 불변/등가성을 보장하기 위해 SE(3)-등가 아키텍처(ChiroPaiNN)를 사용한다.
• 전이 확률을 시간 가변 성분과 시간 불변 성분으로 분해하고 전달 연산자의 고유함수 투영을 통해 훈련을 여러 lagtime(Nτ)과 정렬한다.
• 고유함수 표현을 더 잘 학습하기 위해 분포에서 샘플링된 Lag time의 분포에 노출시키는 데이터 증강 전략으로 학습한다.
• 분자계에 대한 SE3-ITO와 거칠게 표현된 단백질 접힘 데이터용 CG-SE3-ITO의 두 가지 아키텍처를 제공하며, 둘 다 확산 기반 스코어 모델링과 SE(3) 불변성을 보장하기 위한 불변 사전(invariant prior)을 사용한다.

실험 결과

연구 질문

• RQ1 MD 데이터에서 학습된 암시적 전달 연산자 대리모가 다중 시간 스케일에 걸쳐 자가 일관된 확률적 동역학을 생성할 수 있는가?
• RQ2 스톡캐스틱-lag 훈련 전략이 고정-lag 훈련에 비해 메타-안정성 포착과 장시간 동역학을 개선하는가?
• RQ3 SE(3)-등가 ITO 모델이 핵심 역학적 및 정상 관측치를 보전하면서 거칠게 표현된 표현으로 확장할 수 있는가?
• RQ4 ITO 대리모가 재가중 없이 동역학 관측치(예: 접힘/풀림 시간)와 정상 관측치(예: 자유 에너지)를 얼마나 잘 재현하는가?

주요 결과

• ITO 모델은 스톡캐스틱 lag 훈련에서 고정-lag 모델보다 Müller–Brown 벤치마크에서 메타-안정성 예측에 우수하다.
• SE3-ITO는 길게 지속되는 시간 스케일에 대해 자가 일관된 장시간 동역학을 제공하고 알라닌 디페펩타이드에 대해 MD 데이터와 다중 시간 스케일에서 일치한다.
• CG-SE3-ITO는 빠른 접힘 단백질(Chignolin, Trp-Cage, BBA, Villin)에서 주요 동역학적 및 정상 관측치를 긴 모든 원자 MD에 비해 재현한다.
• ITO를 이용한 조상 샘플링은 MD와 느리게 릴랙스하는 좌표에서 일치하는 장기 궤적(마이크로초까지)을 생성할 수 있다.
• CG-SE3-ITO에서 계산된 접힘 자유 에너지 및 평균 첫 통과 시간과 같은 관측값이 참조 MD 통계와 보통의 변동성 내에서 일치한다.
• 프레임워크는 GPU에서 높은 샘플링 속도를 달성하면서 MD 궤적과의 정성적 일치를 유지하는 상당한 효율 향상을 보여준다.