Title: Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics
Venue: NeurIPS 2023
Accepted: yes
Keywords: Molecular Dynamics

Normalizing Flows

MCMC
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

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 fs = 10 -15 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 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.

1 Introduction

Molecular dynamics (MD) is a well-established technique for simulating physical systems at 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-ligand binding [1], and protein-protein association [31]. Many crucial applications of MD boil down to efficiently sampling from the Boltzmann distribution , i.e. , the equilibrium distribution

MD

Timewarp MCMC

a)

c(t)

x(t+ T)

T =0.5×106fs

b)

康

of a molecular system at a temperature T . Let ( x p ( t ) , x v ( t )) := x ( t ) ∈ R 6 N be the state of the molecule at time t , consisting of the positions x p ( t ) ∈ R 3 N and velocities x v ( t ) ∈ R 3 N of the N atoms in Cartesian coordinates. The Boltzmann distribution is given by:

where U ( x p ) is the potential energy, K ( x v ) is the kinetic energy, and k B 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 µ ( x p ) . A popular MD method to sample from µ ( x p ) is Langevin dynamics [19], which obeys the following stochastic differential equation (SDE):

Here i indexes the atoms, m i is the mass of atom i , U ( x p ) is the potential energy, γ is a friction parameter, and d B t is a standard Brownian motion process. Starting from an initial state x (0) , simulating Equation (2), along with the relationship d x p = x v d t , yields values of x ( t ) that are distributed according to the Boltzmann distribution as t →∞ . Standard MD libraries discretise this SDE with a timestep ∆ t , which must be chosen to be ∼ 1 fs = 10 -15 s for stability. Unfortunately, many biomolecules contain metastable states separated by energy barriers that can take milliseconds of MD simulation time ( ∼ 10 12 sequential integration steps) to cross, rendering this approach infeasible. To overcome this, prior work has produced an array of enhanced sampling methods , such as coarse graining [3, 14] and metadynamics [18]. 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 a 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 to study this kind of transferability. We propose Timewarp , a general, transferable enhanced sampling method which uses a normalising flow [35] as a proposal for a Markov chain Monte Carlo (MCMC) method targeting the Boltzmann distribution. Our main contributions are:

We present the first ML algorithm working in general Cartesian coordinates that demonstrates transferability to small peptide systems unseen during training.

We demonstrate, for the first time, wall-clock acceleration of asymptotically unbiased MCMCsampling of the Boltzmann distribution of unseen peptide systems.

We define an MCMC algorithm targeting the Boltzmann distribution using a conditional normalising flow as a proposal distribution, with a Metropolis-Hastings (MH) correction step to ensure detailed balance (Section 3.4).

We design a permutation equivariant, transformer-based normalising flow.

We produce a novel training dataset of MD trajectories of thousands of small peptides.

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: https://github.com/microsoft/timewarp . The datasets are available upon request 3 .

2 Related work

There has recently been a surge of interest in deep learning on molecules. Boltzmann generators [28, 17, 15] use flows to sample from the Boltzmann distribution asymptotically unbiased. There are two ways to generate samples: (i) Produce i.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 in torsion space and use the underlying ODE as a transferable Boltzmann generator. However, in contrast to Timewarp, they use internal coordinates and do not operate in the all atom system. Moreover, [13, 25] recently introduced a Boltzmann generators in Cartesian coordinates for molecules, potentially enabling transferable training. Recently, [46] proposed GeoDiff , a diffusion model that predicts molecular conformations from a molecular graph. Like Timewarp, GeoDiff works in 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, in recent work, [8] trained a transferable ML model to simulate the time-coarsened dynamics of polymers. However, unlike Timewarp, their model acts on coarse grained representations. Additionally, it 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, 40, 10] work by running many short MD simulations, which are used to define a discrete state space, along with an estimated transition probability matrix. Similarly to Timewarp, MSMs estimate the transition probability between the state at a time t and the time t + τ , where τ ≫ ∆ t . Recent work has applied deep learning to MSMs, leading to VAMPnets [23] and deep generative MSMs [45], which replace the MSM data-processing pipeline with deep networks. In contrast to Timewarp, these models are not transferable and model the dynamics in a 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, 21], which use deep generative models as proposal distributions. A-NICE-MC [38] uses a volume-preserving flow trained using a likelihood-free adversarial method. Other methods use objective functions designed to encourage exploration. The entropy term in our objective function is inspired by [41]. 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 x ( t + τ ) induced by an MD simulation of Equation (2) for a time τ ≫ ∆ t , starting from x ( t ) . We denote this conditional distribution by µ ( x ( t + τ ) | x ( t )) . Timewarp uses a deep probabilistic model to approximate µ ( x ( t + τ ) | x ( t )) (see Figure 1). Once trained, the model is used in an MCMC method to sample from the Boltzmann distribution.

3.1 Conditional normalising flows

We fit a conditional normalising flow , p θ ( x ( t + τ ) | x ( t )) , to µ ( x ( t + τ ) | x ( t )) , where θ are learnable parameters. Normalising flows are defined by a base distribution (usually a standard Gaussian), and a diffeomorphism f , i.e. a differentiable bijection with a differentiable inverse. Specifically, we set p θ ( x ( t + τ ) | x ( t )) as the density of the distribution defined by the following generative process:

Here z p ∈ R 3 N and z v ∈ R 3 N . For all settings of θ and x ( t ) , f θ ( · ; x ( t )) is a diffeomorphism that takes the latent variables ( z p , z v ) ∈ R 6 N to ( x p ( t + τ ) , x v ( t + τ )) ∈ R 6 N . The conditioning state x ( t ) parameterises a family of diffeomorphisms, defining a conditional normalising flow [44]. Note that there are no invertibility constraints on the mapping from x ( t ) to the output x ( t + τ ) , only the map from z to x ( t + τ ) must be invertible. Using the change of variables formula, we can evaluate:

where f -1 θ ( · ; x ( t )) : R 6 N → R 6 N is the inverse of the diffeomorphism f θ ( · ; x ( t )) , and J f -1 θ ( · ; x ( t )) ( x ( t + τ )) denotes the Jacobian of f -1 θ ( · ; x ( t )) evaluated at x ( t + τ ) .

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 = {T i } P i =1 , where P is the number of peptides. Each MD trajectory is temporally subsampled with a spacing of τ , so that T i = ( x (0) , x ( τ ) , x (2 τ ) , . . . ) . During training, we randomly sample pairs x ( t ) , x ( t + τ ) from D . Each pair represents a sample from the conditional distribution µ ( x ( t + τ ) | x ( t )) . Additional details are provided in Appendix E. 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 x p , rather than the velocities x v . Thus, it is not necessary for x v ( t ) , x v ( t + τ ) to represent the actual velocities of the atoms in Equation (3). We hence simplify the learning problem by treating x v as non-physical auxiliary variables within the augmented normalising flow framework [9]. For each datapoint x ( t ) = x p ( t ) , x v ( t ) in D , instead of obtaining x v ( t ) by recording the velocities in the MD trajectory, we discard the MD velocity and independently draw x v ( t ) ∼ N (0 , I ) . The auxiliary variables x v ( t ) now contain no information about the future state x p ( t + τ ) , x v ( t + τ ) , since x v ( t ) and x v ( t + τ ) are drawn independently. Hence we can simplify f θ to depend only on x p ( t ) , with x p ( t + τ ) , x v ( t + τ ) := f θ ( z p , z v ; x p ( t )) . We include auxiliary variables for two reasons: First, they increase the expressivity of the distribution for x p without a prohibitive increase in computational cost [9, 2]. Second, constructing a conditional flow that respects permutation equivariance is simplified with auxiliary variables - see Section 4.1.

3.4 Targeting the Boltzmann distribution with asymptotically unbiased MCMC

After training the flow p θ ( x ( t + τ ) | x ( t )) , we use it as a proposal distribution in an MCMC method to target the joint distribution of the positions x p and the auxiliary variables x v , which has density:

Starting from an initial state X 0 = ( X p 0 , X v 0 ) ∈ R 6 N for state m = 0 , we iterate:

where α ( X m , ˜ X m ) is the Metropolis-Hastings (MH) acceptance ratio [24] targeting Equation (4):

The flow used for p θ must allow for efficient sampling and exact likelihood evaluation, which is crucial for fast implementation of Equations (5) and (6). Additionally, after each MH step, we resample the auxiliary variables X v using a Gibbs sampling update:

Iterating these updates yields a sample X p m , X v m ∼ µ aug as m → ∞ . To obtain a Boltzmanndistributed sample of the positions X p m ∼ µ , we simply discard the auxiliary variables X v m . As sending m →∞ is infeasible, we simulate the chain until a 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 1 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 a 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 accurate free energy estimates. It also succeeds in discovering all metastable states orders of magnitude faster than Algorithm 1 and standard MD, which could be used, e.g. , to provide initialisation states for a subsequent MSM method. Pseudocode is given in Algorithm 2 in Appendix D.

4 Model architecture

We now describe the architecture of the flow f θ ( z p , z v ; x p ( t )) , which is shown in Figure 2.

RealNVP coupling flow Our architecture is based on RealNVP [5], which consists of a 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 ℓ th coupling layer of the flow, the following transformations are implemented:

Going forward, we suppress the coupling layer index ℓ . Here ⊙ is the element-wise product, and s p θ : R 3 N → R 3 N is our atom transformer , a neural network based on the transformer architecture [43] that takes the auxiliary latent variables z v and the conditioning state x ( t ) and outputs scaling factors for the position latent variables z p . The function t p θ : R 3 N → R 3 N is implemented as another atom transformer, which uses z v and x ( t ) to output a translation of the position latent variables z p . 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 f θ consists of N coupling stacked coupling layers, beginning from z ∼ N (0 , I ) and ending with a sample from p θ ( x ( t + τ ) | x ( 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 ( t + τ ) -x ( t ) , rather than x ( t + τ ) directly.

Atom transformer We now describe the atom transformer network. Let x p i ( t ) , z p i , z v i , all elements of R 3 , denote respectively the position of atom i 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 z v as input (such as s p θ ( z v , x p ( t )) and t p θ ( z v , x p ( t )) in Equation (8)), we first concatenate the variables associated with atom i . This leads to a vector a i := [ x p i ( t ) , h i , z v i ] T ∈ R H +6 , where z p i has been excluded since s p θ , t p θ are not allowed to depend on z p . Here h i ∈ R H is a learned embedding vector which depends only on the atom type. The vectors a i are fed into an MLP ϕ in : R H +6 → R D , where D is the feature dimension of the transformer. The vectors ϕ in ( a 1 ) , . . . , ϕ in ( a N ) are then fed into N transformer stacked transformer layers. After the transformer layers, they are passed through another atom-wise MLP, ϕ out : R D → R 3 . The final output is in R 3 N as required. This is depicted in Figure 2, Middle. When implementing s v θ and t v θ from Equation (9), a similar procedure is performed on the vector [ x p i ( t ) , h i , z p i ] T , but now including z p i and excluding z v i . There are two key differences between the atom transformer and the architecture in [43]. First, to maintain permutation equivariance,

Atom

transformer

scale

Add & normalize

Add & normalize

Multihead kernel self-attention

Atom-wise MLP

Add & normalize

Atom-wise MLP

Add & normalize

Atom-wise MLP

Outputs from atom-wise MLP

Linear

Kernel self-attention

Attention weight

Attention weight

Attention weight

translate

Atom

transformer

Atom

transformer

scale

Atom

transformer

translate

Atom

transformer

scale

Linear

Linear

Concat &

linear

Concat &

linear

Concat &

linear

we do not use a 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 at time t + τ will be more influenced by atoms that are nearby at time t , compared to atoms that are far away. Thus, we define the attention weight w ij for atom i attending to atom j as follows:

where ℓ is a lengthscale parameter. The outputs { r out ,i } N i =1 , given the inputs { r in ,i } N i =1 , are then:

where V ∈ R d out × d in is a learnable matrix. This kernel self-attention is an instance of the RBF kernel attention investigated in [42]. Similarly to [43], we introduce a multihead version, where each head has a different lengthscale. This is 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 of these symmetries is incorporated in Timewarp.

Permutation equivariance Let σ be a permutation of the N atoms. Since the atoms have no intrinsic ordering, the only effect of a permutation of x ( t ) on the future state x ( t + τ ) is to permute the atoms similarly, i.e. ,

Our conditional flow satisfies permutation equivariance exactly. To show this, we use the following proposition proved in Appendix A.1, which is an extension of [16, 36] for conditional flows:

Proposition 4.1. Let σ be a symmetry action, and let f ( · ; · ) be an equivariant map such that f ( σz ; σx ) = σf ( z ; x ) for all σ, z, x . Further, let the base distribution p ( z ) satisfy p ( σz ) = p ( z ) for all σ, z . Then the conditional flow defined by z ∼ p ( z ) , x ( t + τ ) := f ( z ; x ( t )) satisfies p ( σx ( t + τ ) | σx ( t )) = p ( x ( t + τ ) | x ( t )) .

Our flow satisfies f θ ( σz ; σx ( t )) = σf θ ( z ; x ( t )) since the transformer is permutation equivariant, and permuting z and x ( t ) together permutes the inputs. Furthermore, the base distribution p ( z ) = N (0 , I ) is permutation invariant. Note that the presence of auxiliary variables allows us to easily construct a permutation equivariant coupling layer.

Translation and rotation equivariance Consider a transformation T = ( R,a ) that acts on x p :

where R is a 3 × 3 rotation matrix, and a ∈ R 3 is a translation vector. We would like the model to satisfy p θ ( Tx ( t + τ ) | Tx ( t )) = p θ ( x ( t + τ ) | x ( t )) . We achieve translation equivariance by subtracting the average position of the atoms in the initial state (Appendix A.2). Rotation equivariance is not encoded in the architecture but is handled by data augmentation: each training pair ( x ( t ) , x ( t + τ )) is acted upon by a random rotation matrix R to form ( Rx ( t ) , Rx ( t + τ )) 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 ( k ) ( t ) , x ( k ) ( t + τ )) } K k =1 represents all pairs of states at times τ apart in D . We optimise:

(ii) acceptance training , the model is fine-tuned to maximise the probability of MH acceptance. Let x ( k ) ( t ) be sampled uniformly from D . Then, we use the model to sample ˜ x ( k ) θ ( t + τ ) ∼ p θ ( · | x ( k ) ( t )) using Equation (3). We use this to optimise the acceptance probability in Equation (6) with respect to θ . Let r θ ( X, ˜ X ) denote the model-dependent term in the acceptance ratio in Equation (6):

The acceptance objective is then given by:

Training to maximise the acceptance probability can lead to the model proposing changes that are too small: if ˜ x ( k ) θ ( t + τ ) = x ( k ) ( t ) , then all proposals will be accepted. To mitigate this, during acceptance training, we use an objective which is a weighted average of L acc ( θ ) , L lik ( θ ) and a Monte Carlo estimate of the average differential entropy,

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 to traverse. To find these, we use time-lagged independent component analysis (TICA) [30], a 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

orojections

MD

Timewarp

f)

MD

Free energy/kgT

10

Timewarp

Qw

=

1

5

-2

0

2

TICO

00

0

Psi

MD

Freeenergy/kgT

10

5

Timewarp

0

3

4

TICO

TIC 0

tiontimeinh

a

MD

Timewarp

b)

2

Free energy/kgT

10

0

5

0

-2

0

2

-2

0

2

c)

Phi

Phi

d)

2

0

0

-2

0

-2

-2

2

0

2

2

0

2

0

Phi

Phi

yprojections

MD

Exploration

MCMC

Psi

MD

Timewarp

3

4

where M is the chain length, M eff is the effective number of samples, t sampling is the sampling wallclock time, and ρ τ is the autocorrelation for the lag time τ [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 2 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 ( 50 ns = 10 8 steps), and evaluate on 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 4 ), making the generalisation task significantly more difficult than for 2AA. For more details see Table 2 in Appendix E.

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 Algorithm 1 for a chain length of 10 million, accepting roughly 2% of the proposals. In Figure 3a we visualise the samples using a Ramachandran plot [33], which shows the distribution of the backbone dihedral angles φ and ψ . Each mode in the plot represents a metastable state. We see that the Timewarp samples closely match MD, visiting all the metastable states with the correct relative weights. In Figure 3b we plot the free energy ( i.e. , the relative log probability) of the φ and ψ angles, again showing close agreement. The roughness in the plot is due to some regions of state space having very few samples. In Figure 3c we show, for an initial state x ( t ) , the conditional distribution of MD obtained by integrating Equation (2), µ ( x ( t + τ ) | x ( t )) , compared with the model p θ ( x ( t + τ ) | x ( t )) , demonstrating close agreement. Finally, Figure 3d shows the time-evolution of the φ angle for MD and Timewarp. Timewarp exhibits significantly more transitions between the metastable states than MD. As a result, the autocorrelation along the φ angle decays much faster in terms of wall-clock time, resulting in a ≈ 7 × 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 1 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.1). The results are shown for the dipeptides QW and HT in Figure 3ef, showing close agreement between Timewarp and long MD chains (1 µ s = 2 × 10 9 steps). For these dipeptides, Timewarp achieves ESS/s speed-up factors over MD of 5 and 33 respectively (Appendix B.4). In Figure 4, 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. In 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

a)

0

c)

101

peed-upTIC0

d)4AAexplorat

102

101

100

ptides

Te

0

b)

2A

703

a)

MD

Exploration

MCMC

b)

Free energy/kgT

10

SAEL

TIC 1

5

0

Free energy/kgT

10

5

0

Free energy/kgT

10

TIC 1

5

0

TICO

TICO

projections

C

MD

Exploration

MCMC

Y

distribution

MD

C.3

Exploration

MCMC

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 a median speedup factor of ≈ 600 (Figure 4c). 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 are qualitatively similar to MD, but less accurate than Timewarp MCMC (Figure 3f).

Tetrapeptides (4AA) Finally, we study the more challenging 4AA dataset. After training on the trainset, we sample 20 million Markov chain states for each test tetrapeptide using Algorithm 1 and compare with long MD trajectories ( 1 µ s). In contrast to the simpler dipeptides, both Timewarp MCMCand 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 5a 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 to MD in Figure 4b,d. Although Timewarp MCMC fails to speed up sampling for most tetrapeptides, Timewarp exploration shows a median speedup factor of ≈ 50 . For 8 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 5b shows that when Timewarp MCMC discovers all metastable states, its free energy estimates match those of MD very well. However, it sometimes

misses metastable states leading to poor free energy estimates in those regions. Figure 5c shows that Timewarp MCMC also leads to a 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 is slightly too high relative to MD and Timewarp MCMC.

7 Limitations / Future work

The Timewarp MCMC algorithm generates low acceptance probabilities ( < 1% ) for most peptides (see Appendix B.1). However, this is not a limitation in itself. In general, a larger proposal timestep τ yields smaller acceptance rates as the prediction problem becomes more difficult. However, due to Algorithm 1, we can evaluate multiple samples in parallel at nearly no additional costs. As a result, a lower acceptance rate, when coupled with a larger timestep τ , is often a favorable trade-off. While we speed-up only roughly a 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 a 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 computational cost (Appendix F). Timewarp generates fewer samples compared to 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 to 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 a topic for future research, and there are several promising avenues to consider. One approach is to explore different network architectures, potentially capturing all the symmetries inherent in 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 an MH correction can accelerate asymptotically unbiased sampling on many unseen dipeptides, allowing faster computation of equilibrium expectation values. Although this acceleration was only possible for a minority of the tetrapeptides we considered, we showed that Timewarp used without the MH correction explores the metastable states of both dipeptides and tetrapeptides much faster than standard MD, and we verify the metastable states discovered are 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 to scale Timewarp to larger, more interesting biomolecules, this work clearly demonstrates the ability of deep learning algorithms to leverage transferability to accelerate the MD sampling problem.

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 AI4Science team for insightful discussions and computing help.

References

Ignasi Buch, Toni Giorgino, and Gianni De Fabritiis. Complete reconstruction of an enzymeinhibitor binding process by molecular dynamics simulations. Proceedings of the National Academy of Sciences , 108(25):10184-10189, 2011.

Jianfei Chen, Cheng Lu, Biqi Chenli, Jun Zhu, and Tian Tian. Vflow: More expressive generative flows with variational data augmentation. In International Conference on Machine Learning , pages 1660-1669. PMLR, 2020.

Cecilia Clementi. Coarse-grained models of protein folding: toy models or predictive tools? Current opinion in structural biology , 18(1):10-15, 2008.

Manuel Dibak, Leon Klein, Andreas Krämer, and Frank Noé. Temperature steerable flows and Boltzmann generators. Phys. Rev. Res. , 4:L042005, Oct 2022.

Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations , 2017.

David J. Earl and Michael W. Deem. Parallel tempering: Theory, applications, and new perspectives. Phys. Chem. Chem. Phys. , 7:3910-3916, 2005.

Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle A Beauchamp, Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern, et al. Openmm 7: Rapid development of high performance algorithms for molecular dynamics. PLoS computational biology , 13(7):e1005659, 2017.

Xiang Fu, Tian Xie, Nathan J. Rebello, Bradley Olsen, and Tommi S. Jaakkola. Simulate timeintegrated coarse-grained molecular dynamics with multi-scale graph networks. Transactions on Machine Learning Research , 2023.

C. Huang, Laurent Dinh, and Aaron C. Courville. Augmented normalizing flows: Bridging the gap between generative flows and latent variable models. ArXiv , abs/2002.07101, 2020.

Brooke E Husic and Vijay S Pande. Markov state models: From an art to a science. Journal of the American Chemical Society , 140(7):2386-2396, 2018.

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.

William Thomson Baron Kelvin. The molecular tactics of a crystal . Clarendon Press, 1894.

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

Sebastian Kmiecik, Dominik Gront, Michal Kolinski, Lukasz Wieteska, Aleksandra Elzbieta Dawid, and Andrzej Kolinski. Coarse-grained protein models and their applications. Chemical reviews , 116(14):7898-7936, 2016.

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.

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.

Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems , volume 34, pages 2796-2809. Curran Associates, Inc., 2021.

Alessandro Laio and Michele Parrinello. Escaping free-energy minima. Proceedings of the National Academy of Sciences , 99(20):12562-12566, 2002.

Paul Langevin. Sur la théorie du mouvement brownien. Compt. Rendus , 146:530-533, 1908.

Daniel Levy, Matt D. Hoffman, and Jascha Sohl-Dickstein. Generalizing Hamiltonian Monte Carlo with neural networks. In International Conference on Learning Representations , 2018.

Zengyi Li, Yubei Chen, and Friedrich T. Sommer. A neural network MCMC sampler that maximizes proposal entropy. Entropy , 23(3), 2021.

Kresten Lindorff-Larsen, Stefano Piana, Ron O Dror, and David E Shaw. How fast-folding proteins fold. Science , 334(6055):517-520, 2011.

Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. VAMPnets for deep learning of molecular kinetics. Nature communications , 9(1):1-11, 2018.

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):1087-1092, 1953.

Laurence I Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Se (3) equivariant augmented coupling flows. arXiv preprint arXiv:2308.10364 , 2023.

Radford M Neal. Probabilistic inference using Markov chain Monte Carlo methods . Department of Computer Science, University of Toronto Toronto, ON, Canada, 1993.

Jerome P Nilmeier, Gavin E Crooks, David DL Minh, and John D Chodera. Nonequilibrium candidate monte carlo is an efficient tool for equilibrium simulation. Proceedings of the National Academy of Sciences , 108(45):E1009-E1018, 2011.

Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators - sampling equilibrium states of many-body systems with deep learning. Science , 365:eaaw1147, 2019.

Frank Noé, Christof Schütte, Eric Vanden-Eijnden, Lothar Reich, and Thomas R Weikl. Constructing the equilibrium ensemble of folding pathways from short off-equilibrium simulations. Proceedings of the National Academy of Sciences , 106(45):19011-19016, 2009.

Guillermo Pérez-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 , 139(1):07B604_1, 2013.

Nuria Plattner, Stefan Doerr, Gianni De Fabritiis, and Frank Noé. Complete protein-protein association kinetics in atomic detail revealed by molecular dynamics simulations and Markov modelling. Nature chemistry , 9(10):1005-1011, 2017.

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.

G N Ramachandran, C Ramakrishnan, and V Sasisekharan. Stereochemistry of polypeptide chain configurations. Journal of Molecular Biology , pages 95-99, 1963.

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.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning , pages 1530-1538. PMLR, 2015.

Danilo Jimenez Rezende, Sébastien Racanière, Irina Higgins, and Peter Toth. Equivariant Hamiltonian flows. arXiv preprint arXiv:1909.13739 , 2019.

Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time-resolution surrogates for molecular dynamics. arXiv preprint arXiv:2305.18046 , 2023.

Jiaming Song, Shengjia Zhao, and Stefano Ermon. A-NICE-MC: Adversarial training for MCMC. Advances in Neural Information Processing Systems , 30, 2017.

Robert H Swendsen and Jian-Sheng Wang. Replica monte carlo simulation of spin-glasses. Physical review letters , 57(21):2607, 1986.

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(21):65716581, 2004.

Michalis Titsias and Petros Dellaportas. Gradient-based adaptive Markov chain Monte Carlo. Advances in Neural Information Processing Systems , 32, 2019.

Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-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 Processing (EMNLP-IJCNLP) , pages 4344-4353, Hong Kong, China, November 2019. Association for Computational Linguistics.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems , 30, 2017.

Christina Winkler, Daniel Worrall, Emiel Hoogeboom, and Max Welling. Learning likelihoods with conditional normalizing flows. arXiv preprint arXiv:1912.00042 , 2019.

Hao Wu, Andreas Mardt, Luca Pasquali, and Frank Noé. Deep generative Markov state models. Advances in Neural Information Processing Systems , 31, 2018.

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.

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 X ( t + τ ) x ( t ) denote the random variable obtained by sampling Z ∼ p ( z ) and computing X ( t + τ ) := f ( Z ; x ( t )) . Here we subscript X ( t + τ ) x ( t ) by x ( t ) to emphasize that this is the random variable obtained when conditioning the flow on x ( t ) . We first note that the equivariance condition on the densities p ( σx ( t + τ ) | σx ( t )) = p ( x ( t + τ ) | x ( t )) is equivalent to the following constraint on the random variables:

where d = denotes equality in distribution. To see this, let p X denote the density of the random variable X . Then, in terms of densities, Equation (19) is equivalent to stating that, for all x ,

where in Equation (21) we used the change-of-variables formula, along with the fact that the group actions we consider (rotations, translations, permutations) have unit absolute Jacobian determinant. Redefining x ← σx , we get that for all x ,

Recalling the notation that X ( t + τ ) x ( t ) is interpreted as the random variable obtained by conditioning the flow on x ( t ) , this can be written as

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.

where in Equation (25) we used the fact that the base distribution p ( z ) is σ -invariant.

A.2 Translation equivariance via canonicalisation

We now describe the canonicalisation technique used to make our models translation equivariant. Let q ( x p ( t + τ ) , x v ( t + τ ) | x p ( t )) be an arbitrary conditional density model, which is not necessarily translation equivariant. We can make it translation equivariant in the following way. Let x p denote the average position of the atoms,

Then we define

where the subtraction of x p ( t ) is broadcasted over all atoms. We now consider the effect of translating both x p ( t ) and x p ( t + τ ) by the same amount. Let a be a translation vector in R 3 . Then

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 ( x p ) of a molecular configuration is also invariant under mirroring. However, in the presence of chirality centers, a mirrored configuration is non-superposable to its original image [12]. An example of a chirality center in an amino acid is a 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 a peptide to a D-amino acid in a 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 a change occurs. This does not add any significant computational overhead.

B Additional results

In this section we show additional results like the conditional distribution as well as more peptide examples for experiments discussed in Section 6.

B.1 2AA additional results

More examples from the 2AA test set are presented in Figures 7 and 8. 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 6

B.2 4AA additional results

More examples from the 4AA test set are presented in Figures 9 and 11. The distribution of the acceptance probabilities for the Timewarp MCMC algorithm is shown in Figure 6.

B.3 Conditional distributions

The model was trained to generate samples from the conditional Boltzmann distribution µ ( x ( t + τ ) | x ( t )) . Here we show some examples of the conditional distribution generated by the Timewarp model compared to 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 τ to 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 a 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,

10

a)

MD

Timewarp

b)

Freeenergyprojections

c)

TICOautocorrelations

d)

TICOtimedepedence

Free energy/ksT

10

MD

Timewarp

Autocorrelation

1.0

MD| ESS/t = 0.005/s

Timewarp | ESS/t = 0.024/s

TIC0

MD

TIC 1

5

0.5

0

Timewarp

00

Free energy/kgT

10

0

Autocorrelation

1.0

MD | ESS/t = 0.003/s

Timewarp | ESS/t = 0.096/s

TICO

MD

24

0.5

1

5

0

Timewarp

0.0

0

Free energy/kgT

10

Autocorrelation

1.0

Timewarp | ESS/t = 0.010/s

MD I ESS/t = 5.496/s

TICO

MD

0.5

5

TICO

559957 5984 3218 7257881

Timewarp

0.0

0

TICO

TIC0

TIC 0

TIC 1

0

200

400600800

1000

0

2

4

6

8

Lag time in s

Wall-clock sampling time in h

a)

MD

Timewarp

Free energy/ksT

Free energy projections

c)

TIC O autocorrelations

d)

TICotimedepedence

10

MD

Autocorrelation

1.0

MD|ESS/t = 0.008/s

Timewarp

Timewarp | ESS/t = 0.242/s

TIC 1

5

0.5

0

Timewarp

1.0

0.0

Autocorrelation

MD | ESS/t = 0.004/s

Timewarp | ESS/t = 0.328/s

0.5

5

Timewarp

1.0

0.0

Autocorrelation

MD | ESS/t = 0.006/s

中

MD

Timewarp | ESS/t = 0.042/s

TIC

0.5

FS2

Timewarp

0.0

Free energy/ksT

Autocorrelation

1.0

MD | ESS/t = 0.005/s

0

Timewarp | ESS/t = 0.090/s

TIC1

５

0.5

TICO

0.0

0

00

00

0

TIC 1

0

200

400600800

1000

0

5

10

15

Lag time in s

Wall-clock sampling time in h

a)

MD

Exploration

MCMC

b)

Free energy/kgT

10

KC

5

0

Free energy/kgT

10

0

Free energy/kgT

10

F

5

0L

TICO

TIC 0

TIC 0

projections

c

MD

Exploration

MCMC

MD

Exploration

MCMC

ICO

a)

Boltzmann

MD

Timewarp

b)

2

6

0

（d）

601

4

2

一

2

Phi

0

2

-2

Phi

0

2

-2

Phi

0

2

0

一

2

projections

c

MD

Timewarp

distribution

(p

Phiautocorrelations

MD

ocorrelation

1.0

二

MD|ESS/t= 0.032/s

Timewarp

modelESS/t=0.228/s

0.5

emphasising the importance of the Metropolis-Hastings correction to obtain unbiased samples from the Boltzmann distribution with the Timewarp model.

B.4 Autocorrelations

In Section 6 we compute the speedup of the Timewarp model by comparing the effective sample sizes per second (Equation (18)) for the slowest transition with MD. As the ESS depends on the autocorrelation, it is also insightful to look at the autocorrelation decay in terms of wall-clock time. We show some example autocorrelations for the investigated peptides in Figures 7, 8 and 10. Note that the area under the autocorrleation curve is inversely proportional to 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 in Figure 12a. However, we can easily explore these with the Timewarp exploration algorithm (Algorithm 2). To confirm that these additional metastable states are in fact

a

Boltzmann

MD

Timewarp

b)

(d)60j-

4

2

0

(d)60j-

2

6

AAEW

(d)601-

4

2

AAEW

TIC 1

2

0

6

EASS

(d)601-

4

TIC

2

0

6

(d)60l-

4

2

0

TICO

TICO

TICO

c)

distribution

(P

MD

Timewarp

distribution

MD

model

MD

Timewarp

stable, we run several shorter MD trajectories ( 0 . 5 × 10 6 fs), in the same way as in Appendix B.3, starting in 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 × 10 6 fs) from there as well (Figure 12c). These new short MD trajectories either sample within this previously undiscovered state, or 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

Pseudocode for the algorithm described in Section 3.4 is given in Algorithm 1.

D Exploration of metastable states using Timewarp without MH corrections

We describe the exploration algorithm for Timewarp, which accepts all proposed states unless the 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 or Algorithm 1. The pseudocode is shown in Algorithm 2:

a

MD

Exploration

MCMC

b)

MD

Timewarp

MD

Timewarp

C

MD

Timewarp

CTSA

TIC 1

ICMC

b)

Note that unlike Algorithm 1, there is no 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 a large penalty if the ordering of a chirality center changes as described in Appendix A.3. As sampling proposals from Timewarp can be batched, we can generate B 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 a batch size of 100, and run M = 10000 exploration steps. The maximum allowed energy change cutoff is set at ∆ U max = 300 kJ/mol.

E Dataset details

We evaluate our model on three different datasets, AD, 2AA, and 4AA, as introduced in Section 6. All datasets are simulated in implicit solvent using the openMM library [7]. For all MD simulations we use the parameters shown in Table 1.

We present more dataset details, like simulation times and number of peptides, in Table 2.

F Hyperparameters

Depending on the dataset, different Timewarp model sizes were used, as shown in Table 3. For all datasets the Multihead kernel self-attention layer consists of 6 heads with lengthscales ℓ i = { 0 . 1 , 0 . 2 , 0 . 5 , 0 . 7 , 1 . 0 , 1 . 2 } , given in nanometers.

The ϕ in and ϕ out MLPs use SiLUs as activation functions, while the Transformer MLPs use ReLUs. 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 4.

The first linear layers in the kernel self-attention module always has the shape [128 , 768] (in Section 4 denoted as V ), and the second (after concatenating the output of head head) has the shape [768 , 128] . The transformer feature dimension D is for all datasets 128 .

After likelihood training, we fine-tune the model for the AD and 2AA dataset with a 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 5. We use the

FusedLamb optimizer and the DeepSpeed library [34] for all experiments. The batch size as well as the training times are reported in Table 6. All simulations are started with a learning rate of 5 × 10 -4 ,

the learning rate is then consecutively decreased by a factor of 2 upon hitting training loss plateaus.

G Computing infrastructure

The training was performed on 4 NVIDIA A-100 GPUs for the 2AA and 4AA datasets and on a 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n* 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.\n\nWEAKNESSES:\n* 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nThe limitations are clearly stated in Section 7.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nCan 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\nLIMITATIONS:\nYes, although further discussion of possible avenues for addressing these limitations would be welcome.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\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\nSTRENGTHS:\nThis 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\nWEAKNESSES:\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\nQUESTIONS:\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\nLIMITATIONS:\nNo potential negative societal impact", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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\nSTRENGTHS:\nNOTE: 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.\n\nWEAKNESSES:\n- 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\nQUESTIONS:\nAs 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?\n\nLIMITATIONS:\nThe authors have sufficiently addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}