[Paper Review] Implicit Transfer Operator Learning: Multiple Time-Resolution Surrogates for Molecular Dynamics

Mathias Schreiner, Ole Winther+1 more|arXiv (Cornell University)|May 29, 2023

Machine Learning in Materials Science8 citations

TL;DR

ITO learns multiple time-resolution surrogates for molecular dynamics using SE(3)-equivariant diffusion models, enabling self-consistent long-time dynamics and efficient sampling across scales.

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}.

Motivation & Objective

• Motivate the need to estimate expectations of the Boltzmann distribution in molecular systems and address the challenge of long-time-scale dynamics beyond tiny MD time-steps.
• Introduce the Implicit Transfer Operator (ITO) framework to learn multi-time-resolution surrogates of MD transition densities.
• Develop SE(3)-equivariant generative models to maintain physical consistency across spatial symmetries.
• Demonstrate that ITO can generate self-consistent stochastic dynamics across multiple time scales and, in coarse-grained form, approximate all-atom MD using coarse representations.

Proposed method

• Formulate ITO as learning a conditional transition density p_{Nτ}(x_{Nτ}|x_0) via a conditional denoising diffusion probabilistic model (cDDPM).
• Use an SE(3)-equivariant architecture (ChiroPaiNN) to ensure invariance/equivariance under 3D rotations and translations.
• Decompose the transition probability into time-variant and time-invariant components through eigenfunction projections of the Transfer operator, aligning training with multiple lagtimes (Nτ).
• Train with a data-augmentation strategy exposing the model to a distribution of lag times (N sampled from a distribution) to better learn eigenfunction representations.
• Provide two architectures: SE3-ITO for molecular systems and CG-SE3-ITO for coarse-grained protein folding data, both using diffusion-based score modeling and an invariant prior to guarantee SE(3) invariance.

Experimental results

Research questions

• RQ1 Can an implicit transfer operator surrogate learned from MD data generate self-consistent stochastic dynamics across multiple time scales?
• RQ2 Do stochastic-lag training strategies improve meta-stability capture and long-time dynamics compared to fixed-lag training?
• RQ3 Can SE(3)-equivariant ITO models scale to coarse-grained representations while preserving key dynamic and stationary observables?
• RQ4 How well do ITO surrogates reproduce dynamic observables (e.g., folding/unfolding times) and stationary observables (e.g., free energy) without explicit reweighting?

Key findings

• ITO models with stochastic lag training outperform fixed-lag models in meta-stability prediction on Müller–Brown benchmarks.
• SE3-ITO yields self-consistent long-time dynamics and agrees with MD data across multiple time scales for alanine dipeptide.
• CG-SE3-ITO on fast-folding proteins (Chignolin, Trp-Cage, BBA, Villin) reproduces major dynamic and stationary observables compared to long all-atom MD.
• Ancestral sampling with ITO can generate long trajectories (up to microseconds) that align with MD in slowly relaxing coordinates.
• Observables such as folding free energy and mean first passage times computed from CG-SE3-ITO align with reference MD statistics (within reported variability).
• The framework demonstrates substantial efficiency gains, achieving high sampling rates on GPUs while maintaining qualitative agreement with MD trajectories.