SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Leon Klein*†
Andrew Y. K. Foong*
andrewfoong@microsoft.com
Tor Erlend Fjelde*†
tef30@cam.ac.uk
Bruno Mlodozeniec*†
bkm28@cam.ac.uk
Marc Brockschmidt†
Sebastian Nowozin†
Ryota Tomioka
ryoto@microsoft.com
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\mathrm{fs} = 10^{-15}\mathrm{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.
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
$\dagger$  Work done while at Microsoft Research.
Figure 1: (a) Initial state  $x(t)$  (Left) and accepted proposal state  $x(t + \tau) \sim p_{\theta}(x(t + \tau)|x(t))$  (Right) sampled with Timewarp for the dipeptide HT (unseen during training). (b) TICA projections of simulation trajectories, showing transitions between metastable states, for a short MD simulation (Left) and Timewarp MCMC (Right), both run for 30 minutes of wall-clock time. Timewarp MCMC achieves a speed-up factor of  $\approx 33$  over MD in terms of effective sample size per second.
of a molecular system at a temperature  $T$ . Let  $(x^{p}(t), x^{v}(t)) \coloneqq x(t) \in \mathbb{R}^{6N}$  be the state of the molecule at time  $t$ , consisting of the positions  $x^{p}(t) \in \mathbb{R}^{3N}$  and velocities  $x^{v}(t) \in \mathbb{R}^{3N}$  of the  $N$  atoms in Cartesian coordinates. The Boltzmann distribution is given by:
$$
\mu \left(x ^ {p}, x ^ {v}\right) \propto \exp \left(- \frac {1}{k _ {B} T} \left(U \left(x ^ {p}\right) + K \left(x ^ {v}\right)\right)\right), \quad \mu \left(x ^ {p}\right) = \int \mu \left(x ^ {p}, x ^ {v}\right) \mathrm {d} x ^ {v}. \tag {1}
$$
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  $\mu (x^{p})$ . A popular MD method to sample from  $\mu (x^{p})$  is Langevin dynamics [19], which obeys the following stochastic differential equation (SDE):
$$
m _ {i} \mathrm {d} x _ {i} ^ {v} = - \nabla_ {i} U \mathrm {d} t - \gamma m _ {i} x _ {i} ^ {v} \mathrm {d} t + \sqrt {2 m _ {i} \gamma k _ {B} T} \mathrm {d} B _ {t}. \tag {2}
$$
Here  $i$  indexes the atoms,  $m_{i}$  is the mass of atom  $i$ ,  $U(x^{p})$  is the potential energy,  $\gamma$  is a friction parameter, and  $\mathrm{d}B_{t}$  is a standard Brownian motion process. Starting from an initial state  $x(0)$ , simulating Equation (2), along with the relationship  $\mathrm{d}x^{p} = x^{v}\mathrm{d}t$ , yields values of  $x(t)$  that are distributed according to the Boltzmann distribution as  $t\to \infty$ . Standard MD libraries discretise this SDE with a timestep  $\Delta t$ , which must be chosen to be  $\sim 1\mathrm{fs} = 10^{-15}\mathrm{s}$  for stability. Unfortunately, many biomolecules contain metastable states separated by energy barriers that can take milliseconds of MD simulation time ( $\sim 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:
1. We present the first ML algorithm working in general Cartesian coordinates that demonstrates transferability to small peptide systems unseen during training.
2. We demonstrate, for the first time, wall-clock acceleration of asymptotically unbiased MCMC sampling of the Boltzmann distribution of unseen peptide systems.
3. We define an MCMC algorithm targeting the Boltzmann distribution using a conditional normalising flow as a proposal distribution, with a Metropolis-Hastings (MH) correction step to ensure detailed balance (Section 3.4).
4. We design a permutation equivariant, transformer-based normalising flow.
5. We produce a novel training dataset of MD trajectories of thousands of small peptides.
6. We show that even when deployed without the MH correction (Section 3.5), Timewarp can be used to explore metastable states of new peptides much faster than MD.
The code is available here: https://github.com/microsoft/timewarp. The datasets are available upon request<sup>3</sup>.
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 + \tau$ , where  $\tau \gg \Delta 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.
Consider the distribution of  $x(t + \tau)$  induced by an MD simulation of Equation (2) for a time  $\tau \gg \Delta t$ , starting from  $x(t)$ . We denote this conditional distribution by  $\mu(x(t + \tau)|x(t))$ . Timewarp uses a deep probabilistic model to approximate  $\mu(x(t + \tau)|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_{\theta}(x(t + \tau)|x(t))$ , to  $\mu(x(t + \tau)|x(t))$ , where  $\theta$  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_{\theta}(x(t + \tau)|x(t))$  as the density of the distribution defined by the following generative process:
$$
z ^ {p}, z ^ {v} \sim \mathcal {N} (0, I), \quad x ^ {p} (t + \tau), x ^ {v} (t + \tau) := f _ {\theta} \left(z ^ {p}, z ^ {v}; x ^ {p} (t), x ^ {v} (t)\right). \tag {3}
$$
3Please contact andrewfoong@microsoft.com for dataset access.
Here  $z^p \in \mathbb{R}^{3N}$  and  $z^v \in \mathbb{R}^{3N}$ . For all settings of  $\theta$  and  $x(t)$ ,  $f_{\theta}(\cdot ;x(t))$  is a diffeomorphism that takes the latent variables  $(z^p,z^v) \in \mathbb{R}^{6N}$  to  $(x^{p}(t + \tau),x^{v}(t + \tau)) \in \mathbb{R}^{6N}$ . 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 + \tau)$ , only the map from  $z$  to  $x(t + \tau)$  must be invertible. Using the change of variables formula, we can evaluate:
$$
p _ {\theta} (x (t + \tau) | x (t)) = \mathcal {N} \left(f _ {\theta} ^ {- 1} (x (t + \tau); x (t)); 0, I\right) \left| \det  \mathcal {J} _ {f _ {\theta} ^ {- 1} (\cdot ; x (t))} (x (t + \tau)) \right|,
$$
where  $f_{\theta}^{-1}(\cdot ;x(t)):\mathbb{R}^{6N}\to \mathbb{R}^{6N}$  is the inverse of the diffeomorphism  $f_{\theta}(\cdot ;x(t))$ , and  $\mathcal{I}_{f_{\theta}^{-1}(\cdot ;x(t))}(x(t + \tau))$  denotes the Jacobian of  $f_{\theta}^{-1}(\cdot ;x(t))$  evaluated at  $x(t + \tau)$ .
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  $\mathcal{D} = \{\mathcal{T}_i\}_{i=1}^P$ , where  $P$  is the number of peptides. Each MD trajectory is temporally subsampled with a spacing of  $\tau$ , so that  $\mathcal{T}_i = (x(0), x(\tau), x(2\tau), \ldots)$ . During training, we randomly sample pairs  $x(t), x(t + \tau)$  from  $\mathcal{D}$ . Each pair represents a sample from the conditional distribution  $\mu(x(t + \tau) | 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 + \tau)$  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  $\mathcal{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) \sim \mathcal{N}(0, I)$ . The auxiliary variables  $x^v(t)$  now contain no information about the future state  $x^p(t + \tau)$ ,  $x^v(t + \tau)$ , since  $x^v(t)$  and  $x^v(t + \tau)$  are drawn independently. Hence we can simplify  $f_\theta$  to depend only on  $x^p(t)$ , with  $x^p(t + \tau)$ ,  $x^v(t + \tau) := f_\theta(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_{\theta}(x(t + \tau)|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:
$$
\mu_ {\mathrm {a u g}} \left(x ^ {p}, x ^ {v}\right) \propto \exp \left(- \frac {U \left(x ^ {p}\right)}{k _ {B} T}\right) \mathcal {N} \left(x ^ {v}; 0, I\right). \tag {4}
$$
Starting from an initial state  $X_0 = (X_0^p, X_0^v) \in \mathbb{R}^{6N}$  for state  $m = 0$ , we iterate:
$$
\tilde {X} _ {m} \sim p _ {\theta} (\cdot | X _ {m} ^ {p}), \quad X _ {m + 1} := \left\{ \begin{array}{l l} \tilde {X} _ {m} & \text {w i t h p r o b a b i l i t y} \alpha \left(X _ {m}, \tilde {X} _ {m}\right) \\ X _ {m} & \text {w i t h p r o b a b i l i t y} 1 - \alpha \left(X _ {m}, \tilde {X} _ {m}\right) \end{array} \right. \tag {5}
$$
where  $\alpha(X_m, \tilde{X}_m)$  is the Metropolis-Hastings (MH) acceptance ratio [24] targeting Equation (4):
$$
\alpha (X, \tilde {X}) = \min  \left(1, \frac {\mu_ {\operatorname {a u g}} (\tilde {X}) p _ {\theta} (X \mid \tilde {X} ^ {p})}{\mu_ {\operatorname {a u g}} (X) p _ {\theta} (\tilde {X} \mid X ^ {p})}\right) \tag {6}
$$
The flow used for  $p_{\theta}$  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:
$$
\left(X _ {m} ^ {p}, X _ {m} ^ {v}\right) \leftarrow \left(X _ {m} ^ {p}, \epsilon\right), \quad \epsilon \sim \mathcal {N} (0, I). \tag {7}
$$
Iterating these updates yields a sample  $X_{m}^{p}, X_{m}^{v} \sim \mu_{\mathrm{aug}}$  as  $m \to \infty$ . To obtain a Boltzmann-distributed sample of the positions  $X_{m}^{p} \sim \mu$ , we simply discard the auxiliary variables  $X_{m}^{v}$ . As sending  $m \to \infty$  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_{\theta}(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  $\ell$ th coupling layer of the flow, the following transformations are implemented:
$$
z _ {\ell + 1} ^ {p} = s _ {\ell , \theta} ^ {p} \left(z _ {\ell} ^ {v}; x ^ {p} (t)\right) \odot z _ {\ell} ^ {p} + t _ {\ell , \theta} ^ {p} \left(z _ {\ell} ^ {v}; x ^ {p} (t)\right), \tag {8}
$$
$$
z _ {\ell + 1} ^ {v} = s _ {\ell , \theta} ^ {v} \left(z _ {\ell + 1} ^ {p}; x ^ {p} (t)\right) \odot z _ {\ell} ^ {v} + t _ {\ell , \theta} ^ {v} \left(z _ {\ell + 1} ^ {p}; x ^ {p} (t)\right). \tag {9}
$$
Going forward, we suppress the coupling layer index  $\ell$ . Here  $\odot$  is the element-wise product, and  $s_{\theta}^{p}: \mathbb{R}^{3N} \to \mathbb{R}^{3N}$  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_{\theta}^{p}: \mathbb{R}^{3N} \to \mathbb{R}^{3N}$  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_{\theta}$  consists of  $N_{\mathrm{coupling}}$  stacked coupling layers, beginning from  $z \sim \mathcal{N}(0, I)$  and ending with a sample from  $p_{\theta}(x(t + \tau)|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 + \tau) - x(t)$ , rather than  $x(t + \tau)$  directly.
Atom transformer We now describe the atom transformer network. Let  $x_{i}^{p}(t), z_{i}^{p}, z_{i}^{v}$ , all elements of  $\mathbb{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_{\theta}^{p}(z^{v}, x^{p}(t))$  and  $t_{\theta}^{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_{i}^{p}(t), h_{i}, z_{i}^{v}]^{\top} \in \mathbb{R}^{H + 6}$ , where  $z_{i}^{p}$  has been excluded since  $s_{\theta}^{p}, t_{\theta}^{p}$  are not allowed to depend on  $z^{p}$ . Here  $h_{i} \in \mathbb{R}^{H}$  is a learned embedding vector which depends only on the atom type. The vectors  $a_{i}$  are fed into an MLP  $\phi_{\mathrm{in}}: \mathbb{R}^{H + 6} \to \mathbb{R}^{D}$ , where  $D$  is the feature dimension of the transformer. The vectors  $\phi_{\mathrm{in}}(a_{1}), \ldots, \phi_{\mathrm{in}}(a_{N})$  are then fed into  $N_{\mathrm{transformer}}$  stacked transformer layers. After the transformer layers, they are passed through another atom-wise MLP,  $\phi_{\mathrm{out}}: \mathbb{R}^{D} \to \mathbb{R}^{3}$ . The final output is in  $\mathbb{R}^{3N}$  as required. This is depicted in Figure 2, Middle. When implementing  $s_{\theta}^{v}$  and  $t_{\theta}^{v}$  from Equation (9), a similar procedure is performed on the vector  $[x_{i}^{p}(t), h_{i}, z_{i}^{p}]^{\top}$ , but now including  $z_{i}^{p}$  and excluding  $z_{i}^{v}$ . There are two key differences between the atom transformer and the architecture in [43]. First, to maintain permutation equivariance,
Figure 2: Schematic illustration of the Timewarp conditional flow architecture, described in Section 4. Left: A single conditional RealNVP coupling layer. Middle: A single atom transformer module. Right: the multihead kernel self-attention module.
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  $\tau$  that are not too large, the positions at time  $t + \tau$  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:
$$
w _ {i j} = \frac {\exp \left(- \| x _ {i} ^ {p} - x _ {j} ^ {p} \| _ {2} ^ {2} / \ell^ {2}\right)}{\sum_ {j ^ {\prime} = 1} ^ {N} \exp \left(- \| x _ {i} ^ {p} - x _ {j ^ {\prime}} ^ {p} \| _ {2} ^ {2} / \ell^ {2}\right)}, \tag {10}
$$
where  $\ell$  is a lengthscale parameter. The outputs  $\{r_{\mathrm{out},i}\}_{i = 1}^{N}$ , given the inputs  $\{r_{\mathrm{in},i}\}_{i = 1}^{N}$ , are then:
$$
r _ {\text {o u t}, i} = \sum_ {j = 1} ^ {N} w _ {i j} V \cdot r _ {\text {i n}, j}, \tag {11}
$$
where  $V \in \mathbb{R}^{d_{\mathrm{out}} \times d_{\mathrm{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  $\sigma$  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 + \tau)$  is to permute the atoms similarly, i.e.,
$$
\mu (\sigma x (t + \tau) | \sigma x (t)) = \mu (x (t + \tau) | x (t)). \tag {12}
$$
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: