Implicit Transfer Operator Learning: Multiple Time-Resolution Models for Molecular Dynamics
Authors: Mathias Schreiner
Ole Winther
Simon Olsson
Venue: NeurIPS 2023 poster (2023)
Keywords: AI4Science
Molecular Dynamics
equivariant neural networks
stochastic dynamics
TL;DR: We present an approach to learn multiple time-resolution surrogates of 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}\,\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
Mathias Schreiner"
Ole Winther DTU
Simon Olsson
DTU!
Chalmers University of Technology
Chalmers University of Technology matschreinerdgmai..
simonols@chalmers.cc
Abstract
Computing properties of molecular systems rely estimating expectations of the (unnormalized) Boltzmann distribution. Molecular dynamics (MD) a broadly adopted technique approximate such quantities. However, stable simulations rely on very small integration time-steps (10-15 moments, e-g. binding free energy or rates, might rely on sampling processes time-scales as long as 10-1s, 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 diffusson probabilistic models with new SE(3) equivariant architecture and show the resulting models can generate self-consistent stochastic dynamics across multiple timesscales, even when the system is only partially observed. Finally, we present a coarse-grand -aaure-g 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 available
150-1gg 8), whereas convergence of some
Introduction
Numerical simulation of stochastic differential equations (SDE) is critical in the sciences, including statistics, physics, chemistry, and biology applications [ Molecular dynamics (MD) simulations are an important example of such simu- lations (2]. These simulations prescribe
of mechanis- tic rules governing the time evolution of molecular system through numerical integration of, for example, the Langevin equation 3 MD grants mechanistic insights into experi- mental observables. These observables are expectations, in- cluding time-correlations of observable functions (c.g., pair- wise distances or angles) computed for the Boltzmann dis- tribution p(x)) + of a M-particle molecular system, E Ω C R3M kept at the inverse temperature / = 1/KT.. However, stable numerical integration relies on time steps, т, which are strictly smaller than the fastest characteristic time-
Tg(1)
Ta(N)
exp(=BU(x)) corresponding to the po-
tential U(.)
Figure 1: Implicit Transfer Opera-
tor: multiple time-scale surrogate of stochastic molecular dynamics.
Contributions this work were done while visiting Chalmers University of Technology Technical University of Denmark Coonespondigg author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

scales of the molecular system (10-1ss, s, e.g., bond vibrations), yet many molecular systems are characterized by processes on much longer time-scales (10-3 10 s. e.g. protein-folding, protein- ligand unbinding, regulation). Consequently, we need infeasibly long simulations to characterize many important processes quantitatively due to the slow mixing in 52.
In this work, we present the implicit Transfer Operator (ITO, Fig. D 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)- generative model, we further demonstrate stable long-time-scan 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 not generalize across different thermodynamic ensembles or across chemical space, but show strong generalization across different tinesscales.
dynamics in increasingly difficult settings where an
Our main contributions are
1. the Implicit Transfer Operator (TTO) framework for learning generative models for multiple time resolution molecular dynamics simulations,
2. implementation of ITO using denoising diffusion probabilistic model (DDPM) 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.
Background and Preliminaries
Notation Throughout this work, diffusion time, related to Diffusson Models (see Sec. physical time are represented using superscripts and subscripts, respectively.
2 and
Molecular dynamics and observables Molecular dynamics (MD) is a wide-spread simulation strategy computational chemistry and physics. In this approach, the time-evolution of N particles configuration in Euclidean space x Q C R3M, is modeled via a stochastic differential equation (SDE) with drift term based on a potential energy model U(x): f2 - R. An important aim of MD) is to compute:
1. Stationary observables: O, [x]f(x))
2. Dynamic observables: O((0M(+51)) = //tp:///e/.xaapp////////e///////
where y is the normalized Gibbs measure, and (x x;) is conditional probability density function encoding the time-discrete evolution of the molecular system x, with time-step At NN 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. Ω + RR are observable functions or 'forward models' deseribing the microscopic observation process, e.g. computing distance or an angle. The observables, O(()) and Ofejnii+sa Ac). include binding affinities and binding rates of drug to protein, respectively. Conventionally, these observables are estimated from simulation trajectories using naive Monte Carlo estimators.
For illustrative purposes, we assume the temporal behavior of state, x, follows the Brownian dynamics SDE (Itô form)
= -VUx) dt + v2DaW,
(1)
where D = 1 is diffusion constant, with friction y and anverse-temperauue 8, and dW is Wiener process. Using the Marer-Ma time-discretization, with time-step simulating the SDE corresponds to simulating a Markov chain with the transition probability density
p(x++7 Xx, T) = =N(x Xx: - ~UU(x))> ,
(2)
where N specifies the multi-variate Normal distribution, and law the 3M dimenssioaa identity matrix. If r 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 (x) at B. Consequently, by simulating large number of steps we can draw samples from to compute stationary observables and compute dynamic observables by simulating A/ Nr 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 p specify an initial condition, probability density function on Ω. We can define a Markov operator Ta (9) - L' (ft) using a transition density (e.g.. 2
(Ta p) (xc+r) - 1(xx+7) / ((xx))(xx))(xxxt xx))xx;
which then describes the j--eighedd evolution of absolutely convergent probability density functions on Ω according to eq. D with time-step, т. Such an operator is called the (Ruelle) Transfer Operator
We can express the operator using spectral form
Tx(T) - Σ (T)
where only eigenvalues (T( = exp(-Tk:) depend on the time-step, т. K; are characteristic
'relaxation' rates associated the left and right eigenfunction pair, d; and th 77. We can compute the
operator with time-lag Nr via the Chapman-Kolmogonn Kolmogorov equation (see Sec. A.1 for details)
Ta(NT) Σ
X.(T) (o))
(5)
Equivariant Message Passing Neural Networks In this work, we are concerned with MD, where the eine-evolution of molecule is governed by force field F(.) - (-) derived from central potential (-). While (-) is invariant to group-actions of the Euclidean group three dimensions (E(3)), its corresponding force field is -(3iivaararri We call function, f 'invariant' under group-action g iif f(x) fS5x) and 'equivariant" iff (x) = ((),)), where Sg and Tg are linear representations of the group element g 8
The force field F()) is equivariant under B(3) group-actions. However, in practice, classical molecular dynamics simulations do not change parity during simulation, and consequently, our data distribution only contains single mirror image of molecules.
We extended the PaiNN architecture (). an Equivvaara message passing neural network (MPNN), making it SE(3) equivariant by breaking its symmetry with respect to parity. We in- troduced this minor modification as we experienced sporadic parity changes when sampling with 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 (1i) = (j ||ray S reutoff where rij is the distance between nodes denoted and j, and Feutoff 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- cross-product is an axial vector (i.c., 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 frame- work that learns distributions by modeling gradual denoising process TO In DMs, we pre- specify a forward diffusion process (noising process), which gradually transforms the data distribution p(x") to a simple prior distribution p(x"). e.g.. a standard Gaussian, through a time-inhomogenous Markov process, described by the following SDE (Itô form)
dx' f(x", t) dt + g(t) dW.
(6)
3

---

where 0 <<<<T T is the diffusion time, f and are chosen functions, and dW is Wiener process.
We can generate samples from the data distribution p(xx) by sampling from p(x") and solving the backward diffusion process (denoising process)
dx' = [f((x', 1) Vx) logp(x' [t] dt + g(t) dw
(7)
by approximating the score field xx logp(x' t) - or equivalently a time-dependent Nor-
mal transition kernel ( with a deep neural network surrogate Vx log p(x* t. f. We
can use the learned score field to define neural ordinary differential equation (ODE)
12
13), or probability flow ODE [14] eq. 7 less the term g dW and scaling g 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. Conse- quently, we parameterize the DM us- ing learned Normal transition kernel
A
2
N
No
Fai
B Xn TpoNt N
(ee)
(!)
(l)
(m.3)
Apan
Apus
A
pas
Apes
of a time- ineommooooo diffusion pro- cess. By restricting the transition ker- nels p(x*+1 | x") to be equivariant un- der SE(3) group-actions, the marginal of
(m.
(et.n)
(n)
60
MB-ITO
Concat
Concat
2n)
x is always invariant (15). Combin- ing the equivariant transition kernel with an invariant prior density [16] ensures the whole Markov process is invariant to SE(3) group actions. Consequently, com- bining an isotropic mean-free Gaussian as prior with ChiroPaiNN ChiroPaiN--arramctriie ppaametterized transition kernels, we can construct an SE(3) equivariant diffusion model.
MLP
MLP
(2)
ChiroPaiNN
SE3-ITO
(m,,n)
Concat
(,,3,3)
Um,
MLP
ChiroPaiNN
mm
Figure 2: ITO è networks (A) SE3-ITO used for molec- ular application (B) MB-ITO, used for experiments with the Milllrr-rrww potential. Apes and Anom are positional and nominal embedding respectively, Concat is a concate- nation, and MLP is multi-layer perceptron. Arrows annotated with input and output shapes.
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 aim to approximate the long-time step transition probability pN- (xN= Xx) to allow for one-step sampling of long sonn-timm-ssee time-scale dynamics.
Here, we
As data we consider simulation trajectories. The trajectories are generated by explicit simu- lation which corresponds to sampling ancestrally from the small time-step transition density: X {xr...., .N-) ~ p(xnr ((n-1)r) ), with {1,...,N] .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 - PN- (XNr Xo) from MD data. In practice, we learn a generative model Xi+Nr ~ (x xx,NN denoising diffusion probabilistic model (cDDPM) of the form
with a conditional
/
p p(xiAN) N), &
DT p( Xx, N) dx
(
where xIT are latent variables of the same dimension as our output, and follow a joint density describing the backward diffusion process (eq. 7 and xT ~ (N I). We define a conditional sample likelihood as
((I
&
II pe (x), NN N.) TEI
where I is a list of generated indices i specifying a time 4 and a time-lag (r) integer multiple Nú. 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-liken loglliielhhod
( E
[[lc co(x tNr,r X N1, tdift
(10)
where x aafff X; + v1 C'faiffe, with @f = Il (1 - 8e) and 8 is the variance of the
forward diffusion process diffusion time, v. (g(+) is one of the two ITO neural network model architectures shown in Fig. 2 and is directly related to the score
If the data used train the conditional transition density is generated by MD simulation with time-invariant potential energy (drift), we can express the generating transition probability as decomposition of time-variant and -invariant parts (Proof, see Sec. A.2
p(aNT zo) = Σ (T) (,(IN+))((20)
(11)
timo-variant
time-invariant
where at and B; time-invariant projection coefficients of the state variables on-to the left and right eigenfunctions Φ; and thi of the Transfer operator (T( (5) and [A,(T) S is its i' eigenvalue.
Consequently, call surrogate modeling approach implicit transfer operator learning.
As outlined in Algorithm D we generate indices i, e.g. the tuples (æx,, Nor, N:). such that the model is exposed multiple time-lags, sampled uniformly across orders of magnitude, used for gradient-basee 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 cigenfunctions and, consequently, better generalization across time scales and yield more stable sampling.
manner
optimization with Adam [7]. As result, as illustrated in eq. M the model
3.1 ITO Architectures
We present two architectures for learning cDDPMs encoding ITO models, one for molecular applica- tions SE3-ITO and for the Müller-Brown benchmark system (Fig. The SE3-ITO architecture uses our new SE(3) equivariant MPNN (ChiroPaiNNN, described in sec. 2 to encode x,, N, and atom-types, 2, invariant features, s. and equivariant features, v. We concatenate with encoding of the diffusion- time fdiff and process them through a MLP (multi-layer perceptron). The output from the MLP are passed along with v and x/air as input to a second ChiroPaiNN module which predicts More details on the architecture and hyperparameters
available in Appendices D and
Algorithm Training. DisExp is defined Ap- pendix E
Algorithm Ancestral sampling. Sampling
from is defined in Appendix EL Algorithm
Input:
n MD-trajectories:
X
Input: initial condition Xy, lag N, ancestral
(xg...., .K.....,. xL, In Jz=0= ITO score-model: Eo.
steps n.
Allocate E R(x+1))dim(xx)))
max lag; Nmax
X" Concatenate ((x) xy, -MnnN 15
T X0
for = 1...
n2 do
while not converged do
X; po (T 1], N)
X, Choice(A"')
1 X;
N DisExi(EExiii max)
end for return T
lair ~ UniffrrmO., ?
Take gradient step on:
Vo |||c éo faar
Xt, N, WaAa)))l2))
end while
return ég
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
5

---

Table 1: VAMP2 score-gaps. Difference in SCMP2-ccores of ancestral sampling from ITO models with fixed lag and stochastic lags, compared baseline Langevin simulations. Perfect match negative and negative values correspond to under and over estimation of meta-stability, respectively. Standard deviations on last decimal place are given parentheses.
system lag
10
100
Müller-Broo (fixed)
-0,0351 1 (5)
-0.1189 (2)
Mhiller-Browe (stochastic)
-0.0312 (4)
-0.0970 (5)
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. [ with the Müller-Brown potential energy as (x) (For details, see Appendix B.1. This dataset corresponds to 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 proprietary but available upon request for research purposes. The simulations were performed in explicit solvent with a 2.5 time-step and the positions was saved at 200 ps intervals. We coarse- grainn the simulation by representing each amino-acid by the Euclidean coordinate of their Ca atom done previously [20). leading to 10, 20, 28, and 35 particles in each system respectively. Consequently, these data correspond to mostly unobserved case.
4.2 Stochastic lag improves meta-stability prediction
In sec. 2 we conjecture that exposing an ITT model to multiple lag times during training leads to better and more robust models. To test this, we trained set of models on the fixed constant lags N {10, 100, 1000} (fixed lag) and single model with N DisExxp(1000) (stochastic lag) using the MB-ITO model (Fig.
dataset with
We find that the model trained with stochastic lag systematically outperforms models trained with fixed lag (Table D. We gauge the agreement by comparing Variational Approach Markov Processes (VAMP)-2 scores [2I] (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 tong-time-scalle
that the difference in the ability of fixed and stochastic lag ITO models to capture 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 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- of alanine di-peptide as a function of epoch for models trained with fixed lag and stochastic lag (Fig. ()) 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 self-consistent manner, in the Shapman-Kolmogoro sense, i.e., p(Xst xn) 2 p(xNr xo) = N p(x,, x(-1)r). ), or if samples generated by direct sampling with time-step At NT are

---

distributed similarly to samples generated by performing ancestral sampling N times, each with time- step r. direct sampling, we draw samples for a desired time-step Åt Nr from po (xp, N) and in ancestral sampling we draw n samples with time-step Al = Nr from (x N) in an ancestral manner, e.g. Ii+1 ~ po( N), where [0,...n-1}}
For the fully-observed Müllee-Brown case, we find that the ITO model is self-consistent by the strong overlap in transition densities sampled 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. B 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* to 10 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 150000 samples in total. In the case of ancestral sampling, we sample 4, 64, or 512 steps for 15 000 trajectories, to match At.
Next, we consider four fast-folding proteins [9] where we coarse-grail -raiin the proteins by representing them only with their Ca 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
Ar= 4ps
- 64ps
pixaiXe)
and unfolded states (Fig. 4 projected onto linear subspace determined us- ing time-lagged independent compo- nent analysis (time-lageed indepen- dent components, tIC) [2 (see Ap- pendix B. Further, the long time- scale transition density gradually con- verges to the data distribution pected.
p(xulxa)
At
aoene I 1
Finally, by ancestral sampling (Algo- rithm 2) we perform simulation of Chignolin with the same length as the training trajectory (106 us), us- ing CG-SE3-ITO model, and com- pare with MD. The CG-SE3-ITO sim- ulation is 2120 steps with At 5 ns. Running in parallel, on a sin- gle Titan X GPU we can simulate the CG-SE3-ITO model at a rate of 363 ns/(nn (s.,, 22) where S. denotes seconds wall-time (Appendix C.I. Remarkably, these trajectories are vir- tually 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. aaand TO
Figure 3: Transition probability densities of alanine dipep- tide dynamics with SE3-ITO model; Rows of increasing time-lag (from top to bottom). Contours are samples from SE3-ITO model, and 2D histograms show estimates from MD data. The first column shows conditional transition den- sities projected onto the torsion angles & and &o (inset). The black cross indicates the initial condition. The second and third columns show marginal distributions d and vi, respec- tively, with direct sampling in orange, ancestral sampling in blue, and MD data in black.
Together these results suggest that ITO models accurately and selffconsistently capture the slow dynamics of molecular systems and are robust to situations where the system is only partially observed. In general, we expect robustness sparsely observed representations as long as the input representations are sufficient to span the eigenfunctions of Tr [23 24). Approximation errors will translate into systematic under-estimation of relaxation time-scales 77]. consistent with our slight under-estimation of VAMP-2 scores (Table In future work, combining the learning of SE3-ITO models with systematic scheme for coarse-graining [5 26 could be an avenue for scaling to
large-scale molecular systems at a low computational cost.

---

At = 200ps
A4 = 2nd
At 20ns
4C 200ns
From Unfoldedd
t
Folded DC2
tiC1
t(11
tIC1
tiC1/bC2 tiC1/ti22
MD
-25
V
ITO
UC2
L1C1
Folded
20
40 t/us
62
80
100
Figure 4: Reversible protein folding-unfoldiin of Chignolin with CG-SE3-ITO Conditional probability densities (orange contours) starting from unfolded (upper panels) and folded (lower panels) protein states, at increasing time-lag (left to right), shown on top of data distribution. Below: time-traces of 106 microsecond MD simulations and ITO simulations on tICs and The two dashed lines correspond to the folded state value in tIC I (lower line) and tIC 2 (higher line). Contour lines are based on 10'000 trajectories, generated with ancestral sampling with the length, A/ and time-step 200 For A/ 200 ns this corresponds to 1000 ancestral samples.
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 time-step which sufficiently large the latter akin to Boltzmann generators (27) (See Appendix A. D. We note that, the ITO models surrogates and such without reweighing 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:
also
sample i.i.d from the Boltzmann distribution
Free Energy of Folding, AG = log I-Pr
pt
Mean first passage time, folding, (T)) = Jave if So0 8(x. E s)p(a, xo. dl dro
Mean first passage time, unfolding, (Tw) = - So 6(xz E <-))p(ar zo. f)t dt dao
where -f- Ω are disjoint subsets corresponding to the folded and unfolded states of protein, Py = Szey (a(x) dæ, is the folded state probability and 0(-) is the Dirac delta.
We compute 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 (.6). Strikingly, the observables computed using CG-SE3-ITO models agree well with those computed from long all-atom MD simulations (Table
Finally, we analyzed the robustness, convergence, and consistency 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 indistinguisaable, suggesting consistency, robustness, and convergence. The

---

Table 2: Molecular observables Standard deviations on last decimal place are given in parentheses. Stationary and dynamic observables are denoted $ and d, respectively.
AG (s)
(TT)/Hs (d)
(Tu)/ (d)
Chignolin (MD/ITO)
-1.28(1V-0664(33
(665(4)).02((2))
2.01 12(34)
Trp-Cage (MD/TTO)
1.47(6/2.84(6)
13.6(4)/37(2)
3,4(2)2.85(9)
BBA (MD/TTO)
0.97(3y1152633
11.7()88.6(2)
5.1(1VL177(4)
Villin (MD/ITO)
1.21(222.2(33)
2.41())8.2(77)
0.68())0.354(5)
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 PyTTorch Lightnngg 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 including analysis methods involving deep generative nets (6)., and surrogate models which directly approximate the Boltzmann distribution and allow for recovery of unbiased statistics, including Boltzmann generators 37
methods [32 33
34
35
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
45 46
48
49
and chemoinformatic (50 5T] approaches. Finally, speeding up molecular simulations by reducing the effective number of partieles to simulate through coarse- graiining with special purpose forcefield
models [52] including machine learned variants [53] 54 20 [5] and learned coarse coarse-grainine graining maps
25
26 is orthogonal approach to sample conformation space. Further, several methods to recover
all-atom models from coarse-graiie grained representations through MI. [56] 57] and rule-based approaches available.
58
Transfer Operator surrogates Building transfer operator surrogates is commonly used in molecu- lar modeling including (Deep Generative) Markov state models (MSM) (59
7
60
61 also including
experimental data 62.2 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 644. Klein et al. recently introduced rimewarp which is a flow-based generative model simulate molecular systems with a large (up to 0.5 ns), fixed, time-lag, [68] providing asymptotically unbiased equilibrium samples through a Metropolis- -astiiss correction [69]. While timewarp generates conformers with realistic local structure, has limitations in captur- ing 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.
Limitations
Surrogate model Implicit Transfer Operators are surrogate models of stochastic dynamics' con- ditional 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 misspecifica- tion) uncertainty. We can overcome the latter by reweighing against Markov Chain Monte Carlo acceptance criterion as proposed previously [68], to ensure unbiased dynamics path-reweighing 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 (7T]. Generalization across thermodynamic variables such as temperature and pressure would require using 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 ( ) 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-rann long-range additive term 75 could yield more favorable scaling
76.
8 Conclusions
This paper introduces Implicit Transfer Operators (ITO), an approach building multiple time-scale surrogate models of stochastic molecular dynamics. We implement ITO models with a conditional DDPM using a new dime-angmentttion 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 the marginalization of degrees of freedom in the system, which are unimportant to capture the long-time-scaa 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- corse-grained proteins. As such, we are confident that ITO stepping stone toward general-pune 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, NNF190C0057822), OW's work funded in part by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (NNF200C0062606) 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
[1] Bernt Øksendal. Stochastic Differential Equations. Springer Berlin Heidelberg, 2003. DO1: 10. 10.1007/978.. 642 14394-6)
URL: httt:p:///oll.rrr/00000/////3 6-2 14394
6
[2] Daan Frenkel and Berend Smit. Understanding Molecular Simulation. Elsevier, 2002. DOI:
10. 1- 211--667351111 x.0000-77
URL:
11-2673611-125000-
[3] Paul Langevin. "Sur la théorie du mouvement brownien". In: C. R. Acad. Sci. (Paris) 146 (1908), 530-533.
[4] Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising Diffussio Probabilistic Models. 2020.
DOI: 10. 18550/ARXXIV 2006. 11239 URI.: ltttps://t/tttttte org/ has//bb//2000 11239
[5] David Ruelle. Thermodynanic Formalism. en. Encyclopedia of mathematics and its applica- tions. Harlow, England: Longman Higher Education, Nov. 1978. Christof Schütte et al. "Conformation dynamics". In: Proc. Int. Congr. ICIAM (2009), pp. 297- 336.
[6]
[7] Jan Hendrik Prinz et al. "Markov models of molecular kinetics: Generation and validation". In: The Journal of Chemical Physics 134.17 (May 201 I), 174105. DO1: 10. 1063/ 3565032 URL: httpt::t///// sttps://doi... org/ 10. 1063/1061111 3565032
[8]
Jean-Pierre Serre. Linear Representations of Finite Groups. Springer New York, 1977, DOI:
10. 1-11779988--4644--11
URL: https: 100.://o///....000/0------------
7
10

---

[9] Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger. Equivariant m passing for the prediction of tensorial properties and molecular spectra. 2021. eprint: arXiv: 03150
[10]
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
[11]
Jascha Sohl-Dickster et al. "Deep Unsupervised Learning using Nonequilibrium Thermody- namics". 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: : http:::/prrceedingg mlr press v
sohl minl-dickttein15 html
[12] Ricky T. Q. Chen 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: httpps:/ atttss://proceeedinng
neurips cc/ /cc/paper_.. files /papsr/pppr//eee 2018 file/
[13]
Will Grathwohl et al. "Scalable Reversible Generative Models with Free-form Continuous D Dynamics". In: Internatiional Conference on Learning Representations. 2019. URL: https
/popnneview www.:/eneeoooooumiidneeee.
[14] Yang Song al. "Score- Generative Modeling through Stochastic Differential Equa- tions". In: International Conference on Learning Representations. 2021. URL: https : //
openreview
[15] Minkai Xu et al. GeoDiff: Geometric Diffusion Model for Molecular Conformation Genera-
tion. 2022. eprint: :rXiv:2203 02923
[16) 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. mlr. www.sresssseee. peress/vllne http//111/kkohee0000 html
URL: https:::ps:t /tsto//ut:/tttts////paa p
[17] Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. 2014. eprint:
arXiv: 1412. 6980
[18] URL: https://markovooo//
[19] K. Lindorff-Larsen et al. "How Fast-Folding Proteins Fold". In: Science 334.6055 (Oct. 201 1).
pp. 517-520. DOL: 10. 1126 sts///cienccee 1208351
URL: https: //doi. org/ 10. /126/
science 1208351
[20] Marloes Arts et al. Two for One: Diffusion Models and Force Fields for Coarse-Grainct Molecular Dynamics. 2023, eprint: arXiv: : 2302 00660
[21] Hao Wu and Frank Noé. "Variational Approach for Learning Markov Processes from Time Series Data". In: Journal of Nonlinear Science 30. (2019), pp. 23-66. DOL: 10.0077/s00332 019-09567-y/ URL: https: http://ooooorgrgg org/ /ttt:e:///bi.oorr/0.0000///////0 1ttp:///...000000000033222222-20 01w.:/0003322-199999----9-99
[22] Guillermo Pérez-Hernández al. "Identification of slow molecular order parameters for
Markov model construction".. In: The Journal of Chemical Physics 139.1 (July 2013), 015102. URI.: https:///oli wttpii.rrr1.... httg://0.1066////0 1811111411111111188
DOI: 10. 0...:/..481110 1489
[23] 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. 1O. 11771100998910 URL: lattps:///oi..... www.::....r./1... wttp:/...1137//11 11371100111700008919000
[24] 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. 1037 011770090764043 URL: https: doi .rgg 1137 0011337009774444
[25]
Andreas Krämer et al. "Statistically Optimal Force Aggregation for Coarse- Granning 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
[26] 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/ acscentaci. 2c01200 URL: https: ://d/ org 10. www.1///ccc/cccci.......m. 2c01200

---

[27] 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. 19011119016. DOT: 10. 1073/ / pnas 0905466106
URL.: https: /oi
0905466106
[28]
Adam Paszke et al. "PyTorch: An Imperative Style, High-eefooonee Deep Learning Library". In: 32 (2019). Ed. by H Wallach et al. URL: http:ttsp://occceediiiiaae newr.ed.ddd.ieeeeee..ee....... nwwrips.rippp.. cc/paperr.........ppp.p files/paper/
[29]
William Falcon and The PyTorch Lightning tcam. PyTorch Lighming. Version 1.4. Mar.
2019. DO1: 10 5281 zenodo. 3828935
URL: https // github com Lightning
AI/lightning
[30] James Bradbury et al. JAX: composable transformations of Python+NNmmP programs. Ver- sion 0.3.13. 2018. URL: http://github: http::///ggituuu.cooooooee/e// com/ coom/goolee/jae
[31] Cheng Lu al. "DPM-Solver: Fast ODE Solver for Diffussin Probabilistic Model Sampling
in Around 10 Steps". In: arXiv preprint
(2022).
[32] Jes Frellsen et al. "Bayesian Generalised Ensemble Markov Chain Monte Carlo". In: Pro- ceedings 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: httt:://proceedings. mll. wtw.r..rrss///e..e. wwwss/ /v51/
frellsen16. html
[33] 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:
pnas org/ content/99/20/ 12562 full. CRI:: httpt:ttt////tt www. pnas watesrorgggooooooooeee wwwg///ooeen///eee 90/0//999000/0 20/12562
10. 1 pmas. 209427399
eprint: https:
[34] 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. (2004),
pp. 109411911000 DO1: 10 1119 1707017 eprint: https: // doi org / 10. 1119/1
1707017 URL: https::tt:://o org/10. 1119/1. 1707017
[35] Amey P. Pasarkar et al. "Vendi sampling for molecular simulations: Diversity force for
faster convergence and better exploration". In: The Journal of Chemical Physics 159.14 (Oct. 2023). DOI: [0. 1063/16665500 0166172 URL: TRI.:htpss: htttt::///oooo orog 10. 1063/5.60000 0166172
[36] Yihang Wang, Lukas Herron, and Pratyush Tiwary. "From data noise to data for mixing physics across temperatures with generative artificial intelligence". In: Proceedings of the
National Academy of Sciences 119.32 (Aug. 2022). D01: 10. 10.1107/ppaas http:/pp....22255511111 2203656119 URL:
https://doi.... 0rg/10..r/ 1073 /paa. 2203656119
[37] Frank Noé et al. "Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning". In: Science 365.6657 (Sept. 2019). DOE:
aaw1147
URL: hattps:///ol org/ 10. 1126/ wt///////innc/ ence. aaw1147
[38] Leon Klein, Andreas Krämer, and Frank Noé. Equivariant flow matching. 2023. eprint: arXiv
[39] Juan Viguera Diez et al. "Generation of conformational ensembles of small molecules via Surrogate Model-Assister Molecular Dynamics", In: (Nov. 2023). DOr: 10. 26434/ /ctp/xive URL: http::///:: doi 20g/ 26434/ cheerr//eerxvv 2023-sx61w
[40] Bowen Jing et al. Torsional Diffusion for Molecular Conformer Generation. 2022. eprint:
arXiv:2206 2206601299119 01729
[41] Elman Mansimov et al. "Molecular geometry prediction using a deep generative graph neural network". In: Scientific Reports 9.1 (2019), pp. 1-13.
[42] Chence Shi et al. Learning Gradient Fields for Molecular Conformation Generation. 2021.
eprint: arxiv: 2105. 03902 URL: http: //arxiv....i//////a www. ab1//155./2220.. 03902
[43] Gregor Simm and Jose Miguel Hernandez-Lobate "A Generative Model for Molecular Dis-
tance 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. URL: http:
PMLR, 2020, pp. html
mlr press/v119/ simm20a
[44] 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. URL: httptttt/tptsseeeee httpp:://prcceedingg mlr atres//v13// http:/hhhhbbbheecc
12

---

[45] Tarun Gogineni et al. "TorsionNet: A Reinforcement Learning Approach to Sequential Con-
former Search". In: CoRR abs/2006.07078 (2020). arXiv: 2006 07078
URL: https:
arxiv. org/ abwpaabss.....//.//////////aaa..
[46] Robin Winter, Frank Noé, and Djork-Amé Clevert. Auto-Encoding Molecular Conformations.
2021. URL: http://arxiis wttpa//wwww..gaa/a///ao1/1/ 01618
[47] Victor Garcia Satorras et al. E(n) Equivariant Normalizing Flows for Molecule Generation in
3D. 2021. eprint: 2105 2105.09016 09016 URL: htttpp//aarai..or/aab////////// httg///b/21005
09016
[48] Niklas Gebauer, Michael Gastegger, and Kristof Schütt. "Symmeetry-adapted generation 3d point sets for the targeted discovery of molecules". In: Advances in Neural In- formation Processing Systems. Ed. by H. Wallach et al. Vol. 32. Curran Associates, Inc., 2019. URL: https / proceedings neurips cc paper / 2019 file /
paf
[49] Minkai Xu et al. An EnddttoEEnd Framework for Molecular Conformation Generation via Bilevel Programming. 2021. eprint: arXiv:2105 07246,
[50] 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 DO1: 10. 1021/ acs
jeim 5b00654 URL:
https: /tt.i/doii.oo/ org 10. 1021/ acs. jcim. 5b00654
[51] 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.
URL:: Ittpps://doi. org/ 10. 1021/ ci ci10003xx
[52]
Paulo C. T. Souza et al. "Martini 3: general purpose force field for coarse-grained molecular
dynamics". In: Nature Methods 18.4 (Mar. 2021), pp. 382-388, DOL: [0. 000.10099244592000
01098-3 URL: htttps://dol org/ 10. Jiang Wang et al. "Machine Learning of Coarse- Grained Molecular Dynamics Force Fields". In: ACS Central Science 5.5 (Apr. 2019), pp. 755-767. not: 10. 20t..00022aaccctttcc..0000000111
[53]
URI.: pattps://aon orgg 1021 wttp////accccccccaaccooocaac 88000913
[54] Brooke E. Husic al. "Coarse graining molecular dynamics with graph neural networks". In:
The Journal of Chemical Physics 153.19 19 (Nov. 2020), 194101. DO1: 10. 11111666000011333..
URL: https: //doi. org 10. 1063/5 0026133
[55] Jonas Köhler
"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:
URL: https:///oi. org 10. 10210002sss.... attp///.s.cctc.......o...e.. Jet.s..30c..000015 3c00016 [56] Soojung Yang and Rafael Gómez-Bombardee Bombarelli. Chemically Transferable Generative Backmap-
10. NO.1022/acs.. jetc wtt.jjcctt.o..... 3c00016
ping of Coarse-Grained Proteins. 2023. eprint: arXiv 2303 011669
[57] Yaxin An and Sanket A. Deshmukh. "Machine learning approach for accurate backmap- ping of coarse-grained models to all-atom models". In: Chemical Communications 56.65 (2020), pp. 9312-9315. DOI: 10. 1019000//000 /039//20c0011 ddcc02651dl
URL: https:://oi.org
[58] Leandro E. Lombardi, Marcelo A. Martí, and Luciana Capece. "CG2AA: backmapping protein coarse-grained structures". In: Bioinformatics 32.8 (Dec. 2015), pp. 1235- 1237. Do1: 10
hos//bioihnformatic///br// btv740
URL: https: :/doii:///iii org
[59]
Christof Schütte et al. Hybrid Monte Carlo Method for Essential Molecular Dynamics. eng. Tech. rep. SC-98-04. Takustr. 7, 14195 Berlin: ZIB, 1998.
[60]
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-6588 DOE: 10. 1021 /021/jp00 jp0374p277
org/ 10 URL: https: //doi.
1021/jp037421y
[61] Brooke E. Husic and Vijay S. Pande. "Markov State Models: From an Art Science". In: Journal of the American Chemical Sociery 140.7 (Feb, 2018), pp. 2386-2386 DOI: 10. 1021/
7b1s.7b
CRL.: /ttttsst////oo.... org/ 10. 1021 10cs/jacs...... 7b12191
[62] Christopher Kolloff and Simon Olsson. "Rescuing off 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.
DOT: 10. 00088/2632 2153/ ad10ce. URI.:
http: /ttp://dx dx. doi. org/ 10. 20001000//222211//////// 2d1cee
13

---

[63] 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 0265-8270. ISSN: 1091-6490. DOT: 10. 1073000ass 0etmma...74880050 1704803114
URL: http:
//xx.doi.orgg /orgg//0.... 1
www.1100......... pnas 1704803114
[64] Simon Olsson and Frank Noé. "Dynamic graphical models of molecular kinetics". In: Pro-
ceedings of the National Academy of Sciences 116.30 (July 2019), pp. 15001-15006. DOI:
10.1073/ pma.../maa...mm pmas. 1901692116
URL: httpt:://oooii... org/ 10. 1073/ pma./pma..............p/.. 1901692116
[65] Andreas Mardt et al. "VAMPnets for deep learning of molecular kinetics". In: Nature Commu- nications 9.1 (2018). DOI: 10. 1038 0110://44667-011--0000
URL: htttps://ao..rrg/
1038% 4455225514666 017-7777777----111-
[66]
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. DO1: 10. 1063/06/111 4932406 4932406]
URR: https: bttpss://oiln owg. 1063/1
[67]
Hao Wu et al. Deep Generative Markov State Models. 2018. eprint: 180iv:
[68]
Leon Klein et al. Timewarp: Transferable Acceleration of Molecular Dynamics by Learning
Coarsendd Dynamics. 2023. eprint: arXiv: 01170
[69] W. K. Hastings. "Monte Carlo sampling methods using Markov chains and their applications".
In: Biometrika 57.1 (Apr. 1970), pp. 97-109. DOT: 10 1093 / biomet /57 1. 97 URI.: wtt.:////.me/....../.. https: org/10.
[70] 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. DOr:
10. 10-116/600008000 URL: httpss///od org/ 10. 1061006655... 00004/00000088...
[71] Geemi P Wellawatte, Glen M Hocky, and Andrew D White. "Neural potentials of proteins extrapolate beyond training data". In: (2023).
[72] Manuel Dibak ct al. "Temperature steerable flows and Boltzmann generators". In: Physical Review Research 4.4 (Oct. 2022). DO1: 10. 1103 ptt.s//eerreeeerrchhhhh 1042005 URL:
https:// org/10. 1 http:pyyyeeeeeere.ooo ch 4 1042005
[73] Lukas Herron et al. Inferring phase transitions and critical exponents from limited observations with Thermodynaic Maps. 2023. eprint: arXiv 2308. 14885)
[74] V Rokhlin. "Rapid solution of integral equations of classical potential theory". In: Journal of Computationall Physics 60.2 (Sept. 1985), pp. 187-207. DOI: 10. 1016/0021-999) (85)
90002-6 URL: http::ttt//o/oo/ org/ro.
[75] P. P. Ewald. "Die Berechnung optischer und
Gitterrptttatteeeeeettttntaaaee In: Annalen der Physik 369.3 (1921), pp. 253-287, DOI: 10. 10w1000022n...co 19213690304
URL: https
//doi. orooi.org//n 10. 1 andp. 19213690304
[76] Arthur Kosmala et al. Ewald-based Long-Range Message Passing for Molecular Graphs. 2023.
eprint: arxiv: 2303. 2303.
[77] Benjamin Trendelkamp-Schine
et al. "Estimation and uncertainty of reversible Markoy models". In: The Journal of Chemical Physics 143.17 (Nov. 2015), p. 174101. DOE: 10. 1063/
1.4934536 URL: https:://o// org/ 10. 1063/1 4934536)
[78] 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. DoL: 10. 100.1007/11163--11-00000 013- prg/10.cr 100000/11134-00000
URL: httpp://dag
[79] Igor Mezić. "Spectral Properties of Dynamical Systems, Model Reduction and Decomposi- tions". In: Nonlinear Dynamics 41.1-3 (2005), pp. 309-325. DOL: 10. 1007 0s10///-110010015111110000001111 005-
2824-x URL: https:://oo:: ott/://...ooooggoooo 10. 1007/2Fs11071-005-2220
[80] 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), 1), p. 515009. DOI: 10.. 1088/ / 2632 2153 / ac3dee URL: https: //doi org / 10. 1088 / 2F2632
2153%/Fac3de0
14


Reviews:
[{"review_id": "KNXDxSfNvU", "reviewer": "Reviewer_nRNz", "summary": "In 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.", "questions": "1. 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?", "limitations": "Yes, limitations are addressed.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. 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.", "weaknesses": "#### 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for taking the time to carefully read our manuscript and providing constructive input and criticism as to how we can improve it. Below, you will find a point-by-point reply to your concerns and questions. We are looking forward to continue the discussion, please do not hesitate to ask if things remain unclear.\n\n*Comments to major weaknesses*\n\nPoints 1-2: To clarify, Fig 4, does not show Boltzmann distributions but transition densities. Specifically, we show the log transition probabilities at three given time-lags (or time-horizons) from MD (coarse 2D and 1D histograms). SE3-ITO samples are overlaid with contour lines. As we are amongst the first to present work on models able to do this, we have experimented with visualization to make it as informative as possible. To clarify, the low resolution on the MD histograms are due to the high dimensionality of the transition density; it scales quadratically in the number of bins. Recall that we need to compute the probability from any bin to any other bin. For these 16x16 histograms we divide the data into 65536 transition counts, with 750000 MD samples in our training data we doubling the resolution would give us less than one sample per bin on average. In the plots we show only the transition statistics from the bin to which the initial condition is assigned. \n\nWe would like to emphasize that, we do also compare SE3-ITO samples (blue and orange) and MD samples (black) in the marginal histograms. \nNevertheless, we recognize the reviewers' concern about the comparison of statistics shown with different resolutions, and will include a 1- and 2-dimensional histogram of the model and training samples for more direct comparison, for the camera-ready version. 2D histograms will be separated into individual plots.\n\nPoint 3: In the Global response we show why an SE(3) equivariant model is necessary for the molecular applications pursued here (figs 3-4). Briefly, if we use a parity invariant model, we will sample mirror images of molecules with equal probability a priori, even if one of the mirror images is inaccessible in the physical dynamics used to simulate the system. This is also a shortcoming observed in timewarp (arxiv:2302.01170), that uses a permutation equivariant architecture. While it is interesting to investigate and compare different SE(3) equivariant architectures for SE3-ITO models, we are confident that these are best addressed in future work, as such a benchmark would add little scientific value to the current manuscript\n\nPoint 4: After submission we found our models were slightly underfitted, and after training longer we found improved agreements with observables across all proteins. We are providing new values in the global response table 2, for Chignolin, along with convergence plots. We will update all values for the camera ready version.\n\n\n_Questions:_\n1. Yes this is correct. It is sampled directly with \\Delta t = 1000\\tau.\n2. We have to be careful comparing N across different data-sets, with different resolutions in space and time. In general, we cannot make a statement as to whether fast (‘immediate’) dynamics are captured poorly by our ITO models. Our preliminary investigations of this (Fig 10, in the supplement) suggest that we indeed do not capture very fast dynamics perfectly. However, we argue that these fast dynamics do not need a model like ITO. We can already study fast dynamics very well with conventional MD simulations. Their limitation is the mixing between modes interconnected by low probability barriers. On that task ITO does extremely well."}, {"review_id": "7OmNNTnwKc", "reviewer": "Reviewer_BcfP", "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. ", "questions": "- how are the number of time scales chosen during the training?  \n- could you explain in the caption or here all the input for ITO networks in Figure 2\n- 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. \n- I am curious how it can be generalized to other dynamics like some other physics system that is also multiscale. ", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- a novel framework for multi-scale simulation for molecular dynamics, combined with the SE3 network for stable, long-term results. \n- demonstrate good experimental results on multiple MD datasets and achieve an order of magnitude speedup\n- work on an important task of molecular dynamic modeling \n- use diffusion network for exact likelihood evaluation \n- 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.\n- ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their favorable and supportive evaluation and helpful comments and questions on our paper. We reply to the questions below and welcome a discussion about any outstanding doubts.\n\n_Questions:_\n\n1. The number of time-scales chosen during training, can be understood as how many linear combinations of the eigenfunctions we want the model to see during training (eq.8). We set this number (N_max) to 1000 initially, and we did not experiment with varying this value. We leave optimization of this hyper parameter to future work.\n\n2. z is an atom-embedding (nominal embedding), N is the multiple of the physical time-step \\tau (positional embedding), x_0 is the cartesian coordinates of the conditioning state, t_diff is the diffusion time, and x_{0+N\\tau} is the partially denoised state of the time-lagged configuration given x_0. Note that there is a hat missing on the   x_{0+N\\tau} to be consistent with the main text, we will fix this for the camera ready version. B) Variable names have the same interpretation, and there is also a hat missing on x_{0+N\\tau}.\n\n3.  The diffusion time is the progress along the diffusion process which models the conditional probability density p(x_N\\tau | x_0), the physical time (“actual time”) is the time-step N\\tau, and is connected to the physical process which we are modeling. \n"}, {"review_id": "YzrGyiPc1D", "reviewer": "Reviewer_Dude", "summary": "This 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.", "questions": "- 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?", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough review of our manuscript and the comments, in particular regarding the presentation of the maths and comparison to previous work. We believe the input will help us prepare a much improved version for the camera ready version. \n\n_Comments on weaknesses:_\n\n1. There was indeed a typo in eq. 3 \\rho and f were mixed up in that equation, we will fix this in the camera ready version. Great catch! The reviewer here refers to the propagator which is mathematically different from the transfer operator, but missed the scaling by \\mu which makes the difference. Our eq. 17 and we refer to A.1, for a more verbose derivation in bra-ket notation.\n2. There are no other methods available which allow for sampling on these time-scales and allow an apples to apples comparison. DeepGenMSMs are essentially deep Markov models, which encode configurations as state memberships and then model transition densities between discrete states in a latent space, with a Koopman operator. As shown in their original paper, these methods introduce severe distortions in the local structures when long time-scales are queried. So while DeepGenMSMs have realistic latent space dynamics, the real space dynamics cannot be faithfully reconstructed.  ITO does not suffer from this drawback.  The only methods which allow for sampling of long time-scales faithfully are conventional MD simulations and timewarp (arxiv:2302.01170); the latter paper is in preprint with no code or data available publicly, and suffer other severe drawbacks. Nevertheless, we recognize that expanding the section on prior work will help clarify the ITO contributions to the readership better.\n3. We provide a figure in the global reply showing how a E(3) equivariant model samples both mirror images of alanine dipeptide with equal probabilities, whereas only one mirror image is accessible in the physical dynamics used to train the model. Introducing SE(3) equivariance overcomes this problem.  \n\n_Replies to questions:_\n\n1. Indeed, we do not directly use eq 4, 5 and 8, as they would involve optimization on the Stiefel manifold, and would require us to choose the number of eigenfunctions to estimate. We therefore model these expressions implicitly, by modeling the transition density. We use these equations to motivate the connection between the time-horizon N\\tau and transition probability between two configurations x and y. We conjecture that we can leverage this underlying mathematical structure to instead train a single model which models the transition density at multiple different time horizons [1, N_max], since the proportions of different functions vary exponentially in N. Indeed, we see that training a model with multiple N is an advantage: we get better prediction of meta-stability when training with multiple time horizons.\n2. No, we have not experienced any negative impact of large time-scales in our experiments. Since we converge to the Boltzmann distribution for large N, e.g. the state x_N\\tau ~ \\mu(x) independently of x_0, this is the hallmark of ill posedness, but it is also what we expect from the dynamics: at sufficiently long time-scales we expect our Markov chain to mix instantaneously.\n3. We are illustrating two collective variables tIC1 and tIC2 in our analyses. The dashed lines correspond to the tiC1 and tiC2 coordinates of the folded state. For clarity, we will change the style of the dash to distinguish the two tICs for the camera-ready version.\n4. We expect these values to align within an order of magnitude. However, after submission, we found that training the models longer, led to better agreement between the MD and ITO values. We now provide updated values and show a plot of convergence of observables as a function of training time in the global reply for Chignolin. We will include these updated values for all four proteins in the camera ready version.\n"}, {"review_id": "L3oMlXjB79", "reviewer": "Reviewer_r4Zz", "summary": "The 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!", "questions": "### 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?", "limitations": "The 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.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 1, "contribution": 2, "strengths": "### 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking time to read our manuscript thoroughly, and providing their insightful comments. We in particular appreciate the comments about the presentation, which will help us prepare a much improved camera ready version. We are addressing their highlighted weaknesses and questions in a point-by-point reply below.\n\n_Comments on weaknesses:_\n1. We show for MB as it is the system which can be extensively sampled using conventional methods and where statistics can be computed with confidence. We have trained models for Alanine dipeptide with fixed lags and include VAMP2 gaps calculated for these models as well for completeness, preliminary figures are attached in the Global response (Fig 1), and ecco the results reported for MB. Final figures will be included in the camera ready version. KL divergences are insensitive to low density regions per construction, capturing slow dynamics means capturing low density regions well. VAMP2-gaps compare the singular values of a Koopman operator and directly measures meta-stability, e.g. how well we capture the gaps between high and low density regions of the configuration space. We used trajectories of length 1000\\tau to generate the statistics in the table, to ensure the same statistics across the lag times. The estimation of the Koopman operator is not stable when the lag is equal to the length of the trajectory, so we cannot compute the VAMP2-gap reliably at 1000\\tau. We will redo these calculations for the camera ready versions and include all lags.\n2.  We provide estimates of wall-clock time per conditional sample for the systems in the appendix C.2. We also discuss this briefly in section 4.3 in the main text for one example (Chignolin). The main difference is that for MD, we need to run simulations for N\\tau steps to draw one sample for the Monte carlo estimator (2. Dynamic observables). With ITO we can draw one sample at N\\tau in a 1-shot fashion, and for coarse-grained Chignolin drawing one sample takes ~4 ms. If we very conservatively say we can simulate a coarse-grained Chignolin at 100 microseconds per day with similar GPU resources, and want to evaluate a correlation function at N\\tau = 1 microsecond, it means we can get 100 independent samples per day for a TitanX, for ITO we would get on the order of 10^7 independent samples. We realize that this discussion has been lackluster in the manuscript and will improve it for the camera ready version.\n3. We will provide code when the paper is accepted for publication. \n4.  We have been in active conversation with Prof. Anandkumar and solicited feedback from her group about ITO, and while we got pointers to work on ODEs and PDEs using their Neural Operators approaches, they did not so far investigate transfer operators to our knowledge. If the reviewer can point us to specific references we are happy to expand this discussion.\n\n_Replies to questions:_\n1.  Interesting idea, we have not explored this. We encode the physical time (N) using a positional encoding, which then defines the conditional diffusion model, and in turn the distribution of the random variable x_N\\tau | x_0.\n2.  No. However, we are using ODE solvers to solve the probability flow ODE, for efficient inference of the trained model. We discuss this in the appendix and provide timings for inference.\n3.  The uncertainties are estimated using Bayesian posterior sampling of Markov models consistent with the data (using procedure described in Trendelkamp-Schroer and Noe JCP 2013). We will clarify this in the camera ready version.\n4.  We have not experimented with varying the number of diffusion steps for different N. However, the motivation to keep it the same is that the conditional density at different N are combinations of the same functions but weighed differently (see. Eq. 8, and Appendix A.1.). Since we do not know the weights a priori, we delegate this to a positional encoding. The hope was – and what we see is – that ITO is indeed able to generalize across time horizons (N), even if we do not represent the operator explicitly.\n5. Each plot in Figure 4 is calculated using 15,000 trajectories. For direct sampling, this means that we had to take 15,000 samples in total. In the case of ancestral sampling, we had to sample 4/64/512 steps for 15,000 trajectories, depending on the value of \\Delta t.\nThe contour lines of the fast folding proteins in Figure 5-8 were calculated using 10,000 trajectories. The density in the background are samples from the reference trajectories. The simulation step is 200ps and the simulations were performed with ancestral sampling. This means that in order to plot the contour lines of the \\Delta t = 200ns we had to perform 1000 steps. We will clarify these details for the camera ready version of the manuscript.\n6. The i index in the context of the mini-batch training is indicates a tuple of B_i=(x_{t_i}, x_{t_i+N_i\\tau}, N_i), where the t_i and N_i are sampled randomly for each batch, i. The i is an index in the batch, which is composed of conditioning state, time-lagged state and time-lag (horizon). We realize this notation is not standard and might be confusing. We are revising it for the camera ready version.\n"}, {"review_id": "68HQShujSQ", "reviewer": "Reviewer_Ru8L", "summary": "The 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.", "questions": "- 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.", "limitations": "Limitations 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.", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The presented method has strong theoretical foundations and the underlying theory is well-described.", "weaknesses": "While 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our manuscript and providing insightful comments and questions. We believe your input will help us shape a much improved camera ready version. Please find commentary and replies to your concerns and questions below.\n\n_Comments on weaknesses:_\n1 and 2. We believe that the modification, while indeed “simple” , is still a contribution as it solves a specific problem we had: we needed a highly efficient architecture with SE(3) equivariance (see clarification below). The values in Table 2 are generally very good considering there are no other methods which offer estimates of these values at all. Nevertheless, after submission, we found that training the models longer, led to better agreement between the MD and ITO values, suggesting that we systematically underfitting our models due to time-constraints. We now provide updated values for Chignolin and show a plot of convergence of observables as a function of training time in the global reply (Global reply fig 2 and table 2). We will update all values for the camera ready and provide convergence plots in the updated appendix.\n\n_Replies to questions:_\n1. E(3)-equivariant models do not have the capacity to distinguish between chiral configurations. As a consequence, if we don’t use an SE(3)-equivariant score model, ie. symmetry is not broken when inverting the orientation, half of the predicted configurations will have the opposite chirality of the conditioning configuration. This is a structural problem because the model can not tell a configuration and its mirror image apart. In the Global response we illustrate this problem in Figs. 3 and 4. We will elaborate this discussion for the camera ready version.\n2. This is indeed interesting work (two for one, Arts et al). We expect for large N that -\\nabla_{x_N\\tau} log p(x_N\\tau|x_0) we will approach the effective force field since the density will converge to the Boltzmann distribution (see our appendix A.1). We have not tried to extract the forcefield from our current models, as learning a machine learned potential is not the focus of this work and we believe it would distract attention from the main contributions of this work. We do not believe that the learned forcefield would necessarily be amenable for simulation steps as large as those ITO enables, it is unclear what the reasoning here is.  Since our focus is on sampling and avoiding MD simulations, we leave these experiments for future work. \n3.  As outlined above, we found that we had systematically under fitted our models due to time-constraints. After submission we found much better agreement, and show convergence of observables in our global response.\n4. In the manuscript we do provide timings of sampling. The sampling speed in the current implementation declines quadratically in the number of particles for the simulation (appendix C.3, and limitation section) whereas MD scales from Nlog(N) to N^3 depending on different factors. Even so, with ITO we can stably simulate at a rate 6 orders of magnitude faster than MD. Nevertheless, current timings are not directly comparable, as MD is subject to other degrees of freedom which we ignore for now. To be specific, since the model is not transferable currently, we need MD data before we can do sampling for each system. In the strict sense, this means that in its current instance, there is no gain compared to MD, for an arbitrary new case. Nevertheless, we argue that by systematic data curation and parameter sharing, we can achieve a transferable ITO model in the near future which can capitalize on the multiple orders of magnitude speed ups presented. We transparently discuss this limitation in our manuscript and believe that in spite of it, the advances presented in this paper are extremely interesting for the broader AI4Science community.\n"}]

Author rebuttal:
Global response. see attached pdf.

Meta-review: The reviewers raised many concerns about the presentation of the paper. They all agree that the paper has some value but also agree that the paper, at its present shape, is not ready for submission. 


We encourage the authors to consider the feedback provided in the review process.
Decision: Accept (poster)