SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The model was trained to generate samples from the conditional Boltzmann distribution  $\mu (x(t + \tau)|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  $\tau$  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,
Figure 7: Experiments on 2AA test dipeptides QW (top row), HT (middle row) and GP (bottom row). Comparison of the long MD trajectory and Timewarp MCMC (Algorithm 1). (a) TICA plots. (b) Free energy comparison of the first two TICA components. (c) Autocorrelation for the TIC 0 component. (d) Time dependence of the TIC 0 component.
Figure 8: Experiments for the 2AA test dipeptides DH (first row), GT (second row), TK (third row), and CW (last row). Comparison of the long MD trajectory and Timewarp MCMC (Algorithm 1). (a) TICA plots. (b) Free energy comparison of the first two TICA components. (c) Autocorrelation for the TIC 0 component. (d) Time dependence of the TIC 0 component.
Figure 9: Experiments on 4AA test tetrapeptides AWCK, LYVI and CSFQ (top, middle and bottom rows respectively). Samples were generated via MD, Timewarp exploration (Algorithm 2), and Timewarp MCMC (Algorithm 1). (a) TICA plots of samples. (b) Free energies along the first two TICA components. (c) Potential energy distribution. For AWCK all metastable states are found by all methods, for LYVI the MD trajectory misses one state, and for CSFQ Timewarp MCMC misses the slowest transition. In all cases Timewarp exploration discovers all metastable states.
Figure 10: Comparing the conditional Boltzmann distributions generated with MD trajectories and the Timewarp model for alanine dipeptide. (a) Ramachandran plots for the conditional distributions compared with the equilibrium Boltzmann distribution. The red cross indicates the conditioning state. This is similar to the plot shown in Figure 3c, but here showing a different conditioning state. The match between the conditional distributions is not as close here as it is for Figure 3c, which could be because here the conditioning state is chosen to be in the less likely metastable state. (b) Projections on the first two dihehedral angles for the conditional distributions. (c) Potential energies of the conditional distributions. (d) Autocorrelations for samples generated according to the MCMC algorithm (Algorithm 1) compared with a long MD trajectory. Note that this autocorrelation plot is not for the conditional distribution, but corresponds to the results shown in Figure 3.
emphasising the importance of the Metropolis-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 autocorrelation 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\mu 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
Figure 11: Comparing the conditional distribution of the Timewarp model  $p_{\theta}(x(t + \tau)|x(t))$  with the conditional distribution generated with MD  $\mu (x(t + \tau)|x(t))$ . The rows correspond to the peptides TK, AAEW, EASS, respectively, where we show for each peptide two different conditioning states. (a) TICA plots. First column: samples from the Boltzmann distribution  $\mu (x)$  generated with MD. Second column: samples from  $\mu (x(t + \tau)|x(t))$  generated with MD. The conditioning state  $x(t)$  is indicated with the red cross. Third column: samples from  $p_{\theta}(x(t + \tau)|x(t))$  generated with the Timewarp conditional flow, without MH correction. The conditioning state  $x(t)$  is indicated with the red cross. (b) Projection of the conditional distributions from (a) onto the first two TICA components. (c) Potential energy distributions of the conditional distributions. (d) Conditional bondlength distribution, which for these values of  $\tau$  will be close to the equilibrium distribution. Each mode in the graph represents a different bond type, e.g., C-H.
stable, we run several shorter MD trajectories  $(0.5\times 10^{6}\mathrm{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\times 10^{6}\mathrm{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:
Figure 12: Validation of new metastable states found with Timewarp exploration. The red crosses indicate the different conditioning states. (a) TICA plots for (not conditional) samples generated via a long MD trajectory, Timewarp exploration (Algorithm 2), and Timewarp MCMC (Algorithm 1). Timewarp exploration discovers some metastable states unseen in the long MD trajectories (bottom of LYVI plot, bottom left of CTSA plot). (b) Conditional distributions generated with MD or with the Timewarp model, starting from states visited by the long MD trajectory shown in (a). Some short MD trajectories now discover the new metastable states, verifying that they are indeed valid states. Only the final state of each short MD trajectory is recorded. (c) Conditional distributions generated with MD or with the Timewarp model, starting from the new metastable states discovered by one of the short MD trajectories shown in (b).
Algorithm 1 Timewarp MH-corrected MCMC with batched proposals
Require: Initial state  $X_0 = (X_0^p, X_0^v)$ , chain length  $M$ , proposal batch size  $B$ .  
 $m \gets 0$   
while  $m < M$  do  
Sample  $\tilde{X}_1, \ldots, \tilde{X}_B \sim p_{\theta}(\cdot | X_m^p)$  {Batch sample}  
for  $b = 1, \ldots, B$  do  
 $\epsilon \sim \mathcal{N}(0, I)$  {Resample auxiliary variables}  
 $X_b \gets (X_m^p, \epsilon)$   
Sample  $I_b \sim \text{Bernoulli}(\alpha(X_b, \tilde{X}_b))$   
end for  
if  $S := \{b : I_b = 1, 1 \leq b \leq B\} \neq \emptyset$  then  
 $a = \min(S)$  {First accepted sample}  
 $(X_{m+1}^p, \ldots, X_{m+a-1}^p) \gets (X_m^p, \ldots, X_m^p)$ $X_{m+a}^p \gets \tilde{X}_a^p$ $m \gets m + a$   
else  
 $(X_{m+1}^p, \ldots, X_{m+B}^p) \gets (X_m^p, \ldots, X_m^p)$ $m \gets m + B$   
end if  
end while  
output  $X_0^p, \ldots, X_M^p$
Algorithm 2 Fast, biased exploration of the state space with Timewarp
Require: Initial state  $X_0^p$ , number of steps  $M$ , maximum allowed energy increase  $\Delta U_{\mathrm{max}}$  for  $m = 0, \ldots, M$  do Sample  $\tilde{X}_m^p \sim p_\theta(\cdot \mid X_m^p)$  {Sample from conditional flow} if  $U(\tilde{X}_m^p) - U(X_m^p) < \Delta U_{\mathrm{max}}$  then  $X_{m+1}^p \gets \tilde{X}_m^p$  else  $X_{m+1}^p \gets X_m^p$  {Reject if energy change is too high} end if end for output  $X_0^p, \ldots, X_M^p$
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  $\Delta U_{\mathrm{max}} = 300\mathrm{kJ / mol}$ .
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.
Table 1: OpenMM MD simulation parameters
Force Field	amber-14
Time step	0.5fs
Friction coefficient	0.3 1/ps
Temperature	310K
Integrator	LangevinMiddleIntegrator
We present more dataset details, like simulation times and number of peptides, in Table 2.
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  $\ell_{i} = \{0.1, 0.2, 0.5, 0.7, 1.0, 1.2\}$ , given in nanometers.
Table 2: Dataset details
Dataset name	AD	2AA	4AA
Training set simulation time	100 ns	50 ns	50 ns
Test set simulation time	100 ns	1 μs	1 μs
MD integration step Δt	0.5 fs	0.5 fs	0.5 fs
Timewarp prediction time τ	0.5 × 106fs	0.5 × 106fs	0.5 × 105fs
No. of training peptides	1	200	1400
No. of training pairs per peptide	2 × 105	1 × 104	1 × 104
No. of test peptides	1	100	30
Table 3: Timewarp model hyperparameters
Dataset	RealNVP layers	Transformer layers	Parameters	Atom-embedding dim H
AD	12	6	1 × 108	64
2AA	12	6	1 × 108	64
4AA	16	16	4 × 108	128
The  $\phi_{\mathrm{in}}$  and  $\phi_{\mathrm{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.
Table 4: Timewarp MLP layer sizes
Dataset	φin MLP	φout MLP	Transformer MLP
AD	[70, 256, 128]	[128, 256, 3]	[128, 256, 128]
2AA	[70, 256, 128]	[128, 256, 3]	[128, 256, 128]
4AA	[134, 2048, 128]	[128, 2048, 3]	[128, 2048, 128]
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
Table 5: Timewarp loss weighting factors
Dataset	Llik(θ)	Lacc(θ)	Lent(θ)
AD	0.99	0.01	0.1
2AA	0.9	0.1	0.1
4AA	1	0	0
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 \times 10^{-4}$ ,
Table 6: Timewarp training parameters
Dataset + training method	Batch size	No. of A-100s	Training time
AD — likelihood	256	1	1 week
AD — acceptance	64	1	2 days
2AA — likelihood	256	4	2 weeks
2AA — acceptance	256	4	4 days
4AA — likelihood	256	4	3 weeks
the learning rate is then consecutively decreased by a factor of 2 upon hitting training loss plateaus.
G Computing infrastructure
The training was performed 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.