SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We find that the model trained with a stochastic lag systematically outperforms models trained with fixed lag (Table I). We gauge the agreement by comparing Variational Approach to Markov Processes (VAMP)-2 scores [21] (for details, see Appendix G), between model samples and training data and find that both models tend to underestimate meta-stability compared to training data slightly. However, the model trained with stochastic lag is marginally closer to the reference values. We note that the difference in the ability of fixed and stochastic lag ITO models to capture long-time-scale dynamics is also reflected in the learned transition densities (Fig. 5). Together, these results suggest that lag-time augmentation during training leads to better implicit learning of the Transfer operator's eigenfunctions than training with a fixed lag.
To test whether this phenomena extends in cases where we do not have full observability and to molecular systems we followed the VAMP2-gaps of alanine di-peptide as a function of epoch for models trained with a fixed lag and a stochastic lag (Fig. 6). We find that the VAMP-2 gaps for stochastic lag and fixed lag in this case are statistically indistinguishable across all epochs. These results suggest that we can without compromising on accuracy build multiple time-scale surrogates by training with stochastic lag-times.
4.3 Efficient and accurate self-consistent long time-scale dynamics
We evaluate the ITO models trained with stochastic lags to capture long time-scale dynamics in a self-consistent manner, in the Chapman-Kolmogorov sense, i.e.,  $p(\mathbf{x}_{\Delta t} \mid \mathbf{x}_0) \triangleq p(\mathbf{x}_{N\tau} \mid \mathbf{x}_0) = \prod_{i=1}^{N} p(\mathbf{x}_{i\tau} \mid \mathbf{x}_{(i-1)\tau})$ , or if samples generated by direct sampling with time-step  $\Delta t = N\tau$  are
distributed similarly to samples generated by performing ancestral sampling  $N$  times, each with timestep  $\tau$ . In direct sampling, we draw samples for a desired time-step  $\Delta t = N\tau$  from  $\hat{p}_{\theta}(\mathbf{x}_0,N)$  and in ancestral sampling we draw  $n$  samples with time-step  $\Delta t = N\tau$  from  $\hat{p}_{\theta}(\mathbf{x}_0,N)$  in an ancestral manner, e.g.  $x_{i + 1}\sim \hat{p}_{\theta}(\mathbf{x}_i,N)$ , where  $i = \{0,\dots n - 1\}$ .
For the fully-observed Müller-Brown case, we find that the ITO model is self-consistent by the strong overlap in transition densities sampled in a direct and ancestral manner (Algorithm 2). These results generalize to molecular systems and partially observed systems. Sampling an SE3-ITO model (Fig. 2) trained with alanine dipeptide data, we find strong agreement between the ancestrally and directly sampled transition densities (Fig. 3) and we again have a strong consistency with corresponding transition densities computed from molecular dynamics simulations. Note here, that the time-step of the ITO-sampled transition densities varies from  $10^{4}$  to  $10^{6}$  times the MD integration time-step. The transition densities for alanine di-peptide (Fig. 3) are calculated using 15'000 trajectories. For direct sampling, this means that we draw 15'000 samples in total. In the case of ancestral sampling, we sample 4, 64, or 512 steps for 15'000 trajectories, to match  $\Delta t$ .
Next, we consider four fast-folding proteins [19] where we coarse-grain the proteins by representing them only with their  $C\alpha$  atoms. In this sparsely observed case (CG-SE3-ITO), we find strong model self-consistency, as shown by the comparison between conditional densities from the folded and unfolded states (Fig. 4) projected onto a linear subspace determined using time-lagged independent component analysis (time-lagged independent components, tIC) [22] (see Appendix B.3). Further, the long timescale transition density gradually converges to the data distribution as expected.
Finally, by ancestral sampling (Algorithm 2), we perform a simulation of Chignolin with the same length as the training trajectory  $(106\mu s)$  using a CG-SE3-ITO model, and compare with MD. The CG-SE3-ITO simulation is 2120 steps with  $\Delta t = 5$  ns. Running in parallel, on a single Titan X GPU we can simulate the CG-SE3-ITO model at a rate of  $363\mathrm{ns} / (\mathrm{s}_wM^2)$  where  $\mathrm{s}_w$  denotes seconds wall-time (Appendix C.1). Remarkably, these trajectories are virtually indistinguishable in the slowly
relaxing TICA coordinates, illustrating stability of ITO. These conclusions extend to the proteins Villin, BBA, and Trp-Cage (See Appendix, Figs. 89 and 10)
Together these results suggest that ITO models accurately and self-consistently capture the slow dynamics of molecular systems and are robust to situations where the system is only partially observed. In general, we expect robustness to sparsely observed representations as long as the input representations are sufficient to span the eigenfunctions of  $T_{\Omega}$  [23, 24]. Approximation errors will translate into systematic under-estimation of relaxation time-scales [7], consistent with our slight under-estimation of VAMP-2 scores (Table 1). In future work, combining the learning of SE3-ITO models with a systematic scheme for coarse-graining [25, 26], could be an avenue for scaling to large-scale molecular systems at a low computational cost.
Figure 3: Transition probability densities of alanine dipeptide dynamics with SE3-ITO model; Rows of increasing time-lag (from top to bottom). Contours are samples from SE3-ITO model, and 2D histograms show estimates from MD data. The first column shows conditional transition densities projected onto the torsion angles  $\phi$  and  $\psi$  (inset). The black cross indicates the initial condition. The second and third columns show marginal distributions of  $\phi$  and  $\psi$ , respectively, with direct sampling in orange, ancestral sampling in blue, and MD data in black.
Figure 4: Reversible protein folding-unfolding of Chignolin with CG-SE3-ITO Conditional probability densities (orange contours) starting from unfolded (upper panels) and folded (lower panels) protein states, at increasing time-lag (left to right), shown on top of data distribution. Below: time-traces of 106 microsecond MD simulations and ITO simulations on tICs 1 and 2. The two dashed lines correspond to the folded state value in tIC 1 (lower line) and tIC 2 (higher line). Contour lines are based on  $10^{\prime}000$  trajectories, generated with ancestral sampling with the length,  $\Delta t$  and time-step  $200\mathrm{ps}$ . For  $\Delta t = 200\mathrm{ns}$  this corresponds to 1000 ancestral samples.
5 Prediction of dynamic and stationary observables of using CG-SE3-ITO
As outlined in section 2, an important aim of MD simulations is to compute stationary and dynamic observables, which involves intractable integrals typically approximated via Monte Carlo estimators. Using the trained ITO models we can efficiently sample i.i.d. from the transition density needed for computing dynamic observables, and by choosing a time-step which is sufficiently large we can also sample i.i.d from the Boltzmann distribution  $\mu$ , the latter akin to Boltzmann generators [27] (See Appendix A.1). We note that, the ITO models are surrogates and as such without reweighing we cannot expect unbiased samples from the Boltzmann and dynamic transition densities. Nevertheless, we gauge how accurately ITO models we can compute these observables of interest in the context of protein folding without reweighing:
- Free Energy of Folding,  $\Delta G = -\log \left[\frac{p_{\mathrm{f}}}{1 - p_{\mathrm{f}}}\right]$
- Mean first passage time, folding,  $\langle \tau_f\rangle = \int_{x_0\in \neg f}\int_0^\infty \delta (x_t\in f)p(x_t\mid x_0,t)t  \mathrm{d}t  \mathrm{d}x_0$
- Mean first passage time, unfolding,  $\langle \tau_u\rangle = \int_{x_0\in f}\int_0^\infty \delta (x_t\in \neg f)p(x_t\mid x_0,t)t  \mathrm{d}t  \mathrm{d}x_0$
where  $\{f, \neg f\} \subset \Omega$  are disjoint subsets corresponding to the folded and unfolded states of a protein,  $p_f = \int_{x \in f} \mu(x) \, \mathrm{d}x$ , is the folded state probability and  $\delta(\cdot)$  is the Dirac delta.
We compute these observables using the reference molecular simulation data [19] and sample statistics from the CG-SE3-ITO models of each of the four fast-folding proteins (details in Appendix B.6). Strikingly, the observables computed using CG-SE3-ITO models agree well with those computed from long all-atom MD simulations (Table 2).
Finally, we analyzed the robustness, convergence, and consistency of these observables (Fig. 7). For Chignolin, we trained five models independently and analyzed model checkpoints when the training loss had stabilized. For each checkpoint and each model, we computed the observables. The values predicted are statistically indistinguishable, suggesting consistency, robustness, and convergence. The
Table 2: Molecular observables Standard deviations on last decimal place are given in parentheses. Stationary and dynamic observables are denoted s and d, respectively.
ΔGfold/kT (s)	⟨τf⟩/μs (d)	⟨τu⟩/μs (d)
Chignolin (MD/ITO)	-1.28(1)/-0.64(33)	0.565(4)/1.02(24)	2.01(2)/2.12(34)
Trp-Cage (MD/ITO)	1.47(6)/2.84(6)	13.6(4)/37(2)	3.4(2)/2.85(9)
BBA (MD/ITO)	0.97(3)/1.52(3)	11.7(2)/8.6(2)	5.1(1)/1.75(4)
Villin (MD/ITO)	1.21(2)/2.22(3)	2.41(3)/3.27(7)	0.68(1)/0.354(5)
average predictions closely match the reference values. Nevertheless, we note that the fluctuations in these values are noticeable.
We implemented all experiments using PyTorch [28], PyTorch Lightning [29], JAX [30], and used DPM-Solver [31] for probability flow ODE Sampling.
Molecular sampling Sampling molecular configurations is a broad field and can broadly be divided into two main areas: physically motivated sampling of the Boltzmann distribution and conformer generation. The first area includes algorithmic approaches to sample the Boltzmann distribution including Molecular Dynamics simulations [2], Markov Chain Monte Carlo, extended ensemble methods [32, 33, 34, 35], including analysis methods involving deep generative nets [36], and surrogate models which directly approximate the Boltzmann distribution and allow for recovery of unbiased statistics, including Boltzmann generators [37, 16, 38, 39]. Conformer generation concerns generating physically plausible conformers without explicitly trying to follow the Boltzmann distribution. The latter approaches can be split into ML [15, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49] and chemoinformatic [50, 51] approaches. Finally, speeding up molecular simulations by reducing the effective number of particles to simulate through coarse-graining with special purpose forcefield models [52] including machine learned variants [53, 54, 20, 55] and learned coarse-graining maps [25, 26] is an orthogonal approach to sample conformation space. Further, several methods to recover all-atom models from coarse-grained representations through ML [56, 57] and rule-based approaches [58] are available.
Transfer Operator surrogates Building transfer operator surrogates is commonly used in molecular modeling including (Deep Generative) Markov state models (MSM) [59, 7, 60, 61], also including experimental data [62, 63] dynamic graphical models, [64] VAMPnets [65, 21], observable operator models [66], however, primarily for analysis of molecular dynamics data. Markov state models are time-space discrete approximations of the transfer operator and Deep Generative MSM [67] and VAMPnets [65] are deep learning infused versions, where state discretization is learned by deep nets. Dynamic graphical models reparameterize MSMs as kinetic Markov random fields allowing for scaling to larger systems [64]. Klein et al. recently introduced timewarp which is a flow-based generative model to simulate molecular systems with a large (up to  $0.5\mathrm{ns}$ ), fixed, time-lag, [68] providing asymptotically unbiased equilibrium samples through a Metropolis-Hastings correction [69]. While timewarp generates conformers with realistic local structure, it has limitations in capturing long time-scale dynamics, which is reflected in the predicted transition probability densities. In contrast, our approach captures long time-scale dynamics efficiently allowing for accurate prediction of dynamic observables. However, currently, neither the code nor the data from timewarp is publicly available percluding direct comparisons on the benchmark tasks established in their paper.
**Surrogate model** Implicit Transfer Operators are surrogate models of stochastic dynamics' conditional transition probability densities. We cannot guarantee unbiased sampling of dynamics and the stationary distribution due to aleatoric (e.g., finite data) and epistemic (e.g., model misspecification) uncertainty. We can overcome the latter by reweighing against a Markov Chain Monte Carlo acceptance criterion as proposed previously [68], to ensure unbiased dynamics path-reweighing is necessary, which in turn requires closed-form expressions for the target path probabilities [70].
Transferability and scalability Currently, ITO does not generalize across chemical space and thermodynamic variables. In future work, we anticipate that generalization across chemical space limitations can be overcome by appropriate data set curation and parameter-sharing schemes [71]. Generalization across thermodynamic variables such as temperature and pressure would require using a surrogate model which is steerable under these changes, e.g., temperature steerable flows [72] or thermodynamic maps [73]. Currently, we assume a fully connected graph that scales  $\mathcal{O}(M^2)$  in system size, which limits what systems are practically accessible. Devising new surrogate models which use mean-field approximation approaches from e.g., computational physics [74] or chemistry to truncate the graphs and treat long-range as an additive term [75] could yield more favorable scaling [76].
This paper introduces Implicit Transfer Operators (ITO), an approach to building multiple time-scale surrogate models of stochastic molecular dynamics. We implement ITO models with a conditional DDPM using a new time-augmentation scheme and show how ITO models capture fast and slow dynamics on benchmarks and molecular systems. We show ITO models are self-consistent over multiple time scales and highly robust to the marginalization of degrees of freedom in the system, which are unimportant to capture the long-time-scale dynamics. Combined with a SE(3) variant of the PaiNN architecture [9] (ChiroPaiNN), we further show strong empirical evidence of scaling to molecular applications, such as the folding of coarse-grained proteins. As such, we are confident that ITO is a stepping-stone toward general-purpose surrogates of molecular dynamics.
This work was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation and The Novo Nordisk Foundation (SURE, NNF19OC0057822). OW's work was funded in part by the Novo Nordisk Foundation through the Center for Basic Machine Learning Research in Life Science (NNF20OC0062606) and by the Pioneer Centre for AI, DNRF grant number P1. Preliminary results were enabled by resources provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS) at Alvis (project: NAISS 2023/22-392) partially funded by the Swedish Research Council through grant agreement no. 2022-06725.
[1] Bernt Øksendal. Stochastic Differential Equations. Springer Berlin Heidelberg, 2003. DOI: 10.1007/978-3-642-14394-6. URL: https://doi.org/10.1007/978-3-642-14394-6
[2] Daan Frenkel and Berend Smit. Understanding Molecular Simulation. Elsevier, 2002. DOI: 10.1016/b978-0-12-267351-1.x5000-7. URL: https://doi.org/10.1016/b978-0-12-267351-1.x5000-7.
[3] Paul Langevin. "Sur la théorie du mouvement brownien". In: C. R. Acad. Sci. (Paris) 146 (1908), 530-533.
[4] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. 2020. DOI: 10.48550/ARXIV.2006.11239. URL: https://arxiv.org/abs/2006.11239
[5] David Ruelle. Thermodynamic Formalism. en. Encyclopedia of mathematics and its applications. Harlow, England: Longman Higher Education, Nov. 1978.
[6] Christof Schütte et al. “Conformation dynamics”. In: Proc. Int. Congr. ICIAM (2009), pp. 297-336.
[7] Jan-Hendrik Prinz et al. "Markov models of molecular kinetics: Generation and validation". In: The Journal of Chemical Physics 134.17 (May 2011), p. 174105. DOI: 10.1063/1.3565032. URL: https://doi.org/10.1063/1.3565032.
[8] Jean-Pierre Serre. Linear Representations of Finite Groups. Springer New York, 1977. DOI: 10.1007/978-1-4684-9458-7. URL: https://doi.org/10.1007/978-1-4684-9458-7.