SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Proposition 4.1. Let  $\sigma$  be a symmetry action, and let  $f(\cdot ;\cdot)$  be an equivariant map such that  $f(\sigma z;\sigma x) = \sigma f(z;x)$  for all  $\sigma ,z,x$ . Further, let the base distribution  $p(z)$  satisfy  $p(\sigma z) = p(z)$  for all  $\sigma ,z$ . Then the conditional flow defined by  $z\sim p(z)$ ,  $x(t + \tau)\coloneqq f(z;x(t))$  satisfies  $p(\sigma x(t + \tau)|\sigma x(t)) = p(x(t + \tau)|x(t))$ .
Our flow satisfies  $f_{\theta}(\sigma z; \sigma x(t)) = \sigma f_{\theta}(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) = \mathcal{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$ :
$$
T x _ {i} ^ {p} = R x _ {i} ^ {p} + a, \quad 1 \leq i \leq N, \tag {13}
$$
where  $R$  is a  $3 \times 3$  rotation matrix, and  $a \in \mathbb{R}^3$  is a translation vector. We would like the model to satisfy  $p_{\theta}(Tx(t + \tau)|Tx(t)) = p_{\theta}(x(t + \tau)|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 + \tau))$  is acted upon by a random rotation matrix  $R$  to form  $(Rx(t), Rx(t + \tau))$  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 + \tau))\}_{k=1}^{K}$  represents all pairs of states at times  $\tau$  apart in  $\mathcal{D}$ . We optimise:
$$
\mathcal {L} _ {\mathrm {l i k}} (\theta) := \frac {1}{K} \sum_ {k = 1} ^ {K} \log p _ {\theta} \left(x ^ {(k)} (t + \tau) \mid x ^ {(k)} (t)\right). \tag {14}
$$
(ii) acceptance training, the model is fine-tuned to maximise the probability of MH acceptance. Let  $x^{(k)}(t)$  be sampled uniformly from  $\mathcal{D}$ . Then, we use the model to sample  $\tilde{x}_{\theta}^{(k)}(t + \tau) \sim p_{\theta}(\cdot |x^{(k)}(t))$  using Equation (3). We use this to optimise the acceptance probability in Equation (6) with respect to  $\theta$ . Let  $r_{\theta}(X,\tilde{X})$  denote the model-dependent term in the acceptance ratio in Equation (6):
$$
r _ {\theta} (X, \tilde {X}) := \frac {\mu_ {\mathrm {a u g}} (\tilde {X}) p _ {\theta} (X \mid \tilde {X} ^ {p})}{\mu_ {\mathrm {a u g}} (X) p _ {\theta} (\tilde {X} \mid X ^ {p})}. \tag {15}
$$
The acceptance objective is then given by:
$$
\mathcal {L} _ {\mathrm {a c c}} (\theta) := \frac {1}{K} \sum_ {k = 1} ^ {K} \log r _ {\theta} \left(x ^ {(k)} (t), \tilde {x} _ {\theta} ^ {(k)} (t + \tau)\right). \tag {16}
$$
Training to maximise the acceptance probability can lead to the model proposing changes that are too small: if  $\tilde{x}_{\theta}^{(k)}(t + \tau) = 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  $\mathcal{L}_{\mathrm{acc}}(\theta)$ ,  $\mathcal{L}_{\mathrm{lik}}(\theta)$  and a Monte Carlo estimate of the average differential entropy,
$$
\mathcal {L} _ {\text {e n t}} (\theta) := - \frac {1}{K} \sum_ {k = 1} ^ {K} \log p _ {\theta} \left(\tilde {x} _ {\theta} ^ {(k)} (t + \tau) \mid x ^ {(k)} (t)\right). \tag {17}
$$
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
$$
\mathrm {E S S} / \mathrm {s} = \frac {M _ {\text {e f f}}}{t _ {\text {s a m p l i n g}}} = \frac {M}{t _ {\text {s a m p l i n g}} \left(1 + 2 \sum_ {\tau = 1} ^ {\infty} \rho_ {\tau}\right)}, \tag {18}
$$
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  $\varphi$  and  $\psi$ . (c) Ramachandran plots for the conditional distribution of MD compared with the Timewarp model. Red cross denotes initial state. (d) Time dependence of the  $\varphi$  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 1). (f) Free energy comparison for the MD trajectory, Timewarp MCMC (Algorithm 1), and Timewarp exploration (Algorithm 2).
where  $M$  is the chain length,  $M_{\mathrm{eff}}$  is the effective number of samples,  $t_{\mathrm{sampling}}$  is the sampling wall-clock time, and  $\rho_{\tau}$  is the autocorrelation for the lag time  $\tau$  [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\mathrm{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  $\varphi$  and  $\psi$ . 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  $\varphi$  and  $\psi$  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),  $\mu(x(t + \tau)|x(t))$ , compared with the model  $p_{\theta}(x(t + \tau)|x(t))$ , demonstrating close agreement. Finally, Figure 3d shows the time-evolution of the  $\varphi$  angle for MD and Timewarp. Timewarp exhibits significantly more transitions between the metastable states than MD. As a result, the autocorrelation along the  $\varphi$  angle decays much faster in terms of wall-clock time, resulting in a  $\approx 7\times$  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 \mu \mathrm{s} = 2 \times 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
Figure 4: 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 one. 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 1) on test dipeptides (2AA) and tetrapeptides (4AA), respectively. (b), (d) Speed-up for the Timewarp exploration algorithm (Algorithm 2) on test dipeptides (2AA) and tetrapeptides (4AA), respectively.
Figure 5: Experiments on 4AA test tetrapeptides SAEL, CTSA and LPEM (top, middle and bottom rows respectively). Samples were generated via MD, Timewarp exploration (Algorithm 2), and Timewarp MCMC (Algorithm 1). (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 a median speedup factor of  $\approx 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\mu 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 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  $\approx 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  $\tau$  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  $\tau$ , 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.
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.
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.
[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 H Swendsen and Jian-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(21):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-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.
[43] Ashish Vavwani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N 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.
A Symmetries of the architecture
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 + \tau)_{x(t)}$  denote the random variable obtained by sampling  $Z \sim p(z)$  and computing  $X(t + \tau) \coloneqq f(Z; x(t))$ . Here we subscript  $X(t + \tau)_{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(\sigma x(t + \tau)|\sigma x(t)) = p(x(t + \tau)|x(t))$  is equivalent to the following constraint on the random variables:
$$
X (t + \tau) _ {\sigma x (t)} \stackrel {d} {=} \sigma X (t + \tau) _ {x (t)}, \tag {19}
$$
where  $\stackrel{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$ ,
$$
\begin{array}{l} p _ {X (t + \tau) _ {\sigma x (t)}} (x) = p _ {\sigma X (t + \tau) _ {x (t)}} (x) (20) \\ = p _ {X (t + \tau) _ {x (t)}} (\sigma^ {- 1} x), (21) \\ \end{array}
$$
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 \gets \sigma x$ , we get that for all  $x$ ,
$$
p _ {X (t + \tau) _ {\sigma x (t)}} (\sigma x) = p _ {X (t + \tau) _ {x (t)}} (x). \tag {22}
$$
Recalling the notation that  $X(t + \tau)_{x(t)}$  is interpreted as the random variable obtained by conditioning the flow on  $x(t)$ , this can be written as
$$
p (\sigma x | \sigma x (t)) = p (x | x (t)) \tag {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.
$$
\begin{array}{l} X (t + \tau) _ {\sigma x (t)} := f (Z, \sigma x (t)) (24) \\ \stackrel {d} {=} f (\sigma Z, \sigma x (t)) (25) \\ = \sigma f (Z, x (t)) (26) \\ := \sigma X (t + \tau) _ {x (t)}, (27) \\ \end{array}
$$
where in Equation (25) we used the fact that the base distribution  $p(z)$  is  $\sigma$ -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 + \tau),x^{v}(t + \tau)|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  $\overline{x^p}$  denote the average position of the atoms,
$$
\overline {{x ^ {p}}} := \frac {1}{N} \sum_ {i = 1} ^ {N} x _ {i} ^ {p}. \tag {28}
$$
$$
p \left(x ^ {p} (t + \tau), x ^ {v} (t + \tau) \mid x ^ {p} (t)\right) := q \left(x ^ {p} (t + \tau) - \overline {{x ^ {p} (t)}}, x ^ {v} (t + \tau) \mid x ^ {p} (t) - \overline {{x ^ {p} (t)}}\right) \tag {29}
$$
where the subtraction of  $\overline{x^p(t)}$  is broadcasted over all atoms. We now consider the effect of translating both  $x^p(t)$  and  $x^p(t + \tau)$  by the same amount. Let  $a$  be a translation vector in  $\mathbb{R}^3$ . Then
$$
\begin{array}{l} p \left(x ^ {p} (t + \tau) + a, x ^ {v} (t + \tau) \mid x ^ {p} (t) + a\right) (30) \\ = q \left(x ^ {p} (t + \tau) + a - \left(\overline {{x ^ {p} (t)}} + a\right), x ^ {v} (t + \tau) \mid x ^ {p} (t) + a - \left(\overline {{x ^ {p} (t)}} + a\right)\right) (31) \\ = q \left(x ^ {p} (t + \tau) + a - \overline {{x ^ {p} (t)}} - a, x ^ {v} (t + \tau) \mid x ^ {p} (t) + a - \overline {{x ^ {p} (t)}} - a\right) (32) \\ = q \left(x ^ {p} (t + \tau) - \overline {{x ^ {p} (t)}}, x ^ {v} (t + \tau) \mid x ^ {p} (t) - \overline {{x ^ {p} (t)}}\right) (33) \\ = p \left(x ^ {p} (t + \tau), x ^ {v} (t + \tau) \mid x ^ {p} (t)\right). (34) \\ \end{array}
$$
Figure 6: Acceptance probabilities for samples on unseen test peptides with the Timewarp MCMC algorithm. The red line is at  $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(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_{\alpha}$  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.
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