Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics
Authors: Leon Klein
Andrew Y. K. Foong
Tor Erlend Fjelde
Bruno Kacper Mlodozeniec
Marc Brockschmidt
Sebastian Nowozin
Frank Noe
Ryota Tomioka
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Molecular Dynamics
Normalizing Flows
MCMC
TL;DR: We learn a transferable flow model that predicts the future state of a molecule to accelerate Molecular Dynamics simulations

Abstract: *Molecular dynamics* (MD) simulation is a widely used technique to simulate molecular systems, most commonly at the all-atom resolution where equations of motion are integrated with timesteps on the order of femtoseconds ($1\textrm{fs}=10^{-15}\textrm{s}$). 
MD is often used to compute equilibrium properties, which requires sampling from an equilibrium distribution such as the Boltzmann distribution. 
However, many important processes, such as binding and folding, occur over timescales of milliseconds or beyond, and cannot be efficiently sampled with conventional MD.
Furthermore, new MD simulations need to be performed for each molecular system studied.
We present *Timewarp*, an enhanced sampling method which uses a normalising flow as a proposal distribution in a Markov chain Monte Carlo method targeting the Boltzmann distribution. 
The flow is trained offline on MD trajectories and learns to make large steps in time, simulating the molecular dynamics of $10^{5} - 10^{6} \textrm{fs}$.
Crucially, Timewarp is *transferable* between molecular systems: once trained, we show that it generalises to unseen small peptides (2-4 amino acids) at all-atom resolution, exploring their metastable states and providing wall-clock acceleration of sampling compared to standard MD.
Our method constitutes an important step towards general, transferable algorithms for accelerating MD.

Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics
Leon Klein Freie Universität Berlin leon. klein@fu-berlin.....
Andrew Y. K. Foong"
Tor Erlend Fjelde University of Cambridge tef300cam ac
Microsoft Research A14Science
autreexoong@mlccoooff..com
Bruno M
Mare Brockschmidt
Sebastian Nowozin'
University of Cambridge
bkm288cam ac uk
Frank Noć
Ryota Tomioka
Microsoft Research A14Science
Microsoft Research A14Science
Freie Universität Berlin
ryoto@mi crosoft. com
Rice University
Abstract
Molecular dynamics (MD) simulation a widely used technique to simulate molecular systems, commonly the all-atom resolution where equations of motion are integrated with timesteps on the order of femtoseconds (Ifs = 10-15s). MD often used to compute equilibrium properties, which requires sampling from an equilibrium distribution such as the Boltzmann distribution. However, many important processes, such as binding and folding, occur over timescales of milliseconds or beyond, and cannot be efficiently sampled with conventional
MD, Furthermore, new MD simulations need be performed for each molecular system studied. We present Timewarp, an enhanced sampling method which uses a normalising flow proposal distribution in Markov chain Monte Carlo method targeting the Boltzmann distribution. The flow is trained offline on MD trajectories and learns to make large steps time, simulating the molecular dynamics of 10 1 10 fs. Crucially, Timewarp is transferable between molecular systems: once trained, we show that generalises to unseen small peptides (2-4 amino acids) at all-atom resolution, exploring their metastable states and providing wall-clock acceleration of sampling compared to standard MD, Our method constitutes an important step towards general, transferable algorithms for accelerating MD.
Introduction
Molecular dynamics (MD) is well-established technique for simulating physical systems the atomic level. When performed accurately, it provides unrivalled insight into the detailed mechanics of molecular motion, without the need for wet lab experiments. MD simulations have been used to
understand processes of central interest in molecular biophysics, such as protein folding 29 22] protein-ligam binding [ and protein-protede association [3I. Many crucial applications of MD boil down to efficiently sampling from the Boltzmann distribution, i.e., the equilibrium distribution
"qqual contribution.
Work done while Microsoft Research.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

a)
b)
MD
Timewarp MCMC
xt)
TT)
0.5 x 10
Figure 1: (a) Initial state æxt) (Left) and accepted proposal state ztt po(x po(x (z)((()) (Right) sampled with Timewarp for the dipeptide HT (unseen during training). (b) TICA projections of simulation trajectories, showing transitions between metastable states, for short MD simulation (Leff) and Timewarp MCMC (Righ(), both run for 30 minutes of wall-clock time. Timewarp MCMC achieves speed-up factor of as 33 over MD in terms of effective sample size per second.
of a molecular system at a temperature T. Let (1), (I),=(())) x(() R&N be the state of the molecule at time f. consisting of the positions =²(t) E R3V and velocities x x²(t) E RaN of the N atoms in Cartesian coordinates. The Boltzmann distribution is given by:
1(=,, xx)
x exp (=)TT(U()))) + K(z")))
m(x²) = p(>", "(x",x")ee dx".
xx)
(I)
where U(xP) the potential energy, ((x") the kinetic energy, and kg is Boltzmann's constant.
Many important quantities, such as the free energies of protein folding and protein-ligand binding, can be computed as expectations under ((xx)). popular MD method to sample from (xx"))
Langevin dynamics [19,, which obeys the following stochastic differential equation (SDE):
madx." - d. "m;," di + V2m.7knTdB.
(2)
Here indexes the atoms, m, is the mass of atom i, U(*²) is the potential energy, 7 is a friction parameter, and dB, is a standard Brownian motion process. Starting from an initial state x0)), simulating Equation (2). along with the relationship dx
dt, yields values of xxt) that are distributed according to the Boltzmann distribution as t + co. Standard MD libraries discretise this SDE with a timestep At. which must be chosen to be 1fs- 10-15s for stability. Unfortunately, many biomolecules contain metastable states separated by energy barriers that can take milliseconds of MD simulation time (~ 10 sequential integration steps) to cross, rendering this approach infeasible. To overcome this, prior work has produced an array of enhunced sampling methods, such as coarse graining 3 14 and metadynamics [8]. However, these methods require domain knowledge specific to each molecular system to implement effectively.
Standard MD simulations do not transfer information between molecular systems: for each system studied, a new simulation must be performed. This is wasted opportunity: many molecular systems exhibit closely related dynamics, and simulating one system should yield information relevant to similar systems. In particular, proteins, being comprised of sequences of 20 kinds of amino acids, are prime candidates study this kind of transferability. We propose Timewarp. general, transferable enhanced sampling method which uses normalising flow Monte Carlo (MCMC) method targeting the Boltzmann distribution. Our main contributions are:
(35 as a proposal for Markov chain
1. We present the first MI. algorithm working in general Cartesian coordinates that demon- strates transferability small peptide systems unseen during training.
2. We demonstrate, for the first time, wall-clock accelcration of asymptotically unbiased MCMC sampling of the Boltzmann distribution of unseen peptide systems.
3. We define an MCMC algorithm targeting the Boltzmann distribution using a conditional normalising flow as proposal distribution, with Metropolis Hesttpolls-Hastinss (MH) correction
step to ensure detailed balance (Section[3.4)
4. We design a permutation equivariant, transformer-based normalising flow.
5. We produce a novel training dataset of MD trajectories of thousands of small peptides.
6. We show that even when deployed without the MH correction (Section 3.5), Timewarp can be used to explore metastable states of new peptides much faster than MD.

---

The code is available here: httpes//gtthuu..ooo//cooooo///// able upon reques
The datasets are avail-
2 Related work
There has recently been surge of interest in deep learning on molecules. Bolizmann generators 28 17 [5] use flows to sample from the Boltzmann distribution asymptotically unbiased. There are two ways generate samples: (i) Produce i..d. samples from the flow and use statistical reweighting to debias expectation values. (ii) Use the Boltzmann generator in an MCMC framework [4), as in Timewarp. Currently, Boltzmann generators lack the ability to generalize across multiple molecules, in contrast to Timewarp. The only exception is 11, who propose a diffusion model torsion space and use the underlying ODE as transferable Boltzmann generator. However, in contrast Timewarp, they use internal coordinates and do not operate in the all atom system. Moreover, [13 25 recently introduced Boltzmann generators in Cartesian coordinates for molecules, potentially enabling transferable training. Recently, 46 proposed GeoDiff. a diffusion model that predicts molecular conformations from molecular graph. Like Timewarp, GeoDiff works Cartesian coordinates and generalises to unseen molecules. However, GeoDiff was not applied to proteins. but small molecules, and does not target the Boltzmann distribution. In contrast to Timewarp, [37] learn the transition probability for multiple time-resolutions, accurately capturing the dynamics. However, they do not show transferability between systems. Most similarly to Timewarp, recent work, B trained a transferable ML model to simulate the cine-coareened dynamics of polymers. However, unlike Timewarp, their model acts on coarse grained representations. Additionally, was not applied to proteins, and there is no MH step, which means that errors can accumulate in the simulation without being corrected.
Markov state models (MSMs) 32 10 work by running many short MD simulations, which are used define a discrete state space, along with an estimated transition probability matrix. Similarly to Timewarp, MSMs estimate the transition probability between the state at time and the time t T, where > At. Recent work has applied deep learning to MSMs, leading to VAMPnets 23 and deep generative MSMs 455, which replace the MSM proc--pro pipeline with deep networks. contrast to Timewarp. these models are not transferable and model the dynamics in coarse-grained, discrete state space, rather than in the all-atom coordinate representation.
There has been much previous work on neural adaptive samplers 38 20
which use deep generative models as proposal distributions. A-NICE-MC [38] uses a volume-preserving flow trained using likelihood-free adversarial method. Other methods objective functions designed to encourage exploration. The entropy term in our objective function is inspired by [4]]. In contrast to these methods, Timewarp focuses on generalising to new molecular systems without retraining.
Numerous enhanced sampling methods exist to for MD, such as parallel tempering (39 6 or
proposing updates of collective variables along transition paths [18] 27). Given Timewarp's ability to accelerate MD, it often offers the opportunity to be integrated with these techniques.
3 Method
Consider the distribution of (t++) induced by an MD simulation of Equation (2) for time At, starting from st). We denote this conditional distribution by p(x(t + (+)) Timewarp uses deep probabilistic model to approximate p(r(t + +t)= (see Figure [}. Once trained, the model is used in an MCMC method to sample from the Boltzmann distribution.
3.1 Conditional normalising flows
We fit conditional normalising flow, po + ())= p(x((t + r))=(t)), where are learnable parameters. Normalising flows are defined by a base distribution (usually a standard Gaussian), and diffeomorphism f. i.e. a differentiable bijection with differentiable inverse. Specifically, we po po(x[1 + ())=(()) as the density of the distribution defined by the following generative process:
2, ~,2"~ N(0,1)
(t + = fo(*" ((<",,"" 2"; x (t),x"(t))
(3)
'Please contact anddrewfooogg@miccooott...t com for dataset access.
3

---

Here 2? E R3N and 2" E RAN, For all settings of and x(t), fo(-)x(t)) is a diffeomorphism that takes the latent variables (z², 2²) E RON to 7), 7)) R The conditioning state atl) parameterises a family of diffeomorphisms, defining conditional normalising flow [44]. Note that there are no invertibility constraints on the mapping from æ(t) to the output xxtt + r), only the
map from to stt + 7) must be invertible. Using the change of variables formula, we can evaluate:
po(all + (1)) N (Jo (a(t + +);= det 11'((=200 (z(t + +))
where ( (z())) J6 (=0(t)) (1(t + T)) denotes the Jacobian of (=;()))(() evaluated at sr( T)).
TON - R&N is the inverse of the diffeomorphism sol ;a(()), and
3.2 Dataset generation
We generate MD trajectories by integrating Equation 2 using the OpenMM library 7
We simulate small proteins (peptides) in implicit water, i.e., without explicitly modelling the degrees of freedom of the water molecules. Specifically, we generate a dataset of trajectories D = {Tr) where P is the number of peptides. Each MD trajectory is temporally subsampled with spacing of so that T = (=()),=(T), (2)).....) During training, we randomly sample pairs xxt), att + 7) from D. Each pair represents sample from the conditional distribution juz((t + r) [xx(t). Additional details are provided in Appendix Since the flow is trained on trajectory data from multiple peptides, we can deploy it at test time to generalise to new peptides not seen in the training data.
3.3 Augmented normalising flows
We are typically primarily interested in the distribution of the positions I", rather than the velocities r². Thus, is not necessary for x (t), (t++>) to represent the actual velocities of the atoms Equa- tion (3). We hence simplify the learning problem by treating x" as non non-physical auxiliary variables within the augmented normalising flow framework [9. For each datapoint z(t) = (t), x x²()) in D, instead of obtaining x² (1) by recording the velocities in the MD trajectory, we discard the MD velocity and independently draw x= (1) ~ .N, 7). The auxiliary variables xl) now contain no information about the future state x² (1++)).2" (T) since (t) and x (L++7) are drawn independently. Hence we can simplify fe to depend only on (1), with (t + (t 7) == So(2", x²;x We include auxiliary variables for two reasons: First, they increase the expressivity of the distribution for rx without a prohibitive increase in computational cost [9] 22. Second, constructing a conditional flow that respects permutation equivariance is simplified with auxiliary variables see Section
3.4 Targeting the Boltzmann distribution with asymptotically unbiased MCMC
After training the flow (at( t + (t)),,), we use it as proposal distribution in an MCMC method to target the joint distribution of the positions xP and the auxiliary variables x", which has density:
Paug (x= x²) x exp (-
U(=")) K
N (a
(4)
Starting from an initial state Xo = (XP- X²) E TRON for state m 0, we iterate:
Xm
with probability a(Xm; Xm)
Xm ~ pol (XP), Xm+1 =
(5)
Xm with probability 1 1- 1-a(X a(xmm Xm)
where all Xm, Xm) is the Metropolis-Hasting (MH) acceptance ratio
24 targeting Equation (4
o(x, x) = -min (1, Maug (x) ))) (X XP)
(6)
Maug (X) X)pp X XP)
The flow used for pe must allow for efficient sampling and exact likelihood evaluation, which is crucial for fast implementation of Equations (5) and ( Additionally, after each MH step, we resample the auxiliary variables X using a Gibbs sampling update:
XXP X + (XP G) c c~ I.,
(7)

---

Iterating these updates yields sample Xy, ~ Paug as m - 00. To obtain Boltzmann- distributed sample of the positions X J4, we simply discard the auxiliary variables X As sending + oo is infeasible, we simulate the chain until fixed budget is reached. In practice, we find that acceptance rates for our models can be low, around 1%. However, we stress that even with a low acceptance rate, our models can lead to faster exploration if the changes proposed are large enough, as we demonstrate in Section 6 Furthermore, we introduce a batch sampling procedure which significantly speeds up sampling whilst maintaining detailed balance. This procedure samples a batch of proposals with a single forward pass, and accepts the first proposal that meets the MH acceptance criterion. Pseudocode for the MCMC algorithm is given in Algorithm I in Appendix C
3.5 Fast but biased exploration of the state space without MH corrections
Instead of using the MH correction to guarantee asymptotically unbiased samples, we can opt to use Timewarp in simple exploration algorithm. In this case, we neglect the MH correction and accept all proposals with energy changes below some cutoff. This allows much faster exploration of the state space, and in Section 6 we show that, although technically biased, this often leads to qualitatively ac- curate free energy estimates, It also succeeds in discovering all metastable states orders of magnitude faster than Algorithm and standard MD, which could be used, e.g, provide initialisation states for a subsequent MSM method. Pseudocode is given in Algorithm 2 in Appendix [
4 Model architecture
We now describe the architecture of the flow fo(3P, ="; (1)), which shown in Figure 2
RealNVP coupling flow Our architecture is based on RealNVP [5), which consists of stack of coupling layers which affinely transform subsets of the dimensions of the latent variable based on the other dimensions. Specifically, we transform the position variables based on the auxiliary variables, and vice versa. In the fth coupling layer of the flow, the following transformations are implemented:
z = o (= x² (t)) @ =f + (1)))))()))),
2 = 6 (t)) O o + (
Going forward, we suppress the coupling layer index L. Here @ is the element- product, and 5 -.N + R&N is our atom transformer, neural network based the transformer architecture [31 that takes the auxiliary latent variables 2" and the conditioning state rxt) and outputs scaling factors for the position latent variables 2.. The function th- R3N - R3N is implemented as another atom transformer, which uses z" and xtf) to output a translation of the position latent variables >?.
The affine transformations of the position variables (in Equation (8) are interleaved with similar affine transformations for the auxiliary variables (in Equation (9). Since the scale and translation factors for the positions depend only on the auxiliary variables, and vice versa, the Jacobian of the transformation is lower triangular, allowing for efficient computation of the density. The full flow So consists of Ncoupling stacked coupling layers, beginning from z ~ NN, I) and ending with a sample from pe (z(t + r) (t)) This is depicted in Figure 2 Left. Note that there is a skip connection such
that the output of the flow predicts the change x(f + +T) - x(t,, rather than x(t + 7) directly.
Atom transformer We now describe the atom transformer network. Let az (t), 2f, 2²', all elements of R.3, denote respectively the position of atom in the conditioning state, the position latent variable for atom i, and the auxiliary latent variable for atom i. To implement an atom transformer which takes= 2" as input (such as (2", x." (1)) and 1"((", x>(()) (()) Equation (8) we first concatenate the variables associated with atom i. This leads to a vector a; == (t), hú, Z E #6 where x has been excluded since sy. tlj are not allowed depend on 2". Here h E " is a learned embedding vector which depends only on the atom type. The vectors a; are fed into an ML. Φin : TR. H+ - RR. where D is the feature dimension of the transformer. The vectors onn (a1)..... (a.()) are then fed into Niransformer stacked transformer layers. After the transformer layers, they are passed through another atom-wise MIP, dout RD + RR3 The final output is in I as required. This is depicted in Figure? 2 Middle, When implementing si and lg from Equation ( a similar procedure is performed on the vector [x.] (1), h4, =?]T, but now including =f and excluding 2/- There are two key differences between the atom transformer and the architecture in (3 First, to maintain permutation equivariance,
5

---

())) # ~ (10,1)
2 N(0,1)
#=(6) = [ P 1
Outputs
Linear
Linear
Linear
R
transformer
Kernel sel--atintion
lx]
Norging
transforing mer
Multihedd kernel aulf-atantion
Nuudeas
Atom Insisslormt
Add & normatzee
Adantionn weight ( |=| +71//5)
Altertion vee get oc exy!' ||x +/1
Alomvuten MLP
Afrerton xiexp(--e Lo +118/4)
Add & nomalize
transfororr
Concal
Alam-uiia MLP
tnear
( +)
Figure 2: Schematic illustration of the Timewarp conditional flow architecture, described in Section 4 Left: A single conditional RealNVP coupling layer. Middle: single atom transformer module. Right: the multihead kernel self-attention module.
we do not use positional encoding, Second, instead of dot product attention, we use a simple kernel self-attention module, which we describe next.
Kernel self-attention We motivate the kernel self-attention module with the observation that physical forces acting on the atoms in a molecule are local: i.e., they act more strongly on nearby atoms. Intuitively, for values of that are not too large, the positions time t + will be more influenced by atoms that are nearby at time i, compared to atoms that are far away. Thus, we define the attention weight way for atom attending to atom j as follows:
exp( - | x
Wij
(10)
N Σ exp(-11a3 2 112
where é is lengthscale parameter. The outputs i given the inputs ( are then:
N Fout, ΣE Wij V Fin
(11)
where V [g Xd. learnable matrix. This kernel self-attention is instance of the RBF kernel attention investigated in [42]. Similarly to [43], we introduce multihead version, where each head has different lengthscale. This illustrated in Figure 2 Right. We found that kernel self-attention
was significantly faster to compute than dot product attention, and performed similarly.
4.1 Symmetries
The MD dynamics respects certain physical symmetries that would be advantageous to incorporate. We now describe how each these symmetries incorporated in Timewarp.
Permutation equivariance Let be permutation of the N atoms. Since the atoms have intrinsic ordering, the only effect of a permutation of ztf) on the future state æ(t + 7) is to permute the atoms similarly, i.e.,
p(ar(t + +))oz(()) = =((a(t+))(x(t)
(12)
Our conditional flow satisfies permutation equivariance exactly. To show this, we use the following proposition proved in Appendix A.I which is an extension of [16] 36) for conditional flows:

---

Proposition 4.1. Let be symmerry action. and let f: be an equivariant map such that f(oz; (2;x) for all o, 2, Further, let the base distribution p(z) satisfy p(oz) p(z) for all a, Then the conditional flow defined by ~ p(z), xlt + T) == f(=;a(()) satisfies plax(t ( ())=(())) p(x(t +))((t))
Our flow satisfies (o [o(());xx(()) (z; x())) since the transformer is permutation equivariant, and permuting 2 and 2r(t) together permutes the inputs. Furthermore, the base distribution p(z) =N T) is permutation invariant. Note that the presence of auxiliary variables allows to easily construct permutation equivariant coupling layer.
Translation and rotation equivariance
Consider a transformation T = (R, that acts on xP:
Tx.- = Roç + a,
ISis
(13)
where 3 x 3 rotation matrix, and a E IR3 is a translation vector. We would like the model to (<xr((+)+==(())) We achieve translation equivariance by subtracting the average position of the atoms in the initial state (Appendix [A.)). Rotation equivariance is not encoded in the architecture but is handled by data augmentation: each training pair (a(t), x(t + 7))
satisfy pe (T>x((+7)(7xx(()
is acted upon by a random rotation matrix R to form (Rr(t), Ro(t + r)) in each iteration.
5 Training objective
The model is trained in two stages: (i) likelihood training, the model is trained via maximum likelihood on pairs of states from the trajectories in the dataset. Let k: index training pairs, such that (x(00)) (t), x (k) (t + T)) represents all pairs of states at times T apart in D. We optimise:
(0) = 1/2 Σk log pa (2) (t ( (1)).
(14)
(ii) acceptance training, the model is fine-tuned to maximise the probability of MH acceptance, Let (x() (t) be sampled uniformly from D. Then, we use the model to sample ((k) (t) pa( (=))(((((())
using Equation B). We use this to optimise the acceptance probability in Equation (() with respect to
Let (X, X) denote the model-dependent term in the acceptance ratio in Equation (6):
ro (x X) = Mang ( x)po( X2)
(15)
Maug x))p((e (X XP)
The acceptance objective is then given by:
Lace (a))(0)= = R E log ra (1), (()) (1+
(16)
Training maximise the acceptance probability can lead the model proposing changes that
too small: if (=(k)) (t T) = x (t). then all proposals will be accepted. To mitigate this, during
acceptance training, use objective which is weighted average of (oce Luk (0) and Monte Carlo estimate of the average differential entropy.
Lent (0) == 1 K log po (T (k) (t
T) (t)).
(17)
6 Experiments
We evaluate Timewarp on small peptide systems. To compare with MD, we focus on the slowest transitions between metastable states, as these are the most difficult traverse. To find these, we use time-lagged independent component analysis (TICA) [30]. linear dimensionality reduction technique that maximises the autocorrelation of the transformed coordinates. The slowest components, TIC 0 and TIC 1, are of particular interest. To measure the speed-up achieved by Timewarp, we compute the effective sample size per second of wall-clock time (ESS/s) for the TICA components. The ESS/s is given by
Metf
M
ESS//-
(18)
tsampling
*sampling (1 + 2 2 2
7

---

MID
Timewar
Free energy projections
MD
Timewarp
Free energy By projections MD RORE
Free enegyaat
QW
Aii
11C0 m
PM
Psi
PM PM
Matlehhckk
Figure 3: Left half: Alanine dipeptide experiments. (a) Ramachandran plots for MD and Timewarp samples generated according to Algorithm [1 (b) Free energy comparison for the two dihedral angles sp and v.. (c) Ramachandran plots for the conditional distribution of MD compared with the Timewarp model. Red cross denotes initial state. (d) Time dependence of the ip dihedral angle of MD and the Markov chain generated with the Timewarp model. Right half: Experiments on 2AA test dipeptides QW (top row) and HT (bottom row). (e) TICA plots for a long MD chain and samples generated with the Timewarp MCMC algorithm (Algorithm D. (f) Free energy comparison for the MD trajectory, Timewarp MCMC (Algorithm D. and Timewarp exploration (Algorithm (
where . is the chain length, Mar is the effective number of samples, tsampling is the sampling wall- clock time, and Pr is the autocorrelation for the lag time r [26]. The speed-up factor is defined as the ESS/s achieved by Timewarp divided by the ESS/s achieved by MD. Additional experiments and results can be found in Appendix B We train three flow models on three datasets: (i) AD, consisting of simulations of alanine dipeptide, (ii) 2AA, with peptides with amino acids, and (iii) 4AA, with peptides with 4 amino acids. All datasets are created with MD simulations performed with the same parameters (see Appendix E.. For 2AA and 4AA, we train on a randomly selected trainset of short trajectories (50ns- = 10 steps), and evaluate unseen test peptides. The relative frequencies of the amino acids in 2AA and 4AA are similar across the splits. For 4AA, the training set consists of about 1% of the total number of possible tetrapeptides (20*), making the generalisation task significantly more difficult than for 2AA. For more details see Table 2 in Appendix
Alanine dipeptide (AD) We first investigate alanine dipeptide, a small (22 atoms) single peptide molecule. We train Timewarp on AD as described in Section 5 and sample new states using AlgorithmmI for chain length of 10 million, accepting roughly 2% of the proposals. In Figure 3a visualise the samples using a Ramachandran plor [33), which shows the distribution of the backbone dihedral angles ip and v.. Each mode in the plot represents metastable state. We see that the Timewarp samples closely match MD, visiting all the metastable states with the correct relative weights. In Figure D plot the free energy (i.e., the relative log probability) of the oo and i angles, again showing close agreement. The roughness in the plot is due to some regions of state space having very few samples. In Figure 3 we show, for an initial state xx)), the conditional distribution of MD obtained by integrating Equation (2). (0(x(1 + 7) [=<())) compared with the model pa (z(t + 7) (z)). demonstrating close agreement. Finally, Figure 31 shows the time-evolution of the y angle for MD and Timewarp. Timewarp exhibits significantly more transitions between the metastable states than MD. As a result, the autocorrelation along the i angle decays much faster in terms wall-clock time, resulting in a se 7x speed-up in terms of ESS/s compared to MD (see Appendix B.4).
Dipeptides (2AA) Next, we demonstrate transferability on dipeptides in 2AA. After training on the train dipeptides, we deploy Timewarp with Algorithm on the test dipeptides for a chain length of 20 million. Timewarp achieves acceptance probabilities between 0.03% and 2% and explores all metastable states (Appendix B.I). The results are shown for the dipeptides QW and HT in Figure showing close agreement between Timewarp and long MD chains (1 μss = 10 steps). For these dipeptides, Timewarp achieves ESS/s speed-up factors over MD of 5 and 33 respectively (Appendix [..4). In Figure Left, we show the speed-up factors for Timewarp verses MD for each of the 100 test dipeptides. Timewarp provides a median speed-up factor of about five across these peptides. addition, we generate samples with the Timewarp model without the MH correction as detailed in Section 3.5 We sample 100 parallel chains for only 10000 steps starting from the same
8

---

2AA MCMC - Speed-up TIC D
ZAA exploretion Speed-up TIC 0
c)
4AA MCMC Speed-up TIC
d) 4AA expleration Speediup TIC
Speddus fector
101.
10r
10
1014
10
16h
10
10
10
Dipepiides
Opectides
kirapepiides
Figure Speed-up factors in terms of ESS/s ratios for the slowest TICA component for the Timewarp MCMC and exploration algorithms, compared to MD. The dashed red line shows a speed-up factor of onc. Gray areas depict peptides where Timewarp fails to explore all meta-stable states within 20 million steps, but MD does. Green areas depict peptides where MD fails to find all metastable states, but Timewarp does. (a), (c) Speed-up for the Timewarp MCMC algorithm (Algorithm D on test dipeptides (2AA) and tetrapeptides (4AA), respectively. (b), (d) Speed-up for the Timewarp exploration algorithm (Algorithm on test dipeptides (2AA) and tetrapeptides (4AA), respectively.
MD
Exploration
MCMC
Free onar margglll Free enargyk, Free erengyit, B)
Free energy projections
Energy distribution
10
PON:
Expioration
SAR
nic
MCMC
10
CTSA
TIC
10
LPEE
TIC
NC 110
TIC 0
TCC
-800
600
400
200
Energy kj/mol
Figure 5: Experiments 4AA test tetrapeptides SAEL, CTSA and LPEM (top, middle and bottom rows respectively). Samples were generated via MD, Timewarp exploration (Algorithm ( and Timewarp MCMC (Algorithm[]) I (a) TICA plots of samples. (b) Free energies along the first two TICA components. (c) Potential energy distribution.
initial state for each test peptide. For each peptide we select only one of these chains that finds all meta-stable states for evaluations. As before, we compute the ESS/s to compare with MD, showing median speedup factor of z 600 (Figure[4p) Note that the actual speedup when using all the chains sampled in parallel will be much larger. Timewarp exploration leads to free energy estimates that qualitatively similar to MD, but less accurate than Timewarp MCMC (Figure 3
Tetrapeptides (4AA) Finally, we study the more challenging 4AA dataset. After training on the trainset, sample 20 million Markov chain states for each test tetrapeptide using Algorithm [ and compare with long MD trajectories (1μs). In contrast to the simpler dipeptides, both Timewarp MCMC and the long MD trajectories miss some metastable states. However, Timewarp in exploration mode (Algorithm 2 can be used as a validation tool to quickly verify exploration of the whole state space. Figure 4 shows that metastable states unexplored by MD and Timewarp MCMC can be found by the Timewarp exploration algorithm. We carefully confirm the physical validity of these discovered states by running shorter MD trajectories in their vicinity (see Appendix B.5 to ensure that they are not simply artefacts invented by the model. As with 2AA, we again report the speedup factors for Timewarp relative MD in Figure 4p.d. Although Timewarp MCMC fails to speed up sampling for most tetrapeptides, Timewarp exploration shows a median speedup factor of ~ 50. For test tetrapeptides, MD fails to explore all metastable states, whereas Timewarp succeeds these are marked in green. For 10 tetrapeptides, Timewarp MCMC fails to find all metastable states found by MD - these are marked in grey. Figure shows that when Timewarp MCMC discovers all metastable states, its free energy estimates match those of MD very well. However, sometimes

---

misses metastable states leading to poor free energy estimates in those regions. Figure 5 shows that Timewarp MCMC also leads potential energy distribution that matches MD very closely. In contrast, Timewarp exploration discovers all metastable states (even ones that MD misses), but has less accurate free energy plots. It also has a potential energy distribution that slightly too high relative to MD and Timewarp MCMC.
Limitations / Future work
The Timewarp MCMC algorithm generates low acceptance probabilities (< 1%) for most peptides (see Appendix [.I]] However, this is not limitation in itself. In general, a larger proposal timestep yields smaller acceptance rates as the prediction problem becomes more difficult. However, due to Algorithm [
evaluate multiple samples in parallel nearly additional costs. As result, lower acceptance rate, when coupled with a larger timestep r, is often a favorable trade-off. While speed-up only roughly third of the 4AA peptides when using the MH correction, beating MD in wall-clock time on unseen peptides in the all-atom representation is a challenging task which has not been demonstrated by ML methods before. Furthermore, one could consider targeting systems using a semi-empirical force field instead of a classical one. Given that Timewarp requires considerably fewer energy evaluations than MD simulations, one can anticipate more substantial acceleration in this context.
Although MD and Timewarp MCMC fail to find some metastable states that were found by Timewarp exploration, we refrained from running MD and Timewarp MCMC longer due to the high computa- tional cost (Appendix FF.. Timewarp generates fewer samples compared traditional MD simulations within the same timeframe. Consequently, this scarcity of samples becomes even more pronounced in transition states, which makes Timewarp difficult to apply chemical reactions.
Timewarp could be integrated with other enhanced sampling methods, like parallel tempering or transition path sampling. In the case of parallel tempering, the effective integration requires the training of the Timewarp model across multiple temperatures, which then allows to sample all the replicas with Timewarp instead of MD. We could also alternate Timewarp proposals with learned updates to collective variables, like dihedral angles. These combined steps would still allow unbiased sampling from the target distribution [27].
Moreover, we only studied small peptide systems in this work. Scaling Timewarp to larger systems remains topic for future research, and there are several promising avenues to consider. One approach to explore different network architectures, potentially capturing all the symmetries inherent the system. Another option is to study coarse-grained structures instead of all-atom representations, to reduce the dimensionality of the problem.
8 Conclusion
We presented Timewarp, a transferable enhanced sampling method which uses deep networks to propose large conformational changes when simulating molecular systems. We showed that Timewarp used with MH correction unseen dipeptides, allowing faster computation of equilibrium expectation values. Although this acceleration was only possible for a minority of the letrapeptides we considered, we showed that Timewarp used withou! the MH correction explores the metastable states of both dipeptides and tetrapeptides much faster than standard MD, and we verify the metastable states discovered physically meaningful. This provides a promising method to quickly validate if MD simulations have visited all metastable states. Although further work needs to be done scale Timewarp larger, more interesting biomolecules, this work clearly demonstrates the ability of deep learning algorithms to leverage transferability accelerate the MD sampling problem.
accelerate asymptotically unbiased sampling on many
Acknowledgments
We thank Bas Veeling, Claudio Zeni, Andrew Fowler, Lixin Sun, Chris Bishop, Rianne van den Berg, Hannes Schulz, Max Welling and the entire Microsoft A14Science team for insightful discussions and computing help.
10

---

References
[1] Ignasi Buch, Toni Giorgino, and Gianni De Fabritüis. Complete reconstruction of an enzyme- inhibitor binding process by molecular dynamics simulations. Proceedings of the National Academy of Sciences, (08225):01844-0188) 2011.
[2] Jianfci Chen, Cheng Lu, Biqi Chenli, Jun Zhu, and Tian Tian. Vflow: More expressive generative flows with variational data augmentation. In International Conference Machine Learning, pages 1660-1669. PMLR, 2020.
[3] Cecilia Clementi. Coarse- gaanse-grained models of protein folding: toy models or predictive tools? Current opinion in structural biology. 18(1):10- 15,001155200 2008.
[4] Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and Boltzmann generators. Phys. Rev. Res., 4:1042005, Oct 2022.
[5] Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference Learning Representations, 2017.
[6] David J. Earl and Michael W. Deem. Parallel tempering: Theory, applications, and new perspectives. Phys. Chem. Chem. Phys., 7:3910-3916, 2005.
[7] Peter Eastman, Jason Swails, John D Chodera, Robert McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm 7: Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology, 1007)59,, 2017.
[8] Xiang Fu, Tian Xie, Nathan J. Rebello, Bradley Olsen, and Tommi S. Jaakkola. Simulate time- integrated coarse-grained molecular dynamics with multi-scale graph networks. Transactions on Machine Learning Research, 2023.
[9] C. Huang, Laurent Dinh, and Aaron Courville. Augmented normalizing flows: Bridging the gap between generative flows and latent variable models. ArXiv, abs/2002.07100 101, 2020.
[10] Brooke E Husic and Vijay Pande. Markov state models: From an art to a science. Journal of the American Chemical Society, 140(7):238--239, 2018.
[11] Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi S. Jaakkola. Torsional diffusion for molecular conformer generation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[12] William Thomson Baron Kelvin. The molecular tacties of a crystal. Clarendon Press, 1894.
[13] Leon Klein, Andreas Krämer, and Frank Noć. Equivariant flow matching. arXiv preprint Stt.:23006110330 2023.
[14] Sebastian Kmiecik, Dominik Gront, Michal Kolinski, Lukasz Wieteska, Aleksandra Elzbieta Dawid, and Andrzej Kolinski. Coarse-graing protein models and their applications. Chemical reviews, (1((1))7898-7936 2016.
[15] Jonas Köhler, Michele Invernizzi, Pim de Haan, and Frank Noé. Rigid body flows for sampling molecular crystal structures. In International Conference on Machine Learning. ICML 2023, volume 202 of Proceedings of Machine Learning Research, pages 17301-17326. PMLR, 2023.
[16] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning, pages 5361-5370. PMLR, 2020.
[17] Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and Wortman Vaughan, editors, Advances in Neural Information Processing Systems, volume 34, pages 2796-2809. Curran Associates, Inc., 202 1.
[18] Alessandro Laio and Michele Parrinello. Escaping free-energy minima. Proceedings of the National Academy of Sciences, 12522 2002.

---

[19] Paul Langevin. Sur la théorie du mouvement brownien. Compt. Rendus, 146:530-533 133300533,,1 1908.
[20] Daniel Levy, Matt D. Hoffman, and Jascha Sohl-Dickstein. Generalizing Hamiltonian Monte Carlo with neural networks. In International Conference on Learning Representations, 2018.
[21] Zengyi Li, Yubei Chen, and Friedrich T. Sommer. neural network MCMC sampler that maximizes proposal entropy. Entropy, 23(3), 2021.
[22] Kresten Lindorff-Lareen, Stefano Piana, Ron O Dror, and David E Shaw. How fast-folding proteins fold. Science, 334(6055):517-520 2011.
[23] Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. VAMPnets for deep learning of molecular kinetics. Nature communications, 91):1-111,015
[24] Nicholas Metropolis, Arianna W Rosenbluth. Marshall N Rosenbluth, Augusta H Teller, and Edward Teller. Equation of state calculations by fast computing machines. The Journal of Chemical Physics, 21 (6):1667--102) 1953.
[25] Laurence I Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf. and José Miguel Hernández-Lobate Se (3) equivariant augmented coupling flows. arXiv preprint arXiv:2308. 10364, 2023.
[26] Radford M Neal. Probabilistic inference using Markov chain Monte Carlo methods. Department of Computer Science, University of Toronto Toronto, ON, Canada, 1993.
[27] Jerome Nilmeier, Gavin E Crooks, David DL Minh, and John Chodera. Nonequilibrium candidate monte carlo is an efficient tool for equilibrium simulation. Proceedings of the National Academy of Sciences, 108(45):E1009 E1018, 2011.
[28] Frank Noé, Simon Olsson. Jonas Köhler, and Hao Wu. Boltzmann generators sampling equilibrium states of many-body systems with deep learning. Science, 365: 147, 2019.
[29] Frank Noé, Christof Schüttte, Eric Vanden- Eijnden, Lothar Reich, and Thomas R Weikl. Con- structing the equilibrium ensemble of folding pathways from short off- equilibruum simulations Proceedings of the National Academy of Sciences, 106(45): 1901 1901,,, 2009.
[30] Guillermo Hrezz-Hernández, Fabian Paul, Toni Giorgino, Gianni De Fabritiis, and Frank Noé. Identification of slow molecular order parameters for Markov model construction. The Journal of chemical physics, 1,776400700411,
[31] Nuria Plattner, Stefan Doerr, Gianni De Fabritiis, and Frank Noé. Complete association kinetics in atomic detail revealed by molecular dynamics simulations and Markov modelling. Nature chemistry, 9(10):0055-111 2017.
[32] Jan-Hendrik Prinz, Hao Wu, Marco Sarich, Bettina Keller, Martin Senne, Martin Held, John D Chodera, Christof Schütte, and Frank Noé. Markov models of molecular kinetics: Generation and validation. The Journal of chemical physics, 134(17):174105 2011.
[33] G N Ramachandran, C. RRmakrishhan, and V Sasisekharan. Stereochemistry of polypeptide chain configurations. Journal of Molecular Biology, pages 95-99, 1963.
[34] Jeff Rasley, Samyam Rajbhandari, Olatunji Ruwase, and Yuxiong He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD '20, page 3505-3506, New York, NY, USA, 2020. Association for Computing Machinery.
[35] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference muchine learning, pages 1530-1538. PMLR, 2015.
[36] Danilo Jimenez Rezende, Sébastien Racanière, Irina Higgins, and Peter Toth. Equivariant Hamiltonian flows. arXiv preprint 1909.1990 2019.
[37] Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple dime-resolution surrogates for molecular dynamics. arXiv preprint arXiv:2305.18046, 2023,
12

---

[38] Jiaming Song, Shengjia Zhao, and Stefano Ermon. A-NICE-MC: Adversarial training for
MCMC. Advances in Neural Information Processing Systems, 30, 2017.
[39] Robert Swendsen and Slan-Sheng Wang. Replica monte carlo simulation of spin-glasses. Physical review letters, 57(21):2607, 1986.
[40] William C Swope, Jed W Pitera, and Frank Suits. Describing protein folding kinetics by molecular dynamics simulations. 1. theory. The Journal of Physical Chemistry B, 108((1))6571- 6581, 2004.
[41] Michalis Titsias and Petros Dellaportas. Gradient-based adaptive Markov chain Monte Carlo. Advances in Neural Information Processing Systems, 32, 2019.
[42] Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis Louis-Philiphe Philippe Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Process- ing (EMNLP-IJCNLP), pages 4344-4353, Hong Kong, China, November 2019, Association for Computational Linguistics.
[43] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
[44] Christina Winkler, Daniel Worrall, Emiel Hoogeboom, and Max Welling. Learning likelihoods with conditional normalizing flows. arXiv preprint arXiv:1912.00042. 2019.
[45] Hao Wu, Andreas Mardt, Luca Pasquali, and Frank Noé. Deep generative Markov state models. Advances in Neural Information Processing Systems, 31, 2018
[46] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang, Geodiff: A geometric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
13

---

A Symmetries of the architecture
A.1 Proof of Proposition 4.1
In this appendix we provide more details on the equivariance of the Timewarp architecture. We first prove Proposition 4.1 from the main body:
Proof. Let ((t m) denote the random variable obtained by sampling p(2) and computing ((++T) == f(Z xt)). Here we subseript ( T)n(s) by xxt) emphasize that this the random variable obtained when conditioning the flow s=tt). We first note that the equivariance condition on the densities p(ar(( + ())))(()) = p(x(t + (z))) is equivalent to the following constraint on
the random variables:
X(t + TT A o x +t+
(19)
where á denotes equality in distribution. To see this, let px denote the density of the random variable X. Then, terms of densities, Equation 19 is equivalent to stating that, for all x,
Px (X(++))
(x) (ox (I)
(20)
-x((t+))c)) (o 1x))
(21)
where in Equation 21 we used the change-of-variable formula, along with the fact that the group actions we consider (rotations, translations, permutations) have unit absolute Jacobian determinant.
Redefining z + J.., we get that for all x,
(x((+))az(o) (ax) Px((+=))=co) (2)).
(22)
Recalling the notation that ( ()))))))) is interpreted the random variable obtained by conditioning the flow on (t), this can be written
p(ox|ax(() (()))=((()
(23)
which is exactly the equivariance condition stated in terms of densities above. Having rephrased the
equivariance condition in terms of random variables in Equation (19). the proof of Proposition 4.1 is
straightforward.
[ + T)oz(1) = ((Z,
(24)
- ((oZZor((t)
(25)
=f((Zx(tt)
(26)
== =t
(27)
where in Equation
25) we used the fact that the base distribution p(=) a-invariant.
A.2 Translation equivariance via canonicalisation
We now describe the canonicalisation technique used to make our models translation equivariant. Let q(>²(1 + r), xx (t (()) be an arbitrary conditional density model, which is not necessarily translation equivariant. We can make translation equivariant in the following way. Let x denote the average position of the atoms,
N
(28)
Then define
pl + T), x" (= (t+)))))))) -P)))
(29)
where the subtraction (P()) is broadcasted over all atoms. We now consider the effect of translating both (t) and (t by the same amount. Let be translation vector in R3, Then
p( (1 + T) + a.x" (1+) aa)
(30)
-
(x*(t) +a)
(31)
(32)
= -(0/0++11-1000004+))))0000---
(33)
-(
(34)
14

---

Acceptance probabilities
Acceptance rate in %
2
1
2AA
4AA
Figure 6: Acceptance probabilities for samples on unseen test peptides with the Timewarp MCMC algortihm. The red line is 0.01%, below that efficient sampling becomes difficult.
Hence p is translation equivariant even if q is not.
A.3 Chirality
In addition to the symmetries described in Section 4.1 the potential energy U(*²) of molecular configuration is also invariant under mirroring. However, in the presence of chirality centers, mirrored configuration is non-superposable to its original image [12]. An example of a chirality center an amino acid Carbon atom connected to four different groups, e.g. a C. atom. In nature most amino acids come in one form, namely L-amino acids. Hence, all our datasets consist of peptides containing only L-amino acids. In rare cases, as the model proposes large steps, one step might change one L-amino acid of peptide to a D-amino acid in way that the resulting configuration has a low energy and the step would be accepted. We prevent this by checking all chirality centers for changes at each step and reject samples where such change occurs. This does not add any significant computational overhead.
B Additional results
In this section show additional results like the conditional distribution well as peptide examples for experiments discussed in Section 6
B.1 2AA additional results
More examples from the 2AA test set are presented in Figures and We achieve the worst speed-up for the dipeptide GP Figure) 7 last row) as it does not show any slow transitions. The distribution of the acceptance probabilities for the Timewarp MCMC algorithm is shown in Figure
B.2
4AA additional results
More examples from the 4AA test set are presented Figures and The distribution of the acceptance probabilities for the Timewarp MCMC algorithm is shown in Figure
B.3
Conditional distributions
The model was trained to generate samples from the conditional Boltzmann distribution μ((((t T [t))))) Here we show some examples of the conditional distribution generated by the Timewarp model compared
the conditional distribution induced by MD. While we can generate 5, 000 samples from the conditional distribution of the model in parallel, we require 5, 000 distinct MD trajectories of simulation length sample the conditional distribution with MD. Hence, generating samples from the conditional distribution is several orders of magnitude faster with the model. In Figures 10 to 12 we show example conditional distributions for alanine dipeptide and peptides from the 2AA and 4AA datasets. For all peptides the model learns conditional distribution that is close to the conditional MD distribution. Moreover, the relative weights in the TICA projections as well as the bondlength distributions match very well. Only the energies of the model samples are higher,
15

---

Timewarp
energy projections
autocorreletions
O time depedence
Free eneeryiittyy energyayy
Autecmmationn butscoresstt
Brewwrr Lise cacs
B455 18: Aid
031
Free energylog
30
0SS1 e 5.48%
JL
0344
T10
TC1
200
400
1003
Log Bimre
Walle woc.clocc-kooo mapiing time in
Figure 7: Experiments 2AA test dipeptides QW (top row), HT (middle row) and GP (bottom row). Comparison of the long MD trajectory and Timewarp MCMC (Algorithm D (a) TICA plots. (b) Free energy comparison of the first two TICA components. (c) Autocorrelation for the TIC component. (d) Time dependence of the TIC 0 component.
Timewarp
energy projections
TIC autocorrelations
TIC time depedence
Free energytat Free ereawwerr free exemptor Free everppee
-0
W W
0.5
mco TKC
0.
1.0
Autromeetss
Thnuero ESSA 3 3385
0.5
Aurocindtttonn Autccoeeetto
nco
0.5
TCB
TC0
TIC
200
600
800 1000
Log time in
Whi--lock sampling time
Figure 8: Experiments for the 2AA test dipeptides DH (first row), GT (second row), TK (third row), and CW (last row), Comparison the long MD trajectory and Timewarp MCMC (Algorithm M (a) TICA plots. (b) Free energy comparison of the first two TICA components. (c) Autocorrelation for the TIC 0 component. (d) Time dependence of the TIC 0 component.
16

---

MD
Exploration
MCMC
b)
Free energy. projections
Energy distribution
Froe energyilka
To
MONE
TIC
5
Free energy/la
LYVI
TIC
5
Fres energyital
100
MD
Exploration
CSR TIC
MCMC
D
1 TIC
1 TK1
600
400
-200
Energy in kjimal
Figure 9: Experiments on 4AA test tetrapeptides AWCK, LYVI and CSFQ (top. middle and bottom rows respectively). Samples were generated via MD, Timewarp exploration (Algorithm 2). and Timewarp MCMC (Algorithm D (a) TICA plots of samples. (b) Free energies along the first two TICA components. (c) Potential energy distribution. For AWCK all metastable states are found by all methods, for LYVI the MD trajectory misses one state, and for CSFQ Timewarp MCMC misses the slowest transition. In all cases Timewarp exploration discovers all metastable states.
Boltamann
MD
Timeware
Ramachandran projections 6)
Energy distribution
Phi sutocorreistions
0
25
0.
23
-25
1.0
25
300
54
Mamml
150
150
200
202
400
Figure 10: Comparing the conditional Boltzmann distributions generated with MD trajectories and the Timewarp model for alanine dipeptide. (a) Ramachandran plots for the conditional distributions compared with the equilibrium Boltzmann distribution. The red cross indicates the conditioning state. This is similar to the plot shown in Figure 3b. but here showing different conditioning state. The match between the conditional distributions is not as close here as it is for Figure 3R. which could be because here the conditioning state is chosen be in the less likely metastable state. (b) Projections on the first two dihchdral angles for the conditional distributions. (c) Potential energies of the conditional distributions. (d) Autocorrelations for samples generated according to the MCMC algorithm (Algorithm) I compared with long MD trajectory. Note that this autocorrelation plot is not for the conditional distribution, but corresponds to the results shown in Figure 3
emphasising the importance of the Metropolis-H Hastings correction to obtain unbiased samples from the Boltzmann distribution with the Timewarp model.
B.4 Autocorrelations
In Section 6 compute the speedup of the Timewarp model by comparing the effective sample sizes per second (Equation (8)) for the slowest transition with MD. As the ESS depends on the autocorrelation, it is also insightful to look the autocorrelation decay in terms of wall-clock time.
We show some example autocorrelations for the investigated peptides in Figures 7 and 10 Note that the under the autocorreeation curve is inversely proportional the ESS.
B.5 Exploration of new metastable states
For some tetrapetides in the test set even long MD trajectories (1,s) miss some metastable states, e.g. for LYVI and CSTA shown Figure 12h. However, we can easily explore these with the Timewarp exploration algorithm (Algorithm) 2 To confirm that these additional metastable states are in fact
17

---

Boltzmann
MD
Timewarp
TICA projections
Energy distribution
Rondlength distribuzion
AMA
Cost Timewars
McA
Wan
w
Won
Dooal
BAS
EASS
163
1000
0 1900
TCC
200
500
400
-20 5.00
0.10
015
6:n
Figure 11: Comparing the conditional distribution of the Timewarp model po(x( (z((( with the conditional distribution generated with MD μ(([(t + [=)())) The correspond to the peptides TK, AAEW, EASS, respectively, where we show for each peptide two different conditioning states. (a) TICA plots. First column: samples from the Boltzmann distribution μu()) generated with MD. Second column: samples from uuz(t + r) (t()) generated with MD. The conditioning state xt) is indicated with the red cross. Third column: samples from pe (=((tt + (t)))) generated with the Timewarp conditional flow, without MH correction. The conditioning state ztt) is indicated with the red cross. (b) Projection of the conditional distributions from (a) onto the first two TICA components. (c) Potential energy distributions of the conditional distributions. (d) Conditional bondlength distribution, which for these values of will be close to the equilibrium distribution. Each mode in the graph represents a different bond type, e.g., C-H.
stable, we run several shorter MD trajectories (0.5 10%fs), in the same way as in Appendix B.3 starting a nearby metastable state already discovered with the long MD trajectory (Figure 12b). Once one of them hits the new, previously undiscovered state, we start new short MD trajectories (0.5 x 10 from there as well (Figure 12b). These new short MD trajectories either sample within this previously undiscovered state, transition to the other metastable states. This shows that this metastable state discovered by Timewarp exploration is indeed valid, and was simply undiscovered during the long MD trajectory. In addition, we compare in Figure 12 the conditional MD distributions with that of Timewarp, again showing close agreement.
C Batched sampling algorithm for Timewarp with MH corrections
Pscudocode for the algorithm described Scction 3.4 is given in Algorithm I
D Exploration of metastable states using Timewarp without MH corrections
describe the exploration algorithm for Timewarp, which accepts all proposed states unless energy is above a certain cutoff value. As there is no MH correction, the generated samples will not asymptotically follow the Boltzmann distribution, but the exploration of the state space is much faster than with MD Algorithm The pscudocode is shown in Algorithm
18

---

a)
MD
Exploration
MCMC
MD
Timewarp
Timewarp
DWV Tic1
TCc
TCC
Figure 12: Validation of new metastable states found with Timewarp exploration. The red crosses indicate the different conditioning states. (a) TICA plots for (not conditional) samples generated via a long MD trajectory, Timewarp exploration (Algorithm{), and Timewarp MCMC (Algorithm) []. Timewarp exploration discovers some metastable states unseen in the long MD trajectories (bottom of LYVI plot, bottom left of CTSA plot). (b) Conditional distributions generated with MD or with the Timewarp model, starting from states visited by the long MD trajectory shown in (a). Some short MD trajectories now discover the new metastable states, verifying that they are indeed valid states. Only the final state of each short MD trajectory recorded. (c) Conditional distributions generated with MD or with the Timewarp model, starting from the new metastable states discovered by one of the short MD trajectories shown in (b).
Algorithm Timewarp Morccrre MCMC with batched proposals
Require: Initial state Xo (Xi; X5)), chain length M, proposal batch size B.
m <- 0
while m <M do
Sample X1.... .x po IXP,) [Batch sample)
for 1..... do
~~ N(O.11 I) (Resample auxiliary variables] X, 4- (Xp..e
Sample I, ~ (arnoulliaa(x,.. Xx))
end for
i = {b: Ib 1,, B} + then a min(S) [First accepted sample
m+1 (X.,, X f . XP.
X m+a f- XP
m t- m +
else
(XP
m+ XP x f- ..,))
B
end if
end while
output U: XP....X XP
19

---

Algorithm 2 Fast, biased exploration of the state space with Timewarp
Require: Initial state X number of steps M, maximum allowed energy increase AU max for m 0, O..... M do
Sample X ~ pe(- X (Sample from conditional flow]
if (X m (X < mum then
X m+1 t X
else
X < XP, [Reject if energy change is too high}
end if
end for output xp....X" XM
Note that unlike Algorithm there is need for the auxiliary variables, since the conditional flow only depends on the positions, and no MH acceptance ratio is computed here. The potential energy U includes here also large penalty if the ordering of chirality center changes as described in Appendix A.3 As sampling proposals from Timewarp can be batched, we can generate / chains in parallel, all starting from the same initial state. This batched sampling procedure leads to even further speedups. For all exploration experiments we use batch size of 100, and run M = 10000 exploration steps. The maximum allowed energy change cutoff at AU 300kJ/mol.
E Dataset details
We evaluate model three different datasets, AD, 2AA, and 4AA, introduced in Section 6 All datasets are simulated in implicit solvent using the оpenMM library 7 use the parameters shown in Table
For all MD simulations
Table 1: OpenMM MD simulation parameters
Force Field
amber-14
Time step
0.5fs
Friction coefficient
0.3
Temperature
310K
Integrator
LangevinMiiddeeneerrraao
We present more dataset details, like simulation times and number of peptides, in Table
F Hyperparameters
Depending the dataset, different Timewarp model sizes used, as shown in Table For all datasets the Multihead kernel self-attention layer consists of heads with lengthscales & (0.1, 0.5, 0.7, 1.0, given in nanometers.
Table 2: Dataset details
Dataset name
AD
2AA
4AA
Training set simulation time
100 ns
50 ns
50 ns
Test simulation time
100 ns
1 Jus
1 JuS
MD integration step At Timewarp prediction time
0.5 fs
0.5
0.5 fs
0.5 10
0.5 10
0.5 10
No. of training peptides
1
200
1400
No. of training pairs per peptide
2 x 10
1 104 100
x 1
No. of test peptides
1
30
20

---

Table 3: Timewarp model hyperparameters
Dataset
RealNVP layers
Transformer layers
Parameters
Atom-cmbedding dim H
AD
12
1 x 10$
64
2AA
12
1 x 10a
64
4AA
16
x 10
128
The Pm and Oou MLPs use SiLUs activation functions, while the Transformer MLPs use RcLUs. Note the transformer MLP refers to the atom-wise MLP shown in Figure 2 Middle inside the transformer block. The shapes of these MLPs vary for the different datasets as shown in Table g
Table 4: Timewarp MLP layer sizes
Dataset
iin MLP
Poot MLP
Transformer MLP
AD
70, 256, 128] [70, 256, 128]
128, 256, 3] 128, 256, 3] [128, 2048, 3]
[128, 256, 128] [128, 256, 128] [128, 2048, 128]
2AA
4AA
[134, 2048, 128]
The first linear layers in the kernel self-attention module always has the shape [128, 768] (in Section denoted V), and the second (after concatenating the output of head head) has the shape [768, 128]. The transformer feature dimension
for all datasets 128.
After likelihood training, we fine-tune the model for the AD and 2AA dataset with combination of all three losses discussed in Section 5 We did not perform fine tuning for the model trained on the 4AA dataset. We use a weighted sum of the losses with weights detailed in Table We use the
Table 5: Timewarp loss weighting factors
Dataset
(ik
Lave
Lent (0)
AD
0.99
0.01
0.1
2AA
0.9
0.1
0.1
4AA
1
0
0
FusedLammb optimizer and the DeepSpeed library (34 for all experiments. The batch size as well as the training times are reported in Table All simulations started with learning rate of 5 10-4
Table 6: Timewarp training parameters
Dataset training method
Batch size
No. of A-100s
Training time
AD
likelihood
256
1
1 week
AD
acceptance
64
1
days
2AA
likelihood
256
4
weeks
2AA
acceptance
256
4
days
4AA
likelihood
256
4
weeks
the learning rate is then consecutively decreased by a factor of 2 upon hitting training loss plateaus.
G Computing infrastructure
The training was performed NVIDIA A-100 GPUs for the 2AA and 4AA datasets and single NVIDIA A-100 GPU for the AD dataset. Inference with the model as well as all MD simulations were conducted on single NVIDIA V-100 GPUs for AD and 2AA, and on single NVIDIA A-100 GPUs for 4AA.
21


Reviews:
[{"review_id": "zfdOQqYXXY", "reviewer": "Reviewer_Pymp", "summary": "In this work, the authors present a method for speeding up MD simulations. The speeding up of MD simulations is needed to be able to simulate processes that occur over longer timespans (eg. protein folding). The method proposed by the authors, Timewarp, does so by learning the transition probability of MD simulations over timesteps beyond those normally used. The transition probability is learned using a normalising flow based on a Transformer architecture. Transition steps sampled using the normalising flow are pushed through a Metropolis-Hasting acceptance step to assure asymptotic unbiased samples. Additionally, the authors propose to run the normalising flow without the MH step for faster sampling. The normalising flow is learned using trajectories of multiple structures using traditional MD simulation. By using multiple structures at the same time, the authors aim to make the method generalize to unseen structures. This is one of their major contributions.\n\nThe work experimentally shows the success of the method on 3 different sizes of molecules; 1) Alanine Dipeptide, 2) 2 amino acid peptides and 3) 4 amino acid peptides. The experiments show that\n\n-Timewarp closely matches the samples from MD\n-Timewarp exploration (without MH step) provides estimates of the free energy close to MD, but less accurate then MCMC (with MH step).\n\n", "questions": "As I commented above, from my understanding, the problem discussed and addressed in this paper lies somewhere in between force field learning and conformer generation. The transition step length tau in this case determines which of the two the problem is closer to. If tau is very small, the problem at hand is similar to force field learning. With tau very large, the conditional distribution becomes similar to learning the Boltzmann distribution. Considering this, I'm surprised the work does not contain a study of the influence of tau on the speedup seen by the method over MD. Could the authors comment on this? For example, why was the specific value for tau used in the paper chosen?", "limitations": "The authors have sufficiently addressed the limitations. ", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "NOTE: This is an updated review from a review round of an earlier conference. The review is updated to reflect the changes made by. \n\n**Originality:** The paper discusses a very interesting new intermediate problem in ML for Chem/Bio that lies in between the two more often studied problems of Force Field learning and Conformer generation. I've commented on this below. \n\n**Quality:** The paper provides a good introduction to the problem at hand and as such sufficiently motivates the work. The proposed solution is interesting and makes good use of existing methods for similar problems without overcomplicating it. Experiments are in depth and well executed. \n\n**Clarity:** The paper is well written and structured. Figures are of high quality. \n\n**Significance:** The paper has the potential to be very significant. The problem introduced is novel and important. ", "weaknesses": "- Unfortunately, there's no code available and the datasets are not open sourced. \n- Related work could be improved with more references to work from the computational chemistry literature on enhanced sampling(ie. topic such as Transition Path Sampling and Bias Potential methods).\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review, and appreciate their assessment that the paper has the potential to be very significant. We now address their comments individually.\n\n> Unfortunately, there's no code available and the datasets are not open sourced.\n\nWe will release the code and datasets publicly upon publication of the paper. Unfortunately we cannot release our code and datasets anonymously before publication.\n\n> Related work could be improved with more references to work from the computational chemistry literature on enhanced sampling'\n\nThank you for pointing us to the topics of transition path sampling and bias potential methods. We will include a discussion of this in the related work section in the final version. We believe these could both in practice be used in conjunction with a method like Timewarp in order to further accelerate sampling. One possibility is to alternate Timewarp proposals with some (learned) change of collected variables, like the dihedral angles of peptides. These combined steps would still allow unbiased sampling from the target distribution. Timewarp could also be trained at multiple temperatures, by either training one model at each temperature or having the temperature as an additional conditioning input for a single model, which then could be applied similar to regular parallel tempering. We will include these possible directions for future research in the Conclusion.\n\n> I'm surprised the work does not contain a study of the influence of tau on the speedup seen by the method over MD. Could the authors comment on this? For example, why was the specific value for tau used in the paper chosen?'\n\nThe impact of tau on the speedup seen by the method over MD is indeed an important research question to address. We found that the choice of tau involved a trade-off: if tau was too large, then the prediction problem becomes more difficult. This can lead to low acceptance rates, and hence a smaller speedup. However, if tau is too small, then the model must be run many times in order for the Markov chain to converge, again leading to a smaller speedup. As we introduced the batch acceptance algorithm (Algorithm 1), we can evaluate multiple samples in parallel at nearly no additional costs. Hence, a lower acceptance rate in combination with a larger time step is usually preferable, and an acceptance rate of nearly one would waste the potential for larger time steps. The current values were chosen by manual hyperparameter search guided by these intuitions. Unfortunately, a systematic hyperparameter search by trying a grid of tau values was beyond the computational resources of our project. We will add a further discussion of this to the final version of the paper."}, {"review_id": "2NkDDRPuNj", "reviewer": "Reviewer_heyF", "summary": "\n\nThis paper presents a Boltzmann generator for small molecular systems in implicit solvent.  The generator works demonstrates transferability to test systems of the same class (dipeptides or tetrapeptides).  It is the first Boltzmann generator that operates in Cartesian coordinates and demonstrates transferability.  The authors estimate that for these very small molecular systems in their study, the sampling via a Boltzmann generator could provide wall-time acceleration compared to vanilla MD simulation.  \n", "questions": "\nThe work is limited to molecules in implicit solvent.  How difficult would it be scale these methods to also generate the full environment of water, ions, lipids and small molecules in typical molecular dynamics simulations in the context of drug discovery?  Perhaps in a more practical fashion, are there any limitations in the current architecture \n\nDid the authors try to train the model on explicit solvent simulations and use only the non-water/non-ion coordinates in their systems?  The generation of the training trajectories does not seem particularly time consuming compared to the rest of the compute effort that went into the project and the possible wall-time improvements might be even larger (though of course this depends on the software and hardware that is employed).  The very limited number of total atoms in the current models would suggest that existing enhanced sampling methods for MD simulations (including those that include kinetic considerations) would also work for these small systems; however, such sampling techniques might be trickier to optimize in explicit water than in implicit water.  (Even without any sophisticated techniques, I would guess that the implicit-solvent dipeptides and tetrapeptides would reduce their autocorrelation much more quickly by sampling them at ridiculously high temperatures (or fluctuating temperatures as in simulated tempering), and could then be projected back to the baseline with an MH criterion.)\n\nWhy did the authors not mention the Torsional Diffusion paper in their introduction?  Wasn't that work also a generalizable Boltzmann generator (using internal coordinates), and probably the first of this kind?\n", "limitations": "No potential negative societal impact", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This paper is the first demonstration of a Boltzmann generator in Cartesian coordinates that generalizes to untrained systems.  This is certainly commendable, despite the existence of previous generators for small molecules that relied on internal coordinates.  The method is presented in a clear fashion and the authors plan to release their codes and data upon acceptance.  The experiments are sensible and show that the slow modes of the MD simulations are covered with a similar distribution by the sampler.  This paper is adding to a series of developments in the line of Boltzmann generators, an area that shows great promise, even though it is still at its infancy with respect to practical applications in drug discovery. As such, I think that it is a significant contribution to this conference.\n", "weaknesses": "\nThe correct operation of this generator requires a Metropolis Hastings step, which will limit the application of the current technology to only very small systems.  This is not necessarily a problem with all Boltzmann generators, but this limitation is very clearly demonstrated in the work and in leads to a reduction in the acceptance rate as the number of particles doubled.  The authors don't discuss this possible limitation of the current work in their paper, nor suggest ways to think about scaling to systems of practical interest.\n\nThe work is only dealing with molecules in implicit solvent.  These implicit solvent simulations are not of practical use in drug discovery, so it would be useful to have a discussion of possible generalizations towards explicit simulations (either generating the full environment, or at least learning from such simulations and generating whatever is feasible).\n\nIt is hard to judge reproducibility without the availability of the code/data; I hope this will be alleviated once the paper is published.  \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review and questions. We appreciate their judgement that our paper is a significant contribution to this conference. We now address their questions individually.\n\n> The correct operation of this generator requires a Metropolis Hastings step, which will limit the application of the current technology to only very small systems... The authors don't discuss this possible limitation of the current work in their paper, nor suggest ways to think about scaling to systems of practical interest.\n\nWe agree that the low speed-up for larger systems is a limitation of our work. We note that our limitations section clarifies that we only succeed in speeding up sampling for a third of the 4AA peptides. \nWe leave the problem of scaling to larger systems to future work. A number of routes are possible. We list a few here:\n- We could remove the Metropolis-Hastings correction and focus solely on the exploration algorithm. This would allow us to scale beyond 4AA at the expense of potentially biased simulations. However, we believe this could still be extremely useful as long as methods were developed to validate the correctness of the predictions.\n- Removing the MH correction would also allow for coarse-graining, e.g., representing only all the heavy atoms, or only all the backbone atoms, instead of every single atom in the system.\n- Parameterisations such as those used in AlphaFold2 could also improve performance. There the backbone heavy atoms are represented as rigid frames. Hence the model would not need to learn these kinematic relationships, but these would be hard-coded into the representation.\n\nWe will include a more detailed discussion of this in the final version.\n\n> The work is only dealing with molecules in implicit solvent...\n\nThis is indeed a less realistic scenario than explicit solvent simulations. We focused on implicit solvent simulations in this paper in order to begin with the simplest scenario, and leave the question of explicit solvent to future work. In order to adapt Timewarp for explicit solvent simulation, the most straightforward extension would be to simply add the solvent atoms as additional nodes to be processed by the transformer. This would increase the computational cost of the forward pass. To mitigate this, we could investigate sparse transformer architectures/graph neural networks.\n\n> It is hard to judge reproducibility without the availability of the code/data\n\nWe will release code and data publicly upon publication.\n\n> Why did the authors not mention the Torsional Diffusion paper in their introduction?\n\nWe will include the Torsional Diffusion paper in the related work section. As mentioned by the reviewer, the main differences to our work is that they rely on internal coordinates and do not operate in the all atom system.\n"}, {"review_id": "YYMS2jojpI", "reviewer": "Reviewer_fomi", "summary": "In this paper the authors introduce Timewarp, a method that relies on a normalizing flow as a proposal distribution for enhanced sampling of the Boltzmann distribution. The key results include an acceleration in the MD time step from ~1 fs to 10^5 fs and transferability between peptide systems.  ", "questions": "Can the authors expand on the choice to treat velocity as a non-physical auxiliary variable? It is clear that it makes a permutation equivariant coupling layer easier to construct, but is there no information loss that might contribute to the low MH acceptance rate?\nHow many heads and how were length scales determined for each head in the kernel self-attention? Does ablating this tell you anything about the importance of including long-range interactions, or is the locality assumption implicit in the kernel self-attention sufficient for capturing the dynamics for the peptides considered?\nCan the authors provide additional references and discussion (perhaps in the appendix) of the evaluation metric presented in Eq 18? What ESS/s values are reported when comparing different classical MD algorithms? \nSince Timewarp is evaluated on small peptide systems, conformer generation with a semi-empirical method like xTB might also be an appropriate evaluation. Can the authors discuss their method in the context of such approaches?\nAre there any concrete proposals for future work that would address the current limitations in MH acceptance rate, system size, and the troubling trends in scaling to larger systems?\nWhy do the authors use amber-14 in openMM rather than a more modern force field? \nThe caption in Fig 4 does not match the figure labels.\nIn Fig 3c the red dot indicating initial state is difficult to see. \n", "limitations": "Yes, although further discussion of possible avenues for addressing these limitations would be welcome.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The core demonstration of generalization to unseen peptides in an all-atom representation is an important contribution.\n\nThe method is clearly presented and the value of accelerated sampling is also clear.\n", "weaknesses": "The method seems to break down already for 4-mer peptide systems. The limitations in expressivity of normalizing flows shown throughout the literature cast doubt on the scalability of the method to practical applications.\n\nPerformance benchmarks with respect to classical MD approaches are difficult to evaluate, due to choices in force field, integrator, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review, questions, and suggestions. We appreciate their judgement that our paper is an important contribution. We now address their questions individually.\nTo keep the rebuttal within the character limit, we often only cite parts of each question.\n\n> The method seems to break down already for 4-mer peptide systems...\n\nWe agree on the limitation. However, generative flow models for molecular systems remain an active research field. New architectures might increase the expressiveness allowing scaling to larger systems. \nIn our study, we primarily focused on addressing the transferability challenge, rather than scalability, as prior works in generative modeling for molecules in Cartesian coordinates had not demonstrated transferability before our research.\nMoreover, we  present also the Timewarp exploration algorithm, which does not include the Metropolis acceptance step, and shows significant speed-up over classical MD at the cost of not being asymptotically unbiased. Another common way to scale generative models to larger systems is coarse graining, which could be also used to scale Timewarp to larger systems. \n\n> Performance benchmarks with respect to classical MD approaches are difficult to evaluate, due to choices in force field, integrator, etc.\n\nWe conducted our MD simulations using the widely-used openMM library, employing the popular Amber-14 force field, which is a common choice for unconstrained MD simulations. The simulation parameters we utilized can be found in appendix E and are standard for such simulations.\nWe are confident that our results would not be significantly impacted if alternative MD libraries or classical force fields were employed. However, semi-empirical or other computationally expensive force fields, might yield better results compared to MD. See below for a more detailed discussion.\n\n> Can the authors expand on the choice to treat velocity as a non-physical auxiliary variable? ...\n\nWe initial tried using the physical velocities, but found that replacing them with \nindependent auxiliary variables that target a standard Normal worked better. This is allowed because we are only interested in simulating the position distribution, and the density in Eqn 4 has the same marginal for the positions as the non-augmented density in Eqn 1. The augmented velocities are still required to compute the acceptance ratio, though. Hence, there is no information loss and it should not affect the acceptance rate. \n\n> How many heads and how were length scales determined for each head in the kernel self-attention? ...\n\nWe performed a manual hyperparameter search to determine the appropriate length scales for our Gaussian kernels. The finalized length scales used in our model can be found in appendix F. The usage of Gaussian kernels allows even particles that are far apart to contribute, enhancing its capability to capture relevant information from various distances. We discovered that including both short length scales (approximately 0.1nm) and long length scales (around 1nm) proved to be beneficial for the overall performance of our model. \n\n> Can the authors provide additional references and discussion (perhaps in the appendix) of the evaluation metric presented in Eq 18? ...\n\nThe effective sample size (ESS) is commonly used in the literature to evaluate Markov chains.  A good summary of the ESS is given in the widely used STAN reference manual in 16.4.1. As we are interested in a wall-clock time comparison of our Timewarp algorithms and MD, we compare the ESS per second wall-clock time. To evaluate this metric, we need the sampling wall-clock time as well as the autocorrelations for different lag times (see Eqn. 18). Some example autocorrelation plots are reported in appendix B. We will extend the discussion in the final version of the manuscript.\n\n> Since Timewarp is evaluated on small peptide systems, conformer generation with a semi-empirical method like xTB might also be an appropriate evaluation...\n\nUsing the semi-empirical xtb force field instead of a classical one, would probably yield larger speed-ups compared to MD, as the evaluation of the semi-empirical force field is much more expensive and the Timewarp sampling algorithm requires significantly fewer (about 100 times) energy evaluations. Moreover, some of these energy evaluations can be performed in parallel using the batch sampling algorithm (Alg. 1). Hence, the more expensive the energy function, the more attractive the Timewarp method becomes. We thank the reviewer for bringing this up and will include this idea for future research directions in the Conclusion section. One reason we did not try this is that the data generation and especially the evaluation become orders of magnitude more expensive. \n\n> Are there any concrete proposals for future work that would address the current limitations ...'\n\nThe low acceptance rate is not a limitation in itself if the changes proposed are large enough that the effective sample size per second is fast for the model. The acceptance rate can be easily increased by predicting a smaller amount of time in the future, (e.g. if we set $\\tau$ near zero, we can get high acceptance rates just by predicting $x(t+ \\tau) = x(t)$) but this most likely results in a slower speed-up. \nScaling to larger systems is often done by considering a coarse grained representation.\nWe will include this in the limitations section.\n\n> Why do the authors use amber-14 in openMM rather than a more modern force field?\n\nWe use the amber-14 force field as it is commonly used in the literature and implemented in openMM. We do not think that other classical force fields would change our findings significantly. \n\n> The caption in Fig 4 does not match the figure labels. In Fig 3c the red dot indicating initial state is difficult to see.\n\nThanks for pointing this out, we will incorporate the suggested changes in the final version. "}, {"review_id": "TpGZdz5GN9", "reviewer": "Reviewer_UaAm", "summary": "This paper addresses high-throughput MD simulations.\nThe proposed sampling algorithm simulates the MD of 10^5-10^6 fs timestep by a pre-trained model of a normalizing flow.\nThe normalizing flow can be used as a proposal distribution in an MCMC method.\nThe experimental results show the pre-trained model trained on peptides with two or four amino acids generalized to unseen combinations of amino acids.\nThe proposed method is faster than naive MD in many cases of two amino acids.", "questions": "* Please provide the computational environment for the baseline MD simulation, amber-14. Is the wall clock comparison fair enough in terms of their implementation?\n* Why is the proposed method slower on 4AA? Is this because of O(N^2) computational complexity of the attention mechanism or the size of the neural networks?", "limitations": "The limitations are clearly stated in Section 7.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "* The paper is well-organized and easy to follow.\n* The proposed method can be integrated with a popular sampling method, MCMC.\n* The batched sampling method compensates for the low acceptance ratio of the proposed model. \n* The experimental results of the proposed exploration algorithm are also provided.\n* No speed-up for no slow transitions is reasonable.\n* The limitations are clearly shown in the experimental results.", "weaknesses": "* No comparison with conventional MD-acceleration methods such as metadynamics or parallel tempering. I would like to know how the proposed method is competitive with manually-tuned those methods.\n* The proposed method needs to be evaluated on application-oriented metrics such as the reproducibility of some physical properties.\n* Although the proposed method largely improved the simulation speed, the MH acceptance rates for the model are low, so the diffeomorphism function seems to need to be learned better.\n* According to Figure 3, the proposed method cannot reproduce the transition states, i.e., the hill of MD's energy surfaces, so the method may not be applicable to chemical reactions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful review and questions. We now address their questions individually. To stay within the character limit for our rebuttals, we often cite parts of the questions. \n\n> No comparison with conventional MD-acceleration methods ...\n\nWe do not compare to other MD-acceleration methods as Timewarp does not require any prior knowledge for unseen systems unlike common acceleration methods. Hence, we argue that the comparison provided in the paper is the fairest. \nHowever, Timewarp could also be accelerated by incorporating knowledge about the system into the algorithm. One possibility is to alternate Timewarp proposals with some (learned) change of collected variables, like the dihedral angles for peptides. These combined steps would still allow unbiased sampling from the target distribution. \nTimewarp could also be trained at multiple temperatures, by either training one model at each temperature or having the temperature as an additional conditioning input for a single model, which then could be applied similar to regular parallel tempering.\nWe will include these aspects in the limitations section and the potential applications in the Conclusion.\n\n> The proposed method needs to be evaluated on application-oriented metrics such as the reproducibility of some physical properties.\n\nWe utilize widely accepted evaluation metrics commonly found in the literature (Noé et al. Science 2019). These metrics not only adhere to established standards but also employ a level of rigor surpassing that of many other physical properties. Specifically, we evaluate energies alongside TICA plots, along with corresponding free energy projections. Additionally, we analyze bond lengths and autocorrelations, as detailed in appendix B.\n\n> ... the diffeomorphism function seems to need to be learned better.'\n\nThe low acceptance rate is not a limitation in itself. There is a tradeoff between the achieved acceptance rate and the length of the time step. For smaller time steps the acceptance rate is usually higher and vise versa. Due to  Alg 1, we can evaluate multiple samples in parallel at nearly no additional costs. Hence, a lower acceptance rate in combination with a larger time step can be beneficial.\nIn Figures 11 and 12 in appendix B.3, we show that the learned proposal distribution for unseen peptides matches the target closely. Moreover, the Timewarp exploration algorithm does not diverge. However, we agree that a more expressive proposal kernel would allow for larger time steps and ultimately result in even larger speed-ups compared to MD. We leave the exploration of more expressive architectures for future research.\n\n> the proposed method cannot reproduce the transition states, ..., so the method may not be applicable to chemical reactions.'\n\nWe agree that this is a limitation of the proposed method and will include this in the limitations section. However, this is a common challenge faced by generative machine learning methods, as they often generate fewer samples compared to traditional MD simulations within the same timeframe. Consequently, this scarcity of samples becomes even more pronounced in transition states. Nevertheless, it is crucial to highlight that despite this limitation, our Timewarp MCMC method still allows us to infer unbiased free energy differences.\n\n> Please provide the computational environment for the baseline MD simulation...\n\nThe used parameters for the MD simulations can be found in appendix E and G. We run the Timewarp Model as well as the MD simulations on the same type of GPUs for each experiment to allow for a fair wall-clock time comparison. For both we generate a single trajectory per GPU. A batched simulation is not available in existing fast MD libraries like OpenMM, as they are designed for simulating large systems serially. To parallelise, a custom implementation is needed to run separate simulations on each GPU thread. Also, if multiple trajectories were generated with MD, it is not trivial to combine these into a single unbiased trajectory and compute ESS/s, as we do now. If instead we were not concerned with correctly reweighting multiple trajectories, we could then run Timewarp in exploration mode (Alg 2), which easily produces 100 trajectories in parallel on a single GPU, which is comparable or greater than the number of parallel MD simulations which could be run in theory.\nWe do not expect that our results would change drastically for using a different classical force field or MD library. \n\n> Why is the proposed method slower on 4AA? ...\n\nThis is an excellent question. There is no one reason that the speed-up is slower for the 4AA system. The four main factors are the time step $\\tau$, the acceptance rate, the forward pass through the flow, and the cost of the energy evaluation. As discussed earlier, there is a tradeoff between the time step and the acceptance rate. As described in appendix F the model for 4AA is about four times larger than for 2AA, resulting in an overall slower forward pass. As the attention weights are only computed once during each forward pass, the influence on the forward pass is marginal. The energy evaluation is slightly more expensive for 4AA compared to 2AA, which helps with the overall speed-up compared to MD as Timewarp requires about 100 times less energy evaluation per effective sample. Hence, by using a more expensive energy function, such as a semi-empirical force field, would probably improve the results significantly. We leave this to future work, as the MD simulations for evaluation purposes become orders of magnitude more expensive.\nIn case of the 4AA system the most important factor for the lower speed-up is the smaller time step, which was required to achieve acceptance rates around 1%, as shown in appendix B. More tuning could result in better performance. \nWe will add parts of this discussion to the limitations and conclusion section. "}]

Author rebuttal:
We thank all the reviewers for taking the time to review our paper.
A number of authors have asked for code and datasets to be released. We will release the code upon publication under an open source license. We will also publicly release the datasets that we used to train and evaluate the models. Unfortunately we cannot release our code and datasets anonymously before publication.

Meta-review: All reviewers really liked this paper so this is clearly an accept. 

One may also perhaps argue for a spotlight presentation. However, this AC thinks that the method has the limitation that it does not allow for estimation of physical observables such as folding rates. So there are still room for improvement.
Decision: Accept (spotlight)