SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Implicit Transfer Operator Learning: Multiple Time-Resolution Surrogates for Molecular Dynamics
Mathias Schreiner* DTU†
Chalmers University of Technology matschreiner@gmail.com
Ole Winther DTU
Simon Olsson†
Chalmers University of Technology Simonols@chalmers.se
Computing properties of molecular systems rely on estimating expectations of the (unnormized) Boltzmann distribution. Molecular dynamics (MD) is a broadly adopted technique to approximate such quantities. However, stable simulations rely on very small integration time-steps  $(10^{-15}\mathrm{s})$ , whereas convergence of some moments, e.g. binding free energy or rates, might rely on sampling processes on time-scales as long as  $10^{-1}\mathrm{s}$ , and these simulations must be repeated for every molecular system independently. Here, we present Implicit Transfer Operator (ITO) Learning, a framework to learn surrogates of the simulation process with multiple time-resolutions. We implement ITO with denoising diffusion probabilistic models with a new SE(3) equivariant architecture and show the resulting models can generate self-consistent stochastic dynamics across multiple time-scales, even when the system is only partially observed. Finally, we present a coarse-grained CG-SE3-ITO model which can quantitatively model all-atom molecular dynamics using only coarse molecular representations. As such, ITO provides an important step towards multiple time- and space-resolution acceleration of MD. Code is available at https://github.com/olsson-group/ito
Numerical simulation of stochastic differential equations (SDE) is critical in the sciences, including statistics, physics, chemistry, and biology applications [1]. Molecular dynamics (MD) simulations are an important example of such simulations [2]. These simulations prescribe a set of mechanistic rules governing the time evolution of a molecular system through numerical integration of, for example, the Langevin equation [3]. MD grants mechanistic insights into experimental observables. These observables are expectations, including time-correlations, of observable functions (e.g., pairwise distances or angles) computed for the Boltzmann distribution  $\hat{\mu} (\mathbf{x})\propto \exp [-\beta U(\mathbf{x})]$  corresponding to the potential  $U(\cdot):\Omega \to \mathbb{R}$  of a  $M$  -particle molecular system,  $\mathbf{x}\in \Omega \subset \mathbb{R}^{3M}$  kept at the inverse temperature  $\beta = 1 / kT$ . However, stable numerical integration relies on time steps,  $\tau$ , which are strictly smaller than the fastest characteristic time
Figure 1: Implicit Transfer Operator: A multiple time-scale surrogate of stochastic molecular dynamics.
*Contributions to this work were done while visiting Chalmers University of Technology
$^{\dagger}$ Technical University of Denmark
scales of the molecular system  $(10^{-15}\mathrm{s},$  e.g., bond vibrations), yet many molecular systems are characterized by processes on much longer time-scales  $(10^{-3} - 10^{-1}\mathrm{s},$  e.g. protein-folding, protein-ligand unbinding, regulation). Consequently, we need infeasibly long simulations to characterize many important processes quantitatively due to the slow mixing in  $\Omega$ .
In this work, we present the implicit Transfer Operator (ITO, Fig.1) as an effective way to learn multiple time-step surrogate models of the stochastic generating distribution of MD. To our knowledge, this is the first surrogate modeling approach that allows for the simultaneous generation of stochastic dynamics at multiple different time resolutions. By adopting an SE(3)-equivariant generative model, we further demonstrate stable long-time-scale dynamics in increasingly difficult settings where an increasing number of degrees of freedom are marginalized. Our approach can be several orders of magnitude more efficient than direct MD simulations and can be made asymptotically unbiased if the generative model permits exact likelihood evaluation. Our current results do not generalize across different thermodynamic ensembles or across chemical space, but show strong generalization across different time-scales.
Our main contributions are
1. the Implicit Transfer Operator (ITO) framework for learning generative models for multiple time resolution molecular dynamics simulations,
2. implementation of ITO using a denoising diffusion probabilistic model (DDPM) [4] with strong empirical results across resolutions: SE(3)-equivariant ITO model (SE3-ITO) gives stable long time-scale simulations and self-consistent dynamics across multiple time-scales for molecular benchmarks and Coarse-grained SE3-ITO model (CG-SE3-ITO) trained on large-scale protein folding data sets shows quantitative agreement with major dynamic and stationary observables of interest.
Notation Throughout this work, diffusion time, related to Diffusion Models (see Sec. 2), and physical time are represented using superscripts and subscripts, respectively.
Molecular dynamics and observables Molecular dynamics (MD) is a wide-spread simulation strategy in computational chemistry and physics. In this approach, the time-evolution of  $N$  particles configuration in Euclidean space  $\mathbf{x} \in \Omega \subset \mathbb{R}^{3M}$ , is modeled via a stochastic differential equation (SDE) with a drift term based on a potential energy model  $U(\mathbf{x}) : \Omega \to \mathbb{R}$ . An important aim of MD is to compute:
1. Stationary observables:  $O_{f} = \mathbb{E}_{\mu}[f(\mathbf{x})]$
2. Dynamic observables:  $O_{f(t)h(t + \Delta t)} = \mathbb{E}_{\mathbf{x}_t\sim \mu}[\mathbb{E}_{\mathbf{x}_{t + \Delta t}\sim p_\tau (\mathbf{x}_{t + \Delta t}|\mathbf{x}_t)}[f(\mathbf{x}_t)h(\mathbf{x}_{t + \Delta t})]]$
where  $\mu$  is the normalized Gibbs measure, and  $p_{\tau}(\mathbf{x}_{t + \Delta t}\mid \mathbf{x}_t)$  is a conditional probability density function encoding the time-discrete evolution of the molecular system  $\mathbf{x}$ , with time-step  $\Delta t = N\tau$  as prescribed by a dynamic model, e.g. Langevin dynamics [3], integrated with time-step,  $\tau$ .  $N$  is typically a large integer. The functions  $f,h:\Omega \to \mathbb{R}$  are observable functions or 'forward models' describing the microscopic observation process, e.g. computing a distance or an angle. The observables,  $O_{f(t)}$ , and  $O_{f(t)h(t + \Delta t)}$ , include binding affinities and binding rates of a drug to a protein, respectively. Conventionally, these observables are estimated from simulation trajectories using naive Monte Carlo estimators.
For illustrative purposes, we assume the temporal behavior of a state,  $\mathbf{x}$ , follows the Brownian dynamics SDE (Itô form)
$$
\mathrm {d} \mathbf {x} _ {t} = - \nabla U (\mathbf {x} _ {t}) \gamma^ {- 1} \mathrm {d} t + \sqrt {2 D} \mathrm {d} W, \tag {1}
$$
where  $D = \gamma^{-1}\beta^{-1}$  is a diffusion constant, with friction  $\gamma$  and inverse-temperature  $\beta$ , and  $\mathrm{d}W$  is a Wiener process. Using the Euler-Maruyama time-discretization, with time-step  $\tau$ , simulating the SDE corresponds to simulating a Markov chain with the transition probability density
$$
p \left(\mathbf {x} _ {t + \tau} \mid \mathbf {x} _ {t}, \tau\right) = \mathcal {N} \left(\mathbf {x} _ {t + \tau} \mid \mathbf {x} _ {t} - \tau \nabla U \left(\mathbf {x} _ {t}\right) \gamma^ {- 1}, \tau \sqrt {2 D} \mathbb {I} _ {3 M}\right) \tag {2}
$$
where  $\mathcal{N}$  specifies the multi-variate Normal distribution, and  $\mathbb{I}_{3M}$  is the  $3M$ -dimensional identity matrix. If  $\tau$  is sufficiently small to allow stable simulation, the invariant measure, of the Markov chain
(eq. 2), is the Boltzmann distribution (normalized Gibbs measure) corresponding to the potential energy model  $U(\mathbf{x})$  at  $\beta$ . Consequently, by simulating a large number of steps we can draw samples from  $\mu$  to compute stationary observables and compute dynamic observables by simulating  $\Delta t = N\tau$  steps enough times with initial states distributed according to  $\mu$ . Explicit simulation makes such computations extremely costly, and consequently, there's much interest in speeding up the calculations of these quantities.
Transfer Operators Let  $\rho$  specify an initial condition, a probability density function on  $\Omega$ . We can define a Markov operator  $T_{\Omega}: L^{1}(\Omega) \to L^{1}(\Omega)$  using a transition density (e.g., [2]):
$$
\left[ T _ {\Omega} \circ \rho \right] \left(\mathbf {x} _ {t + \tau}\right) \triangleq \frac {1}{\mu \left(\mathbf {x} _ {t + \tau}\right)} \int_ {x _ {t}} \mu \left(\mathbf {x} _ {t}\right) \rho \left(\mathbf {x} _ {t}\right) p \left(\mathbf {x} _ {t + \tau} \mid \mathbf {x} _ {t}\right) d \mathbf {x} _ {t} \tag {3}
$$
which then describes the  $\mu$ -weighed evolution of absolutely convergent probability density functions on  $\Omega$  according to eq. [1] with time-step,  $\tau$ . Such an operator is called the (Ruelle) Transfer Operator [5,6]. We can express the operator using a spectral form
$$
T _ {\Omega} (\tau) = \sum_ {i = 0} ^ {\infty} \lambda_ {i} (\tau) \left| \psi_ {i} \right\rangle \left\langle \phi_ {i} \right| \tag {4}
$$
where only eigenvalues  $\lambda_{i}(\tau) = \exp (-\tau \kappa_{i})$  depend on the time-step,  $\tau$ .  $\kappa_{i}$  are characteristic 'relaxation' rates associated with the left and right eigenfunction pair,  $\phi_{i}$  and  $\psi_{i}$  [7]. We can compute the operator with time-lag  $N\tau$  via the Chapman-Kolmogorov equation (see Sec. A.1 for details)
$$
T _ {\Omega} (N \tau) = \sum_ {i = 0} ^ {\infty} \lambda_ {i} (\tau) ^ {N} | \psi_ {i} \rangle \langle \phi_ {i} |. \tag {5}
$$
Equivariant Message Passing Neural Networks In this work, we are concerned with MD, where the time-evolution of a molecule is governed by a force field  $\mathcal{F}(\cdot) \triangleq -\nabla U(\cdot)$  derived from a central potential  $U(\cdot)$ . While  $U(\cdot)$  is invariant to group-actions of the Euclidean group in three dimensions (E(3)), its corresponding force field is E(3)-equivariant. We call a function,  $f$  'invariant' under a group-action  $g$  iff  $f(\mathbf{x}) = f(S_g\mathbf{x})$  and 'equivariant' iff  $T_{g}f(\mathbf{x}) = f(S_{g}\mathbf{x})$ , where  $S_{g}$  and  $T_{g}$  are linear representations of the group element  $g$  [8].
The force field  $\mathcal{F}(\cdot)$  is equivariant under E(3) group-actions. However, in practice, classical molecular dynamics simulations do not change parity during simulation, and consequently, our data distribution only contains a single mirror image of molecules.
We extended the PaiNN architecture [9], an E(3)-equivariant message passing neural network (MPNN), making it SE(3) equivariant by breaking its symmetry with respect to parity. We introduced this minor modification as we experienced sporadic parity changes when sampling with a model trained using the PaiNN architecture, and introducing this modification resolved the issue. Briefly, PaiNN embeds a graph  $G = (V,E)$ , where nodes,  $V$ , exchange equivariant messages through edges within a local neighborhood defined as  $\mathcal{N}(i) = \{j\mid \| r_{ij}\| \leq r_{\mathrm{cutoff}}\}$ , where  $r_{ij}$  is the distance between nodes denoted  $i$  and  $j$ , and  $r_{\mathrm{cutoff}}$  is the maximal distance at which nodes are allowed to exchange messages. Messages are pooled and subsequently used to update node features, thereby enabling exchange of equivariant information. We achieve parity symmetry-breaking by constructing the equivariant messages in a manner that depends on cross-products between equivariant node features and direction vectors between interacting nodes. The cross-product is an axial vector (i.e., does not change sign under parity). We combine these vectors with polar vectors (change sign under parity). We refer to this modified PaiNN architecture as ChiroPaiNN (CPaiNN). Further details are in the Appendix D
Diffusion Models The diffusion model (DM) formalism is a powerful generative modeling framework that learns distributions by modeling a gradual denoising process [4,10,11]. In DMs, we pre-spectify a forward diffusion process (noising process), which gradually transforms the data distribution  $p(\mathbf{x}^0)$  to a simple prior distribution  $p(\mathbf{x}^T)$ , e.g., a standard Gaussian, through a time-inhomogeneous Markov process, described by the following SDE (Itô form)
$$
\mathrm {d} \mathbf {x} ^ {t} = f \left(\mathbf {x} ^ {t}, t\right) \mathrm {d} t + g (t) \mathrm {d} W. \tag {6}
$$
where  $0 < t < T$  is the diffusion time,  $f$  and  $g$  are chosen functions, and  $\mathrm{d}W$  is a Wiener process. We can generate samples from the data distribution  $p(\mathbf{x}^0)$  by sampling from  $p(\mathbf{x}^T)$  and solving the backward diffusion process (denoising process)
$$
\mathrm {d} \mathbf {x} ^ {t} = \left[ f (\mathbf {x} ^ {t}, t) - g ^ {2} (t) \nabla_ {\mathbf {x} ^ {t}} \log p (\mathbf {x} ^ {t} \mid t) \right] \mathrm {d} t + g (t) \mathrm {d} W \tag {7}
$$
by approximating the score field  $\nabla_{\mathbf{x}^t}\log p(\mathbf{x}^t\mid t)$  — or equivalently a time-dependent Normal transition kernel [4] — with a deep neural network surrogate  $\nabla_{\mathbf{x}^t}\log \hat{p} (\mathbf{x}^t\mid t,\boldsymbol {\theta})$  . We can use the learned score field to define a neural ordinary differential equation (ODE) [12, 13], or probability flow ODE [14] — eq. 7 less the term  $g(t)\mathrm{d}W$  and scaling  $g^{2}(t)$  by  $1 / 2$  — which we can leverage for efficient sampling and sample likelihood evaluation.
Here, we are concerned with building equivariant probability density functions under SE(3) group actions. Consequently, we parameterize the DM using a learned Normal transition kernel of a time-inhomogeneous diffusion process. By restricting the transition kernels  $p(\mathbf{x}^{t + 1}\mid \mathbf{x}^t)$  to be equivariant under SE(3) group-actions, the marginal of  $\mathbf{x}^{t + 1}$  is always invariant [15]. Combining the equivariant transition kernel with an invariant prior density [16] ensures the whole Markov process is invariant to SE(3) group actions. Consequently, combining an isotropic mean-free Gaussian as prior with ChiroPaiNN-parameterized transition kernels, we can construct an SE(3) equivariant diffusion model.
Figure 2: ITO  $\hat{\epsilon}$  networks (A) SE3-ITO used for molecular application (B) MB-ITO, used for experiments with the Müller-Brown potential.  $\Lambda_{\mathrm{pos}}$  and  $\Lambda_{\mathrm{nom}}$  are positional and nominal embedding respectively, Concat is a concatenation, and MLP is a multi-layer perceptron. Arrows are annotated with input and output shapes.
3 Implicit Transfer Operator
Molecular simulations are Markovian with transition density (e.g. eq. 2) and Normal, however, the latter only for very small physical time-steps  $\tau$ . Here, we aim to approximate the long-time step transition probability  $p_{N\tau}(\mathbf{x}_{N\tau} \mid \mathbf{x}_0)$  to
allow for one-step sampling of long-time-scale dynamics.
As data we consider simulation trajectories. The trajectories are generated by explicit simulation which corresponds to sampling ancestrally from the small time-step transition density:  $\mathbf{X} = \{\mathbf{x}_{\tau},\dots ,\mathbf{x}_{N\tau}\} \sim p(\mathbf{x}_{n\tau}\mid \mathbf{x}_{(n - 1)\tau})$  , with  $n = \{1,\ldots ,N\}$  . In general, the state variable  $\mathbf{x}$  contains both position and velocity information of the particles, along with other details such as box dimensions, depending on the simulation scheme and target ensemble. Throughout this study, we only consider the position information.
We build a surrogate of the conditional transition probability distribution  $p_{N\tau}(\mathbf{x}_{N\tau} \mid \mathbf{x}_0)$  from MD data. In practice, we learn a generative model  $\mathbf{x}_{t+N\tau} \sim p_{\boldsymbol{\theta}}(\mathbf{x}_{t+N\tau} \mid \mathbf{x}_t, N)$  with a conditional denoising diffusion probabilistic model (cDDPM) of the form
$$
p \left(\mathbf {x} _ {t + N \tau} ^ {0} \mid \mathbf {x} _ {t}, N\right) \triangleq \int p \left(\mathbf {x} _ {t + N \tau} ^ {0: T} \mid \mathbf {x} _ {t}, N\right) d \mathbf {x} ^ {1: T} \tag {8}
$$
where  $\mathbf{x}^{1:T}$  are latent variables of the same dimension as our output, and follow a joint density describing the backward diffusion process (eq. 7) and  $\mathbf{x}^T \sim \mathcal{N}(0,\mathbb{I})$ . We define a conditional sample likelihood as
$$
\ell (\mathbf {I}; \theta) \triangleq \prod_ {i \in \mathbf {I}} p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t _ {i} + N _ {i} \tau} ^ {0} \mid \mathbf {x} _ {t _ {i}}, N _ {i}\right) \tag {9}
$$
where  $\mathbf{I}$  is a list of generated indices  $i$  specifying a time  $t_i$  and a time-lag  $(\tau)$  integer multiple  $N_{i}$ , associating two time-points in the trajectory,  $\mathbf{X}$ . Following Ho et al., we train the cDDPM by
optimizing a simplified form of the variational bound of the log-likelihood [4],
$$
\mathcal {L} (\theta) = \mathbb {E} _ {i \sim \mathbf {I}, \epsilon \sim \mathcal {N} (0, \mathbb {I}), t _ {\text {d i f f}} \sim \mathcal {U} (0, T)} \left[ \| \epsilon - \hat {\epsilon} _ {\boldsymbol {\theta}} \left(\widetilde {\mathbf {x}} _ {t _ {i} + N _ {i} T} ^ {t _ {\text {d i f f}}}, \mathbf {x} _ {t _ {i}}, N _ {i}, t _ {\text {d i f f}}\right) \| _ {2} \right], \tag {10}
$$
where  $\widetilde{\mathbf{x}}_t^{t_{\mathrm{diff}}} \triangleq \sqrt{\bar{\alpha}^{t_{\mathrm{diff}}}} \mathbf{x}_t + \sqrt{1 - \bar{\alpha}^{t_{\mathrm{diff}}}} \epsilon$ , with  $\bar{\alpha}^{t_{\mathrm{diff}}} = \prod_i^{t_{\mathrm{diff}}} (1 - \beta_i)$  and  $\beta_i$  is the variance of the forward diffusion process at diffusion time,  $i$ .  $\hat{\epsilon}_{\pmb{\theta}}(\cdot)$  is one of the two ITO neural network model architectures shown in Fig. 2 and is directly related to the score [4].
If the data used to train the conditional transition density is generated by MD simulation with time-invariant potential energy (drift), we can express the generating transition probability as a decomposition of time-variant and -invariant parts (Proof, see Sec. A.2)
$$
p \left(x _ {N \tau} \mid x _ {0}\right) = \sum_ {i = 1} ^ {\infty} \underbrace {\lambda_ {i} ^ {N} (\tau)} _ {\text {t i m e - v a r i a n t}} \underbrace {\alpha_ {i} \left(x _ {N \tau}\right) \beta_ {i} \left(x _ {0}\right)} _ {\text {t i m e - i n v a r i a n t}} \tag {11}
$$
where  $\alpha_{i}$  and  $\beta_{i}$  are time-invariant projection coefficients of the state variables on-to the left and right eigenfunctions  $\phi_{i}$  and  $\psi_{i}$ , of the Transfer operator  $T_{\Omega}(\tau)$  [5] and  $|\lambda_i(\tau)|\leq 1$  is its  $i$ th eigenvalue. Consequently, we call our surrogate modeling approach implicit transfer operator learning.
As outlined in Algorithm  $\boxed{1}$  we generate the indices  $i$ , e.g. the tuples  $(x_{t_i}, x_{t_i + N_i\tau}, N_i)$ , in a manner such that the model is exposed to multiple time-lags, sampled uniformly across orders of magnitude, used for gradient-based optimization with Adam  $\boxed{17}$ . As a result, as illustrated in eq.  $\boxed{11}$  the model will be exposed to multiple different linear combinations of the eigenfunctions of  $T_{\Omega}(\tau)$  in each batch during training. We conjecture that this data augmentation procedure will enable better learning of implicit representations of these eigenfunctions and, consequently, better generalization across time scales and yield more stable sampling.
3.1 ITO Architectures
We present two architectures for learning cDDPMs encoding ITO models, one for molecular applications SE3-ITO and one for the Müller-Brown benchmark system (Fig. ②). The SE3-ITO architecture uses our new SE(3) equivariant MPNN (ChiroPaiNN, described in sec. ②) to encode  $\mathbf{x}_t$ ,  $N$ , and atom-types,  $z$ , to invariant features,  $s$ , and equivariant features,  $v$ . We concatenate  $s$  with an encoding of the diffusion-time  $t_{\mathrm{diff}}$  and process them through a MLP (multi-layer perceptron). The output from the MLP are passed along with  $v$  and  $\widetilde{\mathbf{x}}_t^{t_{\mathrm{diff}}}$  as input to a second ChiroPaiNN module which predicts  $\hat{\epsilon}$ . More details on the architecture and hyperparameters are available in Appendices D and E.
Algorithm 1 Training. DisExp is defined in Appendix E  
Input: n MD-trajectories;  $\mathcal{X} = \{\mathbf{x}_0^j, \dots, \mathbf{x}_{t_j}^j\}_{j=0}^n$ , ITO score-model;  $\hat{\epsilon}_\theta$ , max lag;  $N_{\max}$ $\mathcal{X}' = \text{Concatenate}(\{\mathbf{x}_0^j, \dots, \mathbf{x}_{t_j - N_{\max}}^j\}_{j=0}^n)$   
while not converged do  
 $\mathbf{x}_t \sim \text{Choice}(\mathcal{X}')$ $N \sim \text{DisExp}(N_{\max})$ $t_{\text{diff}} \sim \text{Uniform}(0, T)$   
Take gradient step on:  
 $\nabla_\theta \left[ \| \epsilon - \hat{\epsilon}_\theta(\widetilde{\mathbf{x}}_{t + N_\tau}^{t_{\text{diff}}}, \mathbf{x}_t, N, t_{\text{diff}}) \|_2 \right]$   
end while  
return  $\hat{\epsilon}_\theta$
Algorithm 2 Ancestral sampling. Sampling from  $p_{\theta}$  is defined in Appendix E, Algorithm 4  
Input: initial condition  $\mathbf{x}_0$ , lag  $N$ , ancestral steps  $n$ .  
Allocate  $\mathcal{T} \in \mathbb{R}^{(n+1) \times \dim(\mathbf{x}_0)}$ $\mathcal{T}[0] = \mathbf{x}_0$   
for  $i = 1 \dots n$  do  
 $\mathbf{x}_i \sim \hat{p}_{\theta}(\mathcal{T}[i-1], N)$ $\mathcal{T}[i] = \mathbf{x}_i$   
end for  
return  $\mathcal{T}$
4 Long time-step stochastic dynamics with Implicit Transfer Operators
4.1 Datasets and test-systems
To evaluate how robustly ITO models can model long time-scale dynamics, we conducted three classes of experiments, ranging from fully observed, high time-resolution, to sparsely observed and
Table 1: VAMP2 score-gaps. Difference in VAMP2-scores of ancestral sampling from ITO models with fixed lag and stochastic lags, compared to baseline Langevin simulations. Perfect match is 0, negative and negative values correspond to under and over estimation of meta-stability, respectively. Standard deviations on last decimal place are given in parentheses.
system\lag	10	100
Müller-Brown (fixed)	-0.0351 (5)	-0.1189 (2)
Müller-Brown (stochastic)	-0.0312 (4)	-0.0970 (5)
low time resolution. Details on training and computational resources are available in Appendices E and F, respectively.
Müller-Brown is a 2D potential commonly used for benchmarking molecular dynamics sampling methods. We generate a training data-set by integrating eq. [1] with the Müller-Brown potential energy as  $U(\mathbf{x})$  (For details, see Appendix B.1). This dataset corresponds to a fully observed case.
Alanine dipeptide We use publicly available data from MDshare [18]. Simulation is performed with 2 fs integration time-steps and data is saved at 1 ps intervals. The simulations are performed in explicit solvation, but we only model the 22 atoms of the solute, without considering velocities. Consequently, this dataset is only partially observed.
Fast-folding proteins We use molecular dynamics data previously reported by Lindorff-Larsen et al. on the fast-folding proteins Chignolin, Trp-Cage, BBA, and Villin [19]. The data is proprietary but available upon request for research purposes. The simulations were performed in explicit solvent with a 2.5 fs time-step and the positions was saved at 200 ps intervals. We coarse-grain the simulation by representing each amino-acid by the Euclidean coordinate of their  $C\alpha$  atom as done previously [20], leading to 10, 20, 28, and 35 particles in each system respectively. Consequently, these data correspond to a mostly unobserved case.
4.2 Stochastic lag improves meta-stability prediction
In sec. [2] we conjecture that exposing an ITO model to multiple lag times during training leads to better and more robust models. To test this, we trained a set of models on the Müller-Brown dataset with fixed constant lags  $N = \{10, 100, 1000\}$  (fixed lag) and a single model with  $N \sim \mathrm{DisExp}(1000)$  (stochastic lag) using the MB-ITO model (Fig. [2]).