Title: Implicit Transfer Operator Learning: Multiple Time-Resolution Models for Molecular Dynamics
Venue: NeurIPS 2023
Accepted: yes
Keywords: AI4Science

Molecular Dynamics

equivariant neural networks

stochastic dynamics
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 Implicit 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}.

Implicit Transfer Operator Learning: Multiple Time-Resolution Surrogates for Molecular Dynamics

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 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 s , and these simulations must be repeated for every molecular system independently. Here, we present Implicit 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 https://github.com/olsson-group/ito.

1 Introduction

Numerical simulation of stochastic differential equations (SDE) is critical in the sciences, including statistics, physics, chemistry, and biology applications [1]. Molecular dynamics (MD) simulations are an important example of such simulations [2]. These simulations prescribe a set of mechanistic rules governing the time evolution of a molecular system through numerical integration of, for example, the Langevin equation [3]. MD grants mechanistic insights into experimental observables. These observables are expectations, including time-correlations, of observable functions (e.g., pairwise distances or angles) computed for the Boltzmann distribution ˆ µ ( x ) ∝ exp[ -β U ( x )] corresponding to the potential U ( · ) : Ω → R of a M -particle molecular system, x ∈ Ω ⊂ R 3 M kept at the inverse temperature β = 1 /kT . However, stable numerical integration relies on time steps, τ , which are strictly smaller than the fastest characteristic time-

T

Ω

(1)

T

Ω

(

N

)

Time

μ

p

(

Ω

)

Ω

̂

Chalmers University of Technology simonols@chalmers.se

scales of the molecular system ( 10 -15 s , e.g., bond vibrations), yet many molecular systems are characterized by processes on much longer time-scales ( 10 -3 -10 -1 s , e.g. protein-folding, proteinligand unbinding, regulation). Consequently, we need infeasibly long simulations to characterize many important processes quantitatively due to the slow mixing in Ω .

In this work, we present the implicit Transfer Operator (ITO, Fig. 1) as an effective way to learn multiple time-step surrogate models of the stochastic generating distribution of MD. To our knowledge, this is the first surrogate modeling approach that allows for the simultaneous generation of stochastic dynamics at multiple different time resolutions. By adopting an SE(3)-equivariant generative model, we further demonstrate stable long-time-scale dynamics in increasingly difficult settings where an increasing number of degrees of freedom are marginalized. Our approach can be several orders of magnitude more efficient than direct MD simulations and can be made asymptotically unbiased if the generative model permits exact likelihood evaluation. Our current results do not generalize across different thermodynamic ensembles or across chemical space, but show strong generalization across different time-scales.

Our main contributions are

the Implicit Transfer Operator (ITO) framework for learning generative models for multiple time resolution molecular dynamics simulations,

implementation of ITO using a denoising diffusion probabilistic model (DDPM) [4] with strong empirical results across resolutions: SE(3)-equivariant ITO model (SE3-ITO) gives stable long time-scale simulations and self-consistent dynamics across multiple time-scales for molecular benchmarks and Coarse-grained SE3-ITO model (CG-SE3-ITO) trained on large-scale protein folding data sets shows quantitative agreement with major dynamic and stationary observables of interest.

2 Background and Preliminaries

Notation Throughout this work, diffusion time, related to Diffusion Models (see Sec. 2), and physical time are represented using superscripts and subscripts, respectively.

Molecular dynamics and observables Molecular dynamics (MD) is a wide-spread simulation strategy in computational chemistry and physics. In this approach, the time-evolution of N particles configuration in Euclidean space x ∈ Ω ⊂ R 3 M , is modeled via a stochastic differential equation (SDE) with a drift term based on a potential energy model U ( x ) : Ω → R . An important aim of MD is to compute:

Stationary observables: O f = E µ [ f ( x )]

Dynamic observables: O f ( t ) h ( t + ∆ t ) = E x t ∼ µ [ E x t + ∆ t ∼ p τ ( x t + ∆ t | x t ) [ f ( x t ) h ( x t + ∆ t )]]

where µ is the normalized Gibbs measure, and p τ ( x t + ∆ t | x t ) is a conditional probability density function encoding the time-discrete evolution of the molecular system x , with time-step ∆ t = N τ as prescribed by a dynamic model, e.g. Langevin dynamics [3], integrated with time-step, τ . N is typically a large integer. The functions f, h : Ω → R are observable functions or ' forward models ' describing the microscopic observation process, e.g. computing a distance or an angle. The observables, O f ( t ) , and O f ( t ) h ( t + ∆ t ) , include binding affinities and binding rates of a drug to a protein, respectively. Conventionally, these observables are estimated from simulation trajectories using naive Monte Carlo estimators.

For illustrative purposes, we assume the temporal behavior of a state, x , follows the Brownian dynamics SDE (Itô form)

where D = γ -1 β -1 is a diffusion constant, with friction γ and inverse-temperature β , and d W is a Wiener process. Using the Euler-Maruyama time-discretization, with time-step τ , simulating the SDE corresponds to simulating a Markov chain with the transition probability density

where N specifies the multi-variate Normal distribution, and I 3 M is the 3 M -dimensional identity matrix. If τ is sufficiently small to allow stable simulation, the invariant measure , of the Markov chain

(eq. 2), is the Boltzmann distribution (normalized Gibbs measure) corresponding to the potential energy model U ( x ) at β . Consequently, by simulating a large number of steps we can draw samples from µ to compute stationary observables and compute dynamic observables by simulating ∆ t = N τ steps enough times with initial states distributed according to µ . Explicit simulation make such computations extremely costly, and consequently, there's much interest in speeding up the calculations of these quantities.

Transfer Operators Let ρ specify an initial condition, a probability density function on Ω . We can define a Markov operator T Ω : L 1 ( Ω ) → L 1 ( Ω ) using a transition density (e.g., 2):

which then describes the µ -weighed evolution of absolutely convergent probability density functions on Ω according to eq. 1, with time-step, τ . Such an operator is called the (Ruelle) Transfer Operator [5, 6]. We can express the operator using a spectral form

where only eigenvalues λ i ( τ ) = exp( -τκ i ) depend on the time-step, τ . κ i are characteristic 'relaxation' rates associated the left and right eigenfunction pair, φ i and ψ i [7]. We can compute the operator with time-lag N τ via the Chapman-Kolmogorov equation (see Sec. A.1, for details)

Equivariant Message Passing Neural Networks In this work, we are concerned with MD, where the time-evolution of a molecule is governed by a force field F ( · ) /defines -∇ U ( · ) derived from a central potential U ( · ) . While U ( · ) is invariant to group-actions of the Euclidean group in three dimensions (E(3)), its corresponding force field is E(3)equivariant . We call a function, f ' invariant ' under a group-action g iif f ( x ) = f ( S g x ) and ' equivariant ' iff T g f ( x ) = f ( S g x ) , where S g and T g are linear representations of the group element g [8].

The force field F ( · ) is equivariant under E(3) group-actions. However, in practice, classical molecular dynamics simulations do not change parity during simulation, and consequently, our data distribution only contains a single mirror image of molecules.

We extended the PaiNN architecture [9], an E(3)-equivariant message passing neural network (MPNN), making it SE(3) equivariant by breaking its symmetry with respect to parity. We introduced this minor modification as we experienced sporadic parity changes when sampling with a model trained using the PaiNN architecture, and introducing this modification resolved the issue. Briefly, PaiNN embeds a graph G = ( V, E ) , where nodes, V , exchange equivariant messages through edges within a local neighborhood defined as N ( i ) = { j | ‖ r ij ‖ ≤ r cuto ff } , where r ij is the distance between nodes denoted i and j , and r cuto ff is the maximal distance at which nodes are allowed to exchange messages. Messages are pooled and subsequently used to update node features, thereby enabling exchange of equivariant information. We achieve parity symmetry-breaking by constructing the equivariant messages in a manner that depends on cross-products between equivariant node features and direction vectors between interacting nodes. The cross-product is an axial vector (i.e., does not change sign under parity). We combine these vectors with polar vectors (change sign under parity). We refer to this modified PaiNN architecture as ChiroPaiNN (CPaiNN). Further details are in the Appendix D.

Diffusion Models The diffusion model (DM) formalism is a powerful generative modeling framework that learns distributions by modeling a gradual denoising process [4, 10, 11]. In DMs, we prespecify a forward diffusion process (noising process), which gradually transforms the data distribution p ( x 0 ) to a simple prior distribution p ( x T ) , e.g., a standard Gaussian, through a time-inhomogenous Markov process, described by the following SDE (Itô form)

where 0 < t < T is the diffusion time , f and g are chosen functions, and d W is a Wiener process. We can generate samples from the data distribution p ( x 0 ) by sampling from p ( x T ) and solving the backward diffusion process (denoising process)

by approximating the score field ∇ x t log p ( x t | t ) - or equivalently a time-dependent Normal transition kernel [4] - with a deep neural network surrogate ∇ x t log ˆ p ( x t | t, θ ) . We can use the learned score field to define a neural ordinary differential equation (ODE) [12, 13], or probability flow ODE [14] - eq. 7 less the term g ( t ) d W and scaling g 2 ( t ) by 1 / 2 - which we can leverage for efficient sampling and sample likelihood evaluation.

Here, we are concerned with building equivariant probability density functions under SE(3) group actions. Consequently, we parameterize the DM using a learned Normal transition kernel of a time-inhomogenous diffusion process. By restricting the transition kernels p ( x t +1 | x t ) to be equivariant under SE(3) group-actions, the marginal of x t +1 is always invariant [15]. Combining the equivariant transition kernel with an invariant prior density [16] ensures the whole Markov process is invariant to SE(3) group actions. Consequently, combining an isotropic mean-free Gaussian as prior with ChiroPaiNN-parameterized transition kernels, we can construct an SE(3) equivariant diffusion model.

ChiroPaiNN

/uni03F5

Concat

(

m

,

n

)

¯ 0

(

m

,2

n

)

(

m

,

n

)

(

m

,3)

ChiroPaiNN

z

N

(

m

,3)

(

m

,3)

SE3-ITO

(

m

(1) ,3)

(

m

)

(

m

,

n

,3)

(

m

,

n

,3)

Concat

MLP

MLP

(1)

(

m

,

n

)

(

m

,

n

)

(

m

,

n

)

(

m

,2

n

)

(

m

,

n

)

Concat

(4 + 2

n

)

N

(

m

(1) ,3)

MLP

Λ

pos

Λ

pos

(1)

(

n

)

(

n

)

(2)

/uni03F5

(2)

(2)

MB-ITO

Λ

pos

A

B

x

0

t

di

ff

˜

x

0+

N

τ

x

0

t

di

ff

˜

x

0+

N

τ

Λ

nom

Λ

pos

3 Implicit Transfer Operator

Molecular simulations are Markovian with transition density (e.g. eq. 2) and Normal, however, the latter only for very small physical time-steps τ . Here, we aim to approximate the long-time step transition probability p N τ ( x N τ | x 0 ) to

̂

allow for one-step sampling of long-time-scale dynamics.

As data we consider simulation trajectories. The trajectories are generated by explicit simulation which corresponds to sampling ancestrally from the small time-step transition density: X = { x τ , . . . , x N τ } ∼ p ( x n τ | x ( n -1) τ ) , with n = { 1 , . . . , N } . In general, the state variable x , contains both position and velocity information of the particles, along with other details such as box dimensions, depending on the simulation scheme and target ensemble. Throughout this study, we only consider the position information.

We build a surrogate of the conditional transition probability distribution p N τ ( x N τ | x 0 ) -from MDdata. In practice, we learn a generative model x t + N τ ∼ p θ ( x t + N τ | x t , N ) with a conditional denoising diffusion probabilistic model (cDDPM) of the form

where x 1: T are latent variables of the same dimension as our output, and follow a joint density describing the backward diffusion process (eq. 7) and x T ∼ N (0 , I ) . We define a conditional sample likelihood as

where I is a list of generated indices i specifying a time t i and a time-lag ( τ ) integer multiple N i , associating two time-points in the trajectory, X . Following Ho et al., we train the cDDPM by

̂

optimizing a simplified form of the variational bound of the log-likelihood [4],

˜ where ˜ x t diff t /defines √ ¯ α t diff x t + √ 1 -¯ α t diff /epsilon1 , with ¯ α t diff = ∏ t diff i (1 -β i ) and β i is the variance of the forward diffusion process at diffusion time, i . ˆ /epsilon1 θ ( · ) is one of the two ITO neural network model architectures shown in Fig. 2, and is directly related to the score [4].

If the data used to train the conditional transition density is generated by MD simulation with time-invariant potential energy (drift), we can express the generating transition probability as a decomposition of time-variant and -invariant parts (Proof, see Sec. A.2)

where α i and β i are time-invariant projection coefficients of the state variables on-to the left and right eigenfunctions φ i and ψ i , of the Transfer operator T Ω ( τ ) [5] and | λ i ( τ ) | ≤ 1 is its i 'th eigenvalue. Consequently, we call our surrogate modeling approach implicit transfer operator learning.

As outlined in Algorithm 1, we generate the indices i , e.g. the tuples ( x t i , x t i + N i τ , N i ) , in a manner such that the model is exposed to multiple time-lags, sampled uniformly across orders of magnitude, used for gradient-based optimization with Adam [17]. As a result, as illustrated in eq. 11, the model will be exposed to multiple different linear combinations of the eigenfunctions of T Ω ( τ ) in each batch during training. We conjecture that this data augmentation procedure will enable better learning of implicit representations of these eigenfunctions and, consequently, better generalization across time scales and yield more stable sampling.

3.1 ITO Architectures

We present two architectures for learning cDDPMs encoding ITO models, one for molecular applications SE3-ITO and one for the Müller-Brown benchmark system (Fig. 2). The SE3-ITO architecture uses our new SE(3) equivariant MPNN (ChiroPaiNN, described in sec. 2) to encode x t , N , and atom-types, z , to invariant features, s , and equivariant features, v . We concatenate s with an encoding of the diffusion-time t di ff and process them through a MLP (multi-layer perceptron). The output from the MLP are passed along with v and ˜ x t diff t as input to a second ChiroPaiNN module which predicts ˆ /epsilon1 . More details on the architecture and hyperparameters are available in Appendices D and E.

Algorithm 2

from

p

θ

Ancestral sampling.

Sampling is defined in Appendix E, Algorithm 4

Input:

initial condition

x

0

, lag steps

n

.

Allocate

T

∈

R

(

n

+1)

×

dim(

T

[0] =

x

for

i

= 1

x

i

∼

T

0

. . . n do

p

ˆ

[

i

θ

(

] =

end for return

T

4 Long time-step stochastic dynamics with Implicit Transfer Operators

4.1 Datasets and test-systems

To evaluate how robustly ITO models can model long time-scale dynamics, we conducted three classes of experiments, ranging from fully observed, high time-resolution, to sparsely observed and

x

T

i

[

i

-

N

, ancestral

1]

, N

x

)

0

)

low time resolution. Details on training and computational resources are available in Appendices E and F, respectively.

Müller-Brown is a 2D potential commonly used for benchmarking molecular dynamics sampling methods. We generate a training data-set by integrating eq. 1 with the Müller-Brown potential energy as U ( x ) (For details, see Appendix B.1). This dataset corresponds to a fully observed case.

Alanine dipeptide We use publicly available data from MDshare [18]. Simulation is performed with 2 fs integration time-steps and data is saved at 1 ps intervals. The simulations are performed in explicit solvation, but we only model the 22 atoms of the solute, without considering velocities. Consequently, this dataset is only partially observed.

Fast-folding proteins We use molecular dynamics data previously reported by Lindorff-Larsen et al. on the fast-folding proteins Chignolin, Trp-Cage, BBA, and Villin [19]. The data is proprietary but available upon request for research purposes. The simulations were performed in explicit solvent with a 2 . 5 fs time-step and the positions was saved at 200 ps intervals. We coarse-grain the simulation by representing each amino-acid by the Euclidean coordinate of their C α atom as done previously [20], leading to 10, 20, 28, and 35 particles in each system respectively. Consequently, these data correspond to a mostly unobserved case.

4.2 Stochastic lag improves meta-stability prediction

In sec. 2, we conjecture that exposing an ITO model to multiple lag times during training leads to better and more robust models. To test this, we trained a set of models on the Müller-Brown dataset with fixed constant lags N = { 10 , 100 , 1000 } (fixed lag) and a single model with N ∼ DisExp(1000) (stochastic lag) using the MB-ITO model (Fig. 2).

We find that the model trained with a stochastic lag systematically outperforms models trained with fixed lag (Table 1). We gauge the agreement by comparing Variational Approach to Markov Processes (VAMP)-2 scores [21] (for details, see Appendix G), between model samples and training data and find that both models tend to underestimate meta-stability compared to training data slightly. However, the model trained with stochastic lag is marginally closer to the reference values. We note that the difference in the ability of fixed and stochastic lag ITO models to capture long-time-scale dynamics is also reflected in the learned transition densities (Fig. 5). Together, these results suggest that lag-time augmentation during training leads to better implicit learning of the Transfer operator's eigenfunctions than training with a fixed lag.

To test whether this phenomena extends in cases where we do not have full observability and to molecular systems we followed the VAMP2-gaps of alanine di-peptide as a function of epoch for models trained with a fixed lag and a stochastic lag (Fig. 6). We find that the VAMP-2 gaps for stochastic lag and fixed lag in this case are statistically indistinguishable across all epochs. These results suggest that we can without compromising on accuracy build multiple time-scale surrogates by training with stochastic lag-times.

4.3 Efficient and accurate self-consistent long time-scale dynamics

We evaluate the ITO models trained with stochastic lags to capture long time-scale dynamics in a self-consistent manner, in the Chapman-Kolmogorov sense, i.e., p ( x ∆ t | x 0 ) /defines p ( x N τ | x 0 ) = ∏ N i =1 p ( x i τ | x ( i -1) τ ) , or if samples generated by direct sampling with time-step ∆ t = N τ are

distributed similarly to samples generated by performing ancestral sampling N times, each with timestep τ . In direct sampling, we draw samples for a desired time-step ∆ t = N τ from ˆ p θ ( x 0 , N ) and in ancestral sampling we draw n samples with time-step ∆ t = N τ from ˆ p θ ( x 0 , N ) in an ancestral manner, e.g. x i +1 ∼ ˆ p θ ( x i , N ) , where i = { 0 , . . . n -1 } .

For the fully-observed Müller-Brown case, we find that the ITO model is self-consistent by the strong overlap in transition densities sampled in a direct and ancestral manner (Algorithm 2). These results generalize to molecular systems and partially observed systems. Sampling an SE3-ITO model (Fig. 2) trained with alanine dipeptide data, we find strong agreement between the ancestrally and directly sampled transition densities (Fig. 3) and we again have a strong consistency with corresponding transition densities computed from molecular dynamics simulations. Note here, that the time-step of the ITO-sampled transition densities varies from 10 4 to 10 6 times the MD integration time-step. The transition densities for alanine di-peptide (Fig. 3) are calculated using 15'000 trajectories. For direct sampling, this means that we draw 15'000 samples in total. In the case of ancestral sampling, we sample 4, 64, or 512 steps for 15'000 trajectories, to match ∆ t .

Next, we consider four fast-folding proteins [19] where we coarse-grain the proteins by representing them only with their C α atoms. In this sparsely observed case (CG-SE3-ITO), we find strong model self-consistency, as shown by the comparison between conditional densities from the folded and unfolded states (Fig. 4) projected onto a linear subspace determined using time-lagged independent component analysis (time-lagged independent components, tIC) [22] (see Appendix B.3). Further, the long timescale transition density gradually converges to the data distribution as expected.

Finally, by ancestral sampling (Algorithm 2), we perform a simulation of Chignolin with the same length as the training trajectory ( 106 µs ), using a CG-SE3-ITO model, and compare with MD. The CG-SE3-ITO simulation is 2120 steps with ∆ t = 5 ns . Running in parallel, on a single Titan X GPU we can simulate the CG-SE3-ITO model at a rate of 363 ns / (s w M 2 ) where s w denotes seconds wall-time (Appendix C.1). Remarkably, these trajectories are virtually indistinguishable in the slowly

ψ

ϕ

relaxing TICA coordinates, illustrating stability of ITO. These conclusions extend to the proteins Villin, BBA, and Trp-Cage (See Appendix, Figs. 8,9 and 10)

Together these results suggest that ITO models accurately and self-consistently capture the slow dynamics of molecular systems and are robust to situations where the system is only partially observed. In general, we expect robustness to sparsely observed representations as long as the input representations are sufficient to span the eigenfunctions of T Ω [23, 24]. Approximation errors will translate into systematic under-estimation of relaxation time-scales [7], consistent with our slight under-estimation of VAMP-2 scores (Table 1). In future work, combining the learning of SE3-ITO models with a systematic scheme for coarse-graining [25, 26], could be an avenue for scaling to large-scale molecular systems at a low computational cost.

From Unfolded

From Folded

5 Prediction of dynamic and stationary observables of using CG-SE3-ITO

As outlined in section 2, an important aim of MD simulations is to compute stationary and dynamic observables, which involves intractable integrals typically approximated via Monte Carlo estimators. Using the trained ITO models we can efficiently sample i.i.d. from the transition density needed for computing dynamic observables, and by choosing a time-step which is sufficiently large we can also sample i.i.d from the Boltzmann distribution µ , the latter akin to Boltzmann generators [27] (See Appendix A.1). We note that, the ITO models are surrogates and as such without reweighing we cannot expect unbiased samples from the Boltzmann and dynamic transition densities. Nevertheless, we gauge how accurately ITO models we can compute these observables of interest in the context of protein folding without reweighing:

Free Energy of Folding , ∆ G = -log [ p f 1 -p f ]

Mean first passage time, unfolding , 〈 τ u 〉 = ∫ x 0 ∈ f ∫ ∞ 0 δ ( x t ∈ ¬ f ) p ( x t | x 0 , t ) t d t d x 0

Mean first passage time, folding , 〈 τ f 〉 = ∫ x 0 ∈ ¬ f ∫ ∞ 0 δ ( x t ∈ f ) p ( x t | x 0 , t ) t d t d x 0

where { f, ¬ f } ⊂ Ω are disjoint subsets corresponding to the folded and unfolded states of a protein, p f = ∫ x ∈ f µ ( x ) d x , is the folded state probability and δ ( · ) is the Dirac delta.

Wecompute these observables using the reference molecular simulation data [19] and sample statistics from the CG-SE3-ITO models of each of the four fast-folding proteins (details in Appendix B.6). Strikingly, the observables computed using CG-SE3-ITO models agree well with those computed from long all-atom MD simulations (Table 2).

Finally, we analyzed the robustness, convergence, and consistency of these observables (Fig. 7). For Chignolin, we trained five models independently and analyzed model checkpoints when the training loss had stabilized. For each checkpoint and each model, we computed the observables. The values predicted are statistically indistinguishable, suggesting consistency, robustness, and convergence. The

average predictions closely match the reference values. Nevertheless, we note that the fluctuations in these values are noticeable.

We implemented all experiments using PyTorch [28], PyTorch Lightning [29], JAX [30], and used DPM-Solver [31] for probability flow ODE Sampling.

6 Related Work

Molecular sampling Sampling molecular configurations is a broad field and can broadly be divided into two main areas: physically motivated sampling of the Boltzmann distribution and conformer generation. The first area includes algorithmic approaches to sample the Boltzmann distribution including Molecular Dynamics simulations [2], Markov Chain Monte Carlo, extended ensemble methods [32, 33, 34, 35], including analysis methods involving deep generative nets [36], and surrogate models which directly approximate the Boltzmann distribution and allow for recovery of unbiased statistics, including Boltzmann generators [37, 16, 38, 39]. Conformer generation concerns generating physically plausible conformers without explicitly trying to follow the Boltzmann distribution. The latter approaches can be split into ML [15, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] and chemoinformatic [50, 51] approaches. Finally, speeding up molecular simulations by reducing the effective number of particles to simulate through coarse-graining with special purpose forcefield models [52] including machine learned variants [53, 54, 20, 55] and learned coarse-graining maps [25, 26] is an orthogonal approach to sample conformation space. Further, several methods to recover all-atom models from coarse-grained representations through ML [56, 57] and rule-based approaches [58] are available.

Transfer Operator surrogates Building transfer operator surrogates is commonly used in molecular modeling including (Deep Generative) Markov state models (MSM) [59, 7, 60, 61], also including experimental data [62, 63] dynamic graphical models, [64] VAMPnets [65, 21], observable operator models[66], however, primarily for analysis of molecular dynamics data. Markov state models are time-space discrete approximations of the transfer operator and Deep Generative MSM [67] and VAMPnets [65] are deep learning infused versions, where state discretization is learned by deep nets. Dynamic graphical models reparameterize MSMs as kinetic Markov random fields allowing for scaling to larger systems [64]. Klein et al. recently introduced timewarp which is a flow-based generative model to simulate molecular systems with a large (up to 0 . 5 ns ), fixed, time-lag, [68] providing asymptotically unbiased equilibrium samples through a Metropolis-Hastings correction [69]. While timewarp generates conformers with realistic local structure, it has limitations in capturing long time-scale dynamics, which is reflected in the predicted transition probability densities. In contrast, our approach captures long time-scale dynamics efficiently allowing for accurate prediction of dynamic observables. However, currently, neither the code nor the data from timewarp is publicly available percluding direct comparisons on the benchmark tasks established in their paper.

7 Limitations

Surrogate model Implicit Transfer Operators are surrogate models of stochastic dynamics' conditional transition probability densities. We cannot guarantee unbiased sampling of dynamics and the stationary distribution due to aleatoric (e.g., finite data) and epistemic (e.g., model misspecification) uncertainty. We can overcome the latter by reweighing against a Markov Chain Monte Carlo acceptance criterion as proposed previously [68], to ensure unbiased dynamics path-reweighing is necessary, which in turn requires closed-form expressions for the target path probabilities [70].

Transferability and scalability Currently, ITO does not generalize across chemical space and thermodynamic variables. In future work, we anticipate that generalization across chemical space limitations can be overcome by appropriate data set curation and parameter-sharing schemes [71]. Generalization across thermodynamic variables such as temperature and pressure would require using a surrogate model which is steerable under these changes, e.g., temperature steerable flows [72] or thermodynamic maps [73]. Currently, we assume a fully connected graph that scales O ( M 2 ) in system size, which limits what systems are practically accessible. Devising new surrogate models which use mean-field approximation approaches from e.g., computational physics [74] or chemistry to truncate the graphs and treat long-range as an additive term [75] could yield more favorable scaling [76].

8 Conclusions

This paper introduces Implicit Transfer Operators (ITO), an approach to building multiple time-scale surrogate models of stochastic molecular dynamics. We implement ITO models with a conditional DDPM using a new time-augmentation scheme and show how ITO models capture fast and slow dynamics on benchmarks and molecular systems. We show ITO models are self-consistent over multiple time scales and highly robust to the marginalization of degrees of freedom in the system, which are unimportant to capture the long-time-scale dynamics. Combined with a SE(3) variant of the PaiNN architecture [9] (ChiroPaiNN), we further show strong empirical evidence of scaling to molecular applications, such as the folding of coarse-grained proteins. As such, we are confident that ITO is a stepping-stone toward general-purpose surrogates of molecular dynamics.

Acknowledgments and Disclosure of Funding

This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation and The Novo Nordisk Foundation (SURE, NNF19OC0057822). OW's work was funded in part by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (NNF20OC0062606) and by the Pioneer Centre for AI, DNRF grant number P1. Preliminary results were enabled by resources provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS) at Alvis (project: NAISS 2023/22-392) partially funded by the Swedish Research Council through grant agreement no. 2022-06725.

References

Bernt Øksendal. Stochastic Differential Equations . Springer Berlin Heidelberg, 2003. DOI: 10.1007/978-3-642-14394-6 . URL: https://doi.org/10.1007/978-3-642-143946 .

Daan Frenkel and Berend Smit. Understanding Molecular Simulation . Elsevier, 2002. DOI: 10.1016/b978-0-12-267351-1.x5000-7 . URL: https://doi.org/10.1016/b978-012-267351-1.x5000-7 .

Paul Langevin. 'Sur la théorie du mouvement brownien'. In: C. R. Acad. Sci. (Paris) 146 (1908), 530--533.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models . 2020. DOI: 10.48550/ARXIV.2006.11239 . URL: https://arxiv.org/abs/2006.11239 .

David Ruelle. Thermodynamic Formalism . en. Encyclopedia of mathematics and its applications. Harlow, England: Longman Higher Education, Nov. 1978.

Christof Schütte et al. 'Conformation dynamics'. In: Proc. Int. Congr. ICIAM (2009), pp. 297336.

Jan-Hendrik Prinz et al. 'Markov models of molecular kinetics: Generation and validation'. In: The Journal of Chemical Physics 134.17 (May 2011), p. 174105. DOI: 10.1063/1.3565032 . URL: https://doi.org/10.1063/1.3565032 .

Jean-Pierre Serre. Linear Representations of Finite Groups . Springer New York, 1977. DOI: 10.1007/978-1-4684-9458-7 . URL: https://doi.org/10.1007/978-1-4684-94587 .

Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra . 2021. eprint: arXiv:2102.03150 .

Yang Song et al. Score-Based Generative Modeling through Stochastic Differential Equations . 2020. DOI: 10.48550/ARXIV.2011.13456 . URL: https://arxiv.org/abs/2011. 13456 .

Jascha Sohl-Dickstein et al. 'Deep Unsupervised Learning using Nonequilibrium Thermodynamics'. In: Proceedings of the 32nd International Conference on Machine Learning . Ed. by Francis Bach and David Blei. Vol. 37. Proceedings of Machine Learning Research. Lille, France: PMLR, 2015, pp. 2256-2265. URL: https://proceedings.mlr.press/v37/ sohl-dickstein15.html .

Ricky T. Q. Chen et al. 'Neural Ordinary Differential Equations'. In: Advances in Neural Information Processing Systems . Ed. by S. Bengio et al. Vol. 31. Curran Associates, Inc., 2018. URL: https://proceedings.neurips.cc/paper_files/paper/2018/file/ 69386f6bb1dfed68692a24c8686939b9-Paper.pdf .

Will Grathwohl et al. 'Scalable Reversible Generative Models with Free-form Continuous Dynamics'. In: International Conference on Learning Representations . 2019. URL: https: //openreview.net/forum?id=rJxgknCcK7 .

Yang Song et al. 'Score-Based Generative Modeling through Stochastic Differential Equations'. In: International Conference on Learning Representations . 2021. URL: https:// openreview.net/forum?id=PxTIG12RRHS .

Minkai Xu et al. GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation . 2022. eprint: arXiv:2203.02923 .

Jonas Köhler, Leon Klein, and Frank Noe. 'Equivariant Flows: Exact Likelihood Generative Learning for Symmetric Densities'. In: Proceedings of the 37th International Conference on Machine Learning . Ed. by Hal Daumé III and Aarti Singh. Vol. 119. Proceedings of Machine Learning Research. PMLR, 2020, pp. 5361-5370. URL: https://proceedings. mlr.press/v119/kohler20a.html .

Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization . 2014. eprint: arXiv:1412.6980 .

URL: https://markovmodel.github.io/mdshare/ALA2/#alanine-dipeptide .

K. Lindorff-Larsen et al. 'How Fast-Folding Proteins Fold'. In: Science 334.6055 (Oct. 2011), pp. 517-520. DOI: 10.1126/science.1208351 . URL: https://doi.org/10.1126/ science.1208351 .

Marloes Arts et al. Two for One: Diffusion Models and Force Fields for Coarse-Grained Molecular Dynamics . 2023. eprint: arXiv:2302.00600 .

Hao Wu and Frank Noé. 'Variational Approach for Learning Markov Processes from Time Series Data'. In: Journal of Nonlinear Science 30.1 (2019), pp. 23-66. DOI: 10.1007/s00332019-09567-y . URL: https://doi.org/10.1007%2Fs00332-019-09567-y .

Guillermo Pérez-Hernández et al. 'Identification of slow molecular order parameters for Markov model construction'. In: The Journal of Chemical Physics 139.1 (July 2013), p. 015102. DOI: 10.1063/1.4811489 . URL: https://doi.org/10.1063/1.4811489 .

Natasa Djurdjevac, Marco Sarich, and Christof Schütte. 'Estimating the Eigenvalue Error of Markov State Models'. In: Multiscale Modeling & Simulation 10.1 (Jan. 2012), pp. 61-81. DOI: 10.1137/100798910 . URL: https://doi.org/10.1137/100798910 .

Marco Sarich, Frank Noé, and Christof Schütte. 'On the Approximation Quality of Markov State Models'. In: Multiscale Modeling & Simulation 8.4 (Jan. 2010), pp. 1154-1177. DOI: 10.1137/090764049 . URL: https://doi.org/10.1137/090764049 .

Andreas Krämer et al. 'Statistically Optimal Force Aggregation for Coarse-Graining Molecular Dynamics'. In: The Journal of Physical Chemistry Letters 14.17 (Apr. 2023), pp. 3970-3979. DOI: 10.1021/acs.jpclett.3c00444 . URL: https://doi.org/10.1021/acs. jpclett.3c00444 .

Wangfei Yang et al. 'Slicing and Dicing: Optimal Coarse-Grained Representation to Preserve Molecular Kinetics'. In: ACS Central Science 9.2 (Jan. 2023), pp. 186-196. DOI: 10.1021/ acscentsci.2c01200 . URL: https://doi.org/10.1021/acscentsci.2c01200 .

Frank Noé et al. 'Constructing the equilibrium ensemble of folding pathways from short off-equilibrium simulations'. In: Proceedings of the National Academy of Sciences 106.45 (Nov. 2009), pp. 19011-19016. DOI: 10.1073/pnas.0905466106 . URL: https://doi. org/10.1073/pnas.0905466106 .

Adam Paszke et al. 'PyTorch: An Imperative Style, High-Performance Deep Learning Library'. In: 32 (2019). Ed. by H Wallach et al. URL: https://proceedings.neurips.cc/paper_ files/paper/2019/file/bdbca288fee7f92f2bfa9f7012727740-Paper.pdf .

William Falcon and The PyTorch Lightning team. PyTorch Lightning . Version 1.4. Mar. 2019. DOI: 10.5281/zenodo.3828935 . URL: https://github.com/LightningAI/lightning .

James Bradbury et al. JAX: composable transformations of Python+NumPy programs . Version 0.3.13. 2018. URL: http://github.com/google/jax .

Cheng Lu et al. 'DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps'. In: arXiv preprint arXiv:2206.00927 (2022).

Jes Frellsen et al. 'Bayesian Generalised Ensemble Markov Chain Monte Carlo'. In: Proceedings of the 19th International Conference on Artificial Intelligence and Statistics . Ed. by Arthur Gretton and Christian C. Robert. Vol. 51. Proceedings of Machine Learning Research. Cadiz, Spain: PMLR, 2016, pp. 408-416. URL: http://proceedings.mlr.press/v51/ frellsen16.html .

Alessandro Laio and Michele Parrinello. 'Escaping free-energy minima'. In: Proceedings of the National Academy of Sciences 99.20 (2002), pp. 12562-12566. ISSN: 0027-8424. DOI: 10.1073/pnas.202427399 . eprint: https://www.pnas.org/content/99/20/12562. full.pdf . URL: https://www.pnas.org/content/99/20/12562 .

D. P. Landau, Shan-Ho Tsai, and M. Exler. 'A new approach to Monte Carlo simulations in statistical physics: Wang-Landau sampling'. In: American Journal of Physics 72.10 (2004), pp. 1294-1302. DOI: 10.1119/1.1707017 . eprint: https://doi.org/10.1119/1. 1707017 . URL: https://doi.org/10.1119/1.1707017 .

Amey P. Pasarkar et al. 'Vendi sampling for molecular simulations: Diversity as a force for faster convergence and better exploration'. In: The Journal of Chemical Physics 159.14 (Oct. 2023). DOI: 10.1063/5.0166172 . URL: https://doi.org/10.1063/5.0166172 .

Yihang Wang, Lukas Herron, and Pratyush Tiwary. 'From data to noise to data for mixing physics across temperatures with generative artificial intelligence'. In: Proceedings of the National Academy of Sciences 119.32 (Aug. 2022). DOI: 10.1073/pnas.2203656119 . URL: https://doi.org/10.1073/pnas.2203656119 .

Frank Noé et al. 'Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning'. In: Science 365.6457 (Sept. 2019). DOI: 10.1126/science.aaw1147 . URL: https://doi.org/10.1126/science.aaw1147 .

Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching . 2023. eprint: arXiv: 2306.15030 .

Juan Viguera Diez et al. 'Generation of conformational ensembles of small molecules via Surrogate Model-Assisted Molecular Dynamics'. In: (Nov. 2023). DOI: 10.26434/chemrxiv2023-sx61w . URL: http://dx.doi.org/10.26434/chemrxiv-2023-sx61w .

Bowen Jing et al. Torsional Diffusion for Molecular Conformer Generation . 2022. eprint: arXiv:2206.01729 .

Elman Mansimov et al. 'Molecular geometry prediction using a deep generative graph neural network'. In: Scientific Reports 9.1 (2019), pp. 1-13.

Chence Shi et al. Learning Gradient Fields for Molecular Conformation Generation . 2021. eprint: arXiv:2105.03902 . URL: http://arxiv.org/abs/2105.03902 .

Gregor Simm and Jose Miguel Hernandez-Lobato. 'A Generative Model for Molecular Distance Geometry'. In: Proceedings of the 37th International Conference on Machine Learning . Ed. by Hal Daumé III and Aarti Singh. Vol. 119. Proceedings of Machine Learning Research. PMLR, 2020, pp. 8949-8958. URL: http://proceedings.mlr.press/v119/simm20a. html .

Chence Shi et al. 'Learning Gradient Fields for Molecular Conformation Generation'. In: Proceedings of the 38th International Conference on Machine Learning . Ed. by Marina Meila and Tong Zhang. Vol. 139. Proceedings of Machine Learning Research. PMLR, 2021, pp. 9558-9568. URL: https://proceedings.mlr.press/v139/shi21b.html .

Tarun Gogineni et al. 'TorsionNet: A Reinforcement Learning Approach to Sequential Conformer Search'. In: CoRR abs/2006.07078 (2020). arXiv: 2006.07078 . URL: https:// arxiv.org/abs/2006.07078 .

Robin Winter, Frank Noé, and Djork-Arné Clevert. Auto-Encoding Molecular Conformations . 2021. URL: http://arxiv.org/abs/2101.01618 .

Victor Garcia Satorras et al. E(n) Equivariant Normalizing Flows for Molecule Generation in 3D . 2021. eprint: 2105.09016 . URL: http://arxiv.org/abs/2105.09016 .

Niklas Gebauer, Michael Gastegger, and Kristof Schütt. 'Symmetry-adapted generation of 3d point sets for the targeted discovery of molecules'. In: Advances in Neural Information Processing Systems . Ed. by H. Wallach et al. Vol. 32. Curran Associates, Inc., 2019. URL: https : / / proceedings . neurips . cc / paper / 2019 / file / a4d8e2a7e0d0c102339f97716d2fdfb6-Paper.pdf .

Minkai Xu et al. An End-to-End Framework for Molecular Conformation Generation via Bilevel Programming . 2021. eprint: arXiv:2105.07246 .

Sereina Riniker and Gregory A. Landrum. 'Better Informed Distance Geometry: Using What We Know To Improve Conformation Generation'. In: Journal of Chemical Information and Modeling 55.12 (Nov. 2015), pp. 2562-2574. DOI: 10.1021/acs.jcim.5b00654 . URL: https://doi.org/10.1021/acs.jcim.5b00654 .

Paul C. D. Hawkins et al. 'Conformer Generation with OMEGA: Algorithm and Validation Using High Quality Structures from the Protein Databank and Cambridge Structural Database'. In: Journal of Chemical Information and Modeling 50.4 (Mar. 2010), pp. 572-584. DOI: 10.1021/ci100031x . URL: https://doi.org/10.1021/ci100031x .

Paulo C. T. Souza et al. 'Martini 3: a general purpose force field for coarse-grained molecular dynamics'. In: Nature Methods 18.4 (Mar. 2021), pp. 382-388. DOI: 10.1038/s41592-02101098-3 . URL: https://doi.org/10.1038/s41592-021-01098-3 .

Jiang Wang et al. 'Machine Learning of Coarse-Grained Molecular Dynamics Force Fields'. In: ACS Central Science 5.5 (Apr. 2019), pp. 755-767. DOI: 10.1021/acscentsci.8b00913 . URL: https://doi.org/10.1021/acscentsci.8b00913 .

Brooke E. Husic et al. 'Coarse graining molecular dynamics with graph neural networks'. In: The Journal of Chemical Physics 153.19 (Nov. 2020), p. 194101. DOI: 10.1063/5.0026133 . URL: https://doi.org/10.1063/5.0026133 .

Jonas Köhler et al. 'Flow-Matching: Efficient Coarse-Graining of Molecular Dynamics without Forces'. In: Journal of Chemical Theory and Computation 19.3 (Jan. 2023), pp. 942-952. DOI: 10.1021/acs.jctc.3c00016 . URL: https://doi.org/10.1021/acs.jctc.3c00016 .

Soojung Yang and Rafael Gómez-Bombarelli. Chemically Transferable Generative Backmapping of Coarse-Grained Proteins . 2023. eprint: arXiv:2303.01569 .

Yaxin An and Sanket A. Deshmukh. 'Machine learning approach for accurate backmapping of coarse-grained models to all-atom models'. In: Chemical Communications 56.65 (2020), pp. 9312-9315. DOI: 10.1039/d0cc02651d . URL: https://doi.org/10.1039/ d0cc02651d .

Leandro E. Lombardi, Marcelo A. Martí, and Luciana Capece. 'CG2AA: backmapping protein coarse-grained structures'. In: Bioinformatics 32.8 (Dec. 2015), pp. 1235-1237. DOI: 10. 1093/bioinformatics/btv740 . URL: https://doi.org/10.1093/bioinformatics/ btv740 .

Christof Schütte et al. A Hybrid Monte Carlo Method for Essential Molecular Dynamics . eng. Tech. rep. SC-98-04. Takustr. 7, 14195 Berlin: ZIB, 1998.

William C. Swope, Jed W. Pitera, and Frank Suits. 'Describing Protein Folding Kinetics by Molecular Dynamics Simulations. 1. Theory'. In: The Journal of Physical Chemistry B 108.21 (Apr. 2004), pp. 6571-6581. DOI: 10.1021/jp037421y . URL: https://doi.org/10. 1021/jp037421y .

Brooke E. Husic and Vijay S. Pande. 'Markov State Models: From an Art to a Science'. In: Journal of the American Chemical Society 140.7 (Feb. 2018), pp. 2386-2396. DOI: 10.1021/ jacs.7b12191 . URL: https://doi.org/10.1021/jacs.7b12191 .

Christopher Kolloff and Simon Olsson. 'Rescuing off-equilibrium simulation data through dynamic experimental data with dynAMMo'. In: Machine Learning: Science and Technology 4.4 (Dec. 2023), p. 045050. ISSN: 2632-2153. DOI: 10.1088/2632-2153/ad10ce . URL: http://dx.doi.org/10.1088/2632-2153/ad10ce .

Simon Olsson et al. 'Combining experimental and simulation data of molecular processes via augmented Markov models'. In: Proceedings of the National Academy of Sciences 114.31 (July 2017), 8265-8270. ISSN: 1091-6490. DOI: 10.1073/pnas.1704803114 . URL: http: //dx.doi.org/10.1073/pnas.1704803114 .

Simon Olsson and Frank Noé. 'Dynamic graphical models of molecular kinetics'. In: Proceedings of the National Academy of Sciences 116.30 (July 2019), pp. 15001-15006. DOI: 10.1073/pnas.1901692116 . URL: https://doi.org/10.1073/pnas.1901692116 .

Andreas Mardt et al. 'VAMPnets for deep learning of molecular kinetics'. In: Nature Communications 9.1 (2018). DOI: 10.1038/s41467-017-02388-1 . URL: https://doi.org/10. 1038%2Fs41467-017-02388-1 .

Hao Wu, Jan-Hendrik Prinz, and Frank Noé. 'Projected metastable Markov processes and their estimation with observable operator models'. In: The Journal of Chemical Physics 143.14 (Oct. 2015), p. 144101. DOI: 10.1063/1.4932406 . URL: https://doi.org/10.1063/1. 4932406 .

Hao Wu et al. Deep Generative Markov State Models . 2018. eprint: arXiv:1805.07601 .

Leon Klein et al. Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics . 2023. eprint: arXiv:2302.01170 .

W. K. Hastings. 'Monte Carlo sampling methods using Markov chains and their applications'. In: Biometrika 57.1 (Apr. 1970), pp. 97-109. DOI: 10.1093/biomet/57.1.97 . URL: https://doi.org/10.1093/biomet/57.1.97 .

S. Kieninger and B. G. Keller. 'Path probability ratios for Langevin dynamics-Exact and approximate'. In: The Journal of Chemical Physics 154.9 (Mar. 2021), p. 094102. DOI: 10.1063/5.0038408 . URL: https://doi.org/10.1063/5.0038408 .

Geemi P Wellawatte, Glen M Hocky, and Andrew D White. 'Neural potentials of proteins extrapolate beyond training data'. In: (2023).

Manuel Dibak et al. 'Temperature steerable flows and Boltzmann generators'. In: Physical Review Research 4.4 (Oct. 2022). DOI: 10.1103/physrevresearch.4.l042005 . URL: https://doi.org/10.1103/physrevresearch.4.l042005 .

Lukas Herron et al. Inferring phase transitions and critical exponents from limited observations with Thermodynamic Maps . 2023. eprint: arXiv:2308.14885 .

V Rokhlin. 'Rapid solution of integral equations of classical potential theory'. In: Journal of Computational Physics 60.2 (Sept. 1985), pp. 187-207. DOI: 10.1016/0021-9991(85) 90002-6 . URL: https://doi.org/10.1016/0021-9991(85)90002-6 .

P. P. Ewald. 'Die Berechnung optischer und elektrostatischer Gitterpotentiale'. In: Annalen der Physik 369.3 (1921), pp. 253-287. DOI: 10.1002/andp.19213690304 . URL: https: //doi.org/10.1002/andp.19213690304 .

Arthur Kosmala et al. Ewald-based Long-Range Message Passing for Molecular Graphs . 2023. eprint: arXiv:2303.04791 .

Benjamin Trendelkamp-Schroer et al. 'Estimation and uncertainty of reversible Markov models'. In: The Journal of Chemical Physics 143.17 (Nov. 2015), p. 174101. DOI: 10.1063/ 1.4934536 . URL: https://doi.org/10.1063/1.4934536 .

Susanna Röblitz and Marcus Weber. 'Fuzzy spectral clustering by PCCA+ application to Markov state models and data classification'. In: Advances in Data Analysis and Classification 7.2 (May 2013), pp. 147-179. DOI: 10.1007/s11634-013-0134-6 . URL: https://doi. org/10.1007/s11634-013-0134-6 .

Igor Mezi´ c. 'Spectral Properties of Dynamical Systems, Model Reduction and Decompositions'. In: Nonlinear Dynamics 41.1-3 (2005), pp. 309-325. DOI: 10.1007/s11071-0052824-x . URL: https://doi.org/10.1007%2Fs11071-005-2824-x .

Moritz Hoffmann et al. 'Deeptime: a Python library for machine learning dynamical models from time series data'. In: Machine Learning: Science and Technology 3.1 (2021), p. 015009. DOI: 10.1088/2632- 2153/ac3de0 . URL: https://doi.org/10.1088%2F26322153%2Fac3de0 .

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose Implicit Transfer Operator (ITO) learning, which aims to learn a surrogate of a molecular dynamics (MD) simulation process. Since standard MD simulations integrate Newton's equations of motion numerically, small integration time-steps are necessary, making simulations costly when processes on long time-scales (requiring many integration steps) are of interest. The proposed ITO method learns to simulate the dynamics across multiple time-scales, allowing to study processes on long time-scales more efficiently by using large time steps. It is also shown that ITO models are able to learn surrogate dynamics using partial observations, which is useful e.g. in the context of coarse-graining.\n\nSTRENGTHS:\nThe presented method has strong theoretical foundations and the underlying theory is well-described.\n\nWEAKNESSES:\nWhile the premise of the paper is interesting, the results and contributions are not particularly impressive:\n1. ChiroPaiNN is an extremely simple modification of the existing PaiNN architecture. So simple in fact, that I find it hard to argue this is a separate contribution (the only modification seems to be an added cross product between vector features, followed by a scalar product).\n2. The authors write that their models show quantitative agreement with dynamic and stationary observables, which is not backed by the results shown in Table 2. Free energy differences of folding have relative errors in excess of 400% and also absolute errors are extremely large with over 3 kT in some cases. Mean first passage times have similarly large errors. These results cast doubt on the practical usefulness of the proposed method.\n\nQUESTIONS:\n- It is unclear to me why symmetry with respect to parity inversion needs to be broken for the task at hand. Can the authors explain why this is necessary? What happens when an E(3)-equivariant models is used instead of an SE(3)-equivariant one? Additional ablations studies in this direction would be insightful.\n\n- Recent work (see arXiv:2302.00600) has pointed out an interesting connection between the score field $\\nabla \\log p$ (see Eq.7) in diffusion models and force-fields. Are the authors aware of this connection? Have the authors tried to \"extract\" the effective force-field learned by their model? My gut feeling is that it might be possible to run ordinary MD simulations with the learned force-field, but possibly using much larger time steps than would typically be feasible. I encourage the authors to investigate and run additional experiments in this direction.\n\n- Can the authors think of ways to improve their method to achieve better quantitative agreement for observables like free energy differences with conventional MD?\n\n- Direct efficiency comparisons to conventional MD would be interesting. What is the expected effective speed up achievable with the ITO method? Considering that training data for ITO needs to be generated, the models needs to be trained, and finally evaluated, I wonder how big the advantage is compared to directly running MD simulations.\n\nLIMITATIONS:\nLimitations with respect to transferability and scalability of the proposed method and the accuracy of the surrogate dynamics model are mentioned. In some cases it is unclear how the limitations can be reconciled in the future (e.g. the requirement of a closed-form expression for target path probabilities). Potential societal impacts of the work are not discussed, which however, I find acceptable considering the topic of the manuscript. A discussion of societal impacts would probably appear contrived.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper develops a conditional diffusion model to sample unnormalized densities - e.g., the Boltzmann distribution of molecules to replace molecular dynamics simulations. Given a molecule structure, noise is added to it and the diffusion model denoises this distribution to the distribution generated by MD. The diffusion model is not only conditioned on diffusion time, but also on MD time, such that a single model can be used to make larger or smaller MD steps.\n\nThe method is not generalizable - training data needs to be collected via the expensive MD that the method aims to replace. The paper is transparent about this. Experiments are performed on a toy system, alanine dipeptide, and 5 fast folding proteins.\n\nThe paper has strong contributions but is very unpolished. I think it can be substantially improved in the rebuttal, especially by addressing my first two content concerns in the weaknesses.\n\nI thank the authors for any time they take to answer my questions!\n\nSTRENGTHS:\n### Strengths\n\n1. The authors tackle an important and impactful task (sampling Boltzmann distributions of molecules).\n2. The idea is clever. Since it is hard to sample the Boltzmann distribution of a molecule unconditionally, the authors instead construct a noisy distribution that still shares structure with the bolztmann distribution such that the generative model has an easier job at producing a sample from the less complicated conditional distribution instead of the unconditional one. **************************************************Notably this was already done by TimeWarp************************************************** but the authors bring diffusion models to this task which seem better fit for it (if no reweighting is done).\n3. It is an interesting hypothesis that training with multiple MD times improves stability and would perform better than a single MD timestep size. Furthermore, this is justified and illustrated with interesting transfer operator theory. I think there is great value in bringing this theory to the attention of the ML for MD and boltzmann generator community.\n4. The model is self consistent. The experiments set up for this are ideal and illustrate the self consistency very well.\n\nWEAKNESSES:\nThe first three content concerns are by far the most relevant concerns\n\n### Content concerns (approximately in order of importance)\n\n1. The main claim (as I see it), that training with multiple MD times helps, is not sufficiently investigated:\n    1. Why do you only show results for the Muller Potential for this? I do not see why not to evaluate it for all systems or at least for some molecule instead of only the toy system.\n    2. Why do you only show VAMP2 score gaps and not kl divergences?\n    3. What about 1000 lag in Table 1?\n2. What is the comparison in wallclock time for sampling the whole distribution compared to MD? I would expect to see e.g., a plot of how the KL divergence of some ovservable goes down in MD vs. with your method in wall-clock time. Running diffusion models also takes quite some time I imagine, so is this not one of the main things that need to be investigated? Am I missing this information somewhere in the paper?\n3. The paper claims to be reproducible, but there is no code provided.\n4. Is there related work on ML for transfer operators (e.g., by Anima Anandkumar) that should be explained in the related work? I am not familiar with the field - just a potential pointer.\n\n\n### Presentation Concerns\n\nI think the presentation has to be improved and currently suffers from missing information, information relegated to the appendix, distraction with irrelevant material, and unnecessarily complicated explanations. The first two concerns are the most relevant.\n\n1. Are we obfuscating the method unnecessarily with math to make it seem more sophisticated? I think it would be tremendously helpful if the transfer operator theory is explained and justifies correctness, but the explanation of what is actually done in the end can be vastly simplified and made much more transparent by just stating that it is now a diffusion model that is conditioned on the MD timestep as well and can therefore make larger or smaller MD timesteps during inference.\n2. The “novel” architecture explanation is unnecessary and it is really not needed to claim this as a contribution instead of just focusing on the clearly good contributions and the point of the paper (to me at least) in using multiple MD times and a conditional diffusion model to replace MD. These are two novel **************actual************** contributions. How about just saying you use Painn and there is a little tweak similar to what you can do with e.g. e3nn to deal with chirality that the reader can look up in the appendix? \n3. (would take quite some rewriting I suppose): I think the presentation might benefit a lot from distinguishing between e.g., an “MD SDE” and a “Diffusion SDE”. Then you can distinguish between an MD time and a diffusion time and say that in the text next to your superscript, subscript notation. \n4. Figure 3 does not show any visible difference between fixed and stochastic lag? My recommendation would be to put it in the appendix and use the space for something useful.\n5. I would recommend not talking about a special coarse grained approach in the contributions introduction and abstract. I went into the paper thinking you do something interesting with coarse graining.\n6. You talk about nested sampling vs. direct sampling. While it is clear to me what is meant, this is not explained anywhere. A simple sentence would suffice.\n7. More explicitly point out the dependence of the added noise on the size of the MD timestep size.\n8. It should be mentioned in the main text that the number if diffusion ODE steps is the same independent of the size of N.\n9. Maybe fix all the typos especially in the related works section for the next/camera-ready version :sweatsmile:\n10. The appendix pointer for Algorithm 1 is wrong.\n\nQUESTIONS:\n### Questions\n\n1. For the dependence of the diffusion/the amount of added noise on the size N: What is the connection between increasing the maximum diffusion time T and making beta_i dpend on i instead and controlling the variance with that? Is one correct and the other is not or are they equivalent?\n2. Did you try an SDE solver for the diffusion SDE?\n3. How are the standard deviations in Table 2 calculated?\n4. Why do you use the same number of diffusion steps/discretization steps if you have a higher or lower N?\n5. How many samples do you draw in all your experiments? How long do you run the “MD” with the diffusion model?\n6. Why does the same i index both t and N? Wouldnt that mean that a particular step in the trajectory always uses the same N?\n\nLIMITATIONS:\nThe paper is transparent with the main limitation, that there is currently no capability to generalize. I think the extent to which this is made clear could go further though: point out that the model is only **********************overfitting********************** on a distribution and that there is **********no \"real-world use value\" (e.g., speeding up MD) in its current form.**********\n\nThe paper also points out the important difference to e.g. Timewarp that there is no direct way to do reweighting and get exact convergence to the boltzmann distribution in the limit. However, in the discussion it is made to seem like obtaining exact likelihoods from a diffusion model is very much feasible as well and would not come with a large amount of difficulty.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents implicit transfer operator (ITO) learning for the simulation in molecular dynamics. Their approach adopts the SE3 equivariant MPNN architecture (ChiroPaiNN) to parameterize the transition kernels in the denoising diffusion probabilistic model. The method displays a decent performance on several all-atom molecular simulation tasks with only coarse-grained representations.\n\nSTRENGTHS:\n- The problem is very relevant and the motivation is clear. Most notably, the method takes SE3-equivariance into consideration, which I think is a very good practice to combine state-of-the-art generative models with the simulation of molecular dynamics.\n- The paper presents a good balance between the introduction, methodology, and experiments. It is an interesting paper to read.\n- Illustrative examples are shown to demonstrate the effectiveness of their approach with detailed implementation details.\n\nWEAKNESSES:\n- The theory part concerning transfer operators is chaotic. Eq. (3) has a typo in it ($\\rho$ should be $f$). And in Appendix A.2, I don't think Eq. (17) should be correct if the transfer operator is defined as in (16). It seems they mix up the definition in Eq. (16) with that defined in Eq. (5) of [1]. \n- The experiments seem sound but it is very hard for me to evaluate the effectiveness of their approach since there is a lack of comparisons. As listed in the related works section by the authors, the idea of transfer operator surrogates has already been commonly used in molecular modeling as well as deep generative Markov state models. I believe there would be state-of-the-art methods other than all-atom MD simulations, to which they are comparing.\n- As I mentioned before, the use of an SE3 equivariant message-passing neural network seems one of the main contributions of this work. Following the last point, I would suggest at least one experiment that could explicitly demonstrate how and why this architecture is beneficial to the diffusion model. Will the results be much worsened if we remove this SE3-equivariance from the architecture, given a similar size of the NN?\n\n[1] Prinz, Jan-Hendrik, et al. \"Markov models of molecular kinetics: Generation and validation.\" The Journal of chemical physics 134.17 (2011).\n\nQUESTIONS:\n- I don't see a direct connection between Eq. (4) and (5) and their method. Would you explain more about the motivation of the algorithmic design from this \"relaxation\"?\n- Following the discussion in Appendix A.1, it seems that when $N$ is very large, the spectrum of the transfer operator $T_\\Omega$ will become very ill-posed (with one eigenvalue being 1 and all others converging to 0). Will this affect the performance of the method? Or does this mean that the $N_{\\rm max}$ in Algorithm 1 should not be too large?\n- I don't quite understand the results in Figure 5. There are two dashed lines in the figure all labeled \"Folded\". \n- Following the previous point, there seems to be a large discrepancy between the value of $\\langle \\tau_f \\rangle$ for Trp-Cage, BBA, and Villin in Table 2. Would you explain how should we interpret this? Is the MD result generally accurate or it is also just a reference? If so, how could we know that the simulation yields reasonable results?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

comments: SUMMARY:
The proposed approach aims to learn molecular dynamics via an implicit transfer operator framework that can perform modeling at multiple time scales. The framework is using diffusion model with SE3 equivariant architecture. The approach has shown capability in stable and self-consistent modeling at multiple time and space resolutions for molecular dynamic modeling.

STRENGTHS:
- a novel framework for multi-scale simulation for molecular dynamics, combined with the SE3 network for stable, long-term results. 
- demonstrate good experimental results on multiple MD datasets and achieve an order of magnitude speedup
- work on an important task of molecular dynamic modeling 
- use diffusion network for exact likelihood evaluation 
- honest and clear notes on limitations - appreciated

WEAKNESSES:
- in terms of presentation, it'd be nice to make each figure self-contained in terms of points to make and the meaning of the notations.
-

QUESTIONS:
- how are the number of time scales chosen during the training?  
- could you explain in the caption or here all the input for ITO networks in Figure 2
- what's the difference between the diffused time and actual time

LIMITATIONS:
- this paper may be intense for people noting working on MD. it'd be nice to include more descriptions in captions. 
- I am curious how it can be generalized to other dynamics like some other physics system that is also multiscale.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-09

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this work, the authors proposed a framework that combines SE(3)-equivariant MPNN and conditional DDPM, called Implicit Transfer Operator (ITO) learning, as a method to efficiently sample observales from MD simulation trajectories. Such a framework is validated on Muller-Brown potential data generated by the authors, as well as the commonly used alanine dipeptide and fast-folding protein dataset generated using MD simulation. ITO learning framework is shown to bear the potential of accelerating or surrogate MD simulation.\n\nSTRENGTHS:\n1. It is a nice effort to formulate the MD trajectory using the ITO such that the sampling can be learned by the conditional DDPM.\n\n2. Proposed a SE(3) version of the PaiNN.\n\nWEAKNESSES:\n#### Major\nMy major concerns are related to the experiment and results sections of the manuscript.\n1. In Figure 4, the $\\phi$-$\\psi$ plot or Ramachandran plot is often colored by the free energy of the system (e.g. [Köhler et al. 2023](https://pubs.acs.org/doi/epdf/10.1021/acs.jctc.3c00016), [Marloes et al. 2023](https://arxiv.org/abs/2302.00600) and many more), which is directly related to the probability of the state given Boltzmann distribution. From the current Figure 4, the physical property of the system, energy, is not directly visualized. \n\n2. Another issue with Figure 4 is the large bin size of the 2D histogram. The current binsize of 2D histogram is about 0.4 rad, which makes it really hard for readers to understnad the performance of the SE3-ITO model. Moreover, the MD simulation data are represented in the form of 2D histogram while the model sampled data are in the form of contours. I would highly recommend that the authors to show the MD data in separate subfigures with the same form of the model sampled data. If 2D histogram is to be used, a much smaller bin size should be used for clarity.\n\n3. Although the effort of proposing a SE(3)-equivariant version of PaiNN (ChiroPaiNN) should be recognized, the necessity of CPaiNN is not well estabilished in the manuscript. As the author mentioned in the manuscript, there is no parity change during MD simulation. I am wondering if there will be significant difference in accuracy if the original PaiNN is used in the ITO learning framework. Assuming an SE(3)-equivariant model is absolutely necessary, the authors have not shown any comparison between CPaiNN and other established SE(3)-equivariant model such as the [SE(3)-Transformer](https://arxiv.org/abs/2006.10503). Such a benchmark can definitely help to improve the manuscript.\n\n4. For the fast-folding protein experiments, CG-SE3-ITO model is compared with MD data. In Table 2, the error of the model is significantly higher for proteins with more residues or $C_{\\alpha}$ atoms. Yet, the discussion about the high error is limited in the manuscript. The purpose of coarse graining is to achieve relatively high accuracy with large system. Underwhelming accuracy on large system impairs the applicability of the SE3-ITO model on coarse graining.\n\n\n#### Minor\n5. Line 65, the probability is written as $p_{\\tau} (x_{t+\\tau} | x_{t})$, which is inconsistent with the notation in the dynamics observables equation ($p_{\\tau} (x_{t+\\Delta t} | x_{t})$) in line 64.\n\nQUESTIONS:\n1. In Figure 3, I assume that the green curve ($N=1000$) is not sampled from nested conditional probability. Can you verify my understanding?\n\n2. In Figure 3, the model with ($N=10$ and $\\Delta t=1000\\tau$) is lower comparing to higher $N$ values. In Figure 4, $\\Delta t=4$ps seems to lead to less accurate samples from model. My interpretation of those results is that the ITO learning framework suffers higher error when predicting more immediate conformational change of the molecule in MD simulation. If so, the dynamics of the molecule when moving between high probability state might be missed when using ITO framework. Can the authors discuss more about this?\n\nLIMITATIONS:\nYes, limitations are addressed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-22"}