Boltzmann priors for Implicit Transfer Operators

## Abstract

Abstract Accurate prediction of thermodynamic properties is essential in drug discovery and materials science. Molecular dynamics (MD) simulations provide a principled approach to this task, yet they typically rely on prohibitively long sequential simulations. Implicit Transfer Operator (ITO) Learning offers a promising approach to address this limitation by enabling stable simulation with time steps orders of magnitude larger than MD. However, to train ITOs, we need extensive, unbiased MD data, limiting the scope of this framework. Here, we introduce Boltzmann Priors for ITO (BoPITO) to enhance ITO learning in two ways. First, BoPITO enables more efficient data generation, and second, it embeds inductive biases for long-term dynamical behavior, simultaneously improving sample efficiency by one order of magnitude and guaranteeing asymptotically unbiased equilibrium statistics. Further, we showcase the use of BoPITO in a new tunable sampling protocol interpolating ITO models trained on off-equilibrium simulation data and an unbiased equilibrium distribution to solve inverse problems in molecular science.

## 1 Introduction

Efficient molecular dynamics (MD) simulation on long time-scales is critical to a large number of scientific and engineering applications. Stable simulations rely on solvers taking tiny integration time-steps, making the simulation of most phenomena impractical with current methods. Since these simulations are stochastic, a simulation step corresponds to drawing a sample from a transition probability density $p(\mathbf{x}_{t+\tau}\mid\mathbf{x}_{t})$, where $\tau$ is a tiny time-step. Recently, deep generative models have emerged as a promising strategy to potentially speed up these simulations by learning transition probability densities where $\tau$ is much larger (Schreiner et al., ; Klein et al., ; Hsu et al., ; Fu et al., ) and thereby allow efficient sampling at long time-scales. Implicit Transfer Operator (ITO) learning (Schreiner et al., ) learns such surrogate models at multiple time-resolutions. While ITO has shown promise in accelerating simulations, it relies on extensive unbiased simulation data which may not always be available. In Markov state models, this limitation can be mitigated by integrating off-equilibrium simulations and enhanced sampling simulations (Trendelkamp-Schroer et al., ; Rosta & Hummer, ; Wu et al., ; ). However, such estimators are so far unavailable for deep generative surrogates of the transition density.

Figure: Figure 1: Boltzmann Priors for Implicit Transfer Operators (BoPITO) leverage pre-trained Boltzmann Generators to enable data-efficient training of surrogate models of the transition density after $N$ simulation steps. BoPITO is implemented using Score-Based Diffusion Models. The score of the model, $s_{\mathbf{\theta}}$, separates the contributions from the first eigenfunction of the transfer operator (equilibrium density), $s_{\text{eq}}$, from the rest, $s_{\text{dyn}}$. BoPITO embeds inductive biases for long-term dynamical behavior and enables interpolation between off-equilibrium and equilibrium models.
Refer to caption: x1.png

Here, we introduce Boltzmann Priors for Implicit Transfer Operator (BoPITO) learning (Figure 1). BoPITO leverages pre-trained Boltzmann Generators (BG) (Noé et al., ; Viguera Diez et al., ; Klein & Noé, ; Köhler et al., ; ; Midgley et al., ; ) to enable data-efficient training of ITO models, leading to one order of magnitude reduction in the simulation data needed for training. As the BG encodes the invariant measure or Boltzmann distribution of the dynamics encoded by the transition density, BoPITO, by construction, guarantees asymptotically unbiased equilibrium statistics. In this way, BoPITO can combine off-equilibrium data and biased data encoded into a BG prior to train ITO models that predict MD across multiple time-scales. Using BoPITO we introduce a new sampling strategy to recover approximate dynamics from biased off-equilibrium data, a BG prior, and unbiased time-correlation data, providing a new method for inverse problems for molecular systems.

Our main contributions are

- 1.
Boltzmann Priors for Implicit Transfer Operator Learning (BoPITO): We provide a principled way to leverage Boltzmann Generators as priors for Implicit Transfer Operator Learning, boosting sample efficiency, ensuring the generation of uncorrelated data and allowing asymptotically unbiased equilibrium statistics for $\tau\rightarrow\infty$.
- 2.
BoPITO interpolators: A tunable sampling protocol facilitating interpolation of BoPITO models trained on off-equilibrium simulation data to an unbiased equilibrium distribution.
- 3.
We show that BoPITO interpolators can recover approximate dynamics from models trained on biased simulations. We frame the optimization of interpolation parameters as an inverse problem and select the interpolated ensemble that is most consistent with unbiased observables.

## 2 Background and preliminaries

### 2.1 Molecular dynamics and observables

Molecular dynamics (MD) is a widely used simulation method in chemistry, physics, and biology. It combines a mathematical model of the dynamics — e.g. Langevin dynamics (Langevin, ) — with a potential energy model, $U(\mathbf{x}):\Omega\rightarrow\mathbb{R}$, of a system of interest, in turn providing detailed mechanistic insights of molecular systems, through the time evolution of particles in configuration space, $\mathbf{x}\in\Omega$. Practically, MD is solved by numerical integration, and time is discretized, with step $\tau$. Consequently, we can understand MD as a Markov process with a Normal transition density $p(\mathbf{x}_{t+\tau}\mid\mathbf{x}_{t})$, whose associated Markov operator has the Boltzmann distribution,

$$ $\displaystyle\mu(\mathbf{x})=\mathcal{Z}^{-1}\exp(-\beta U(\mathbf{x})),\text{ with }\mathcal{Z}=\int\mathrm{d}\mathbf{x}\,\exp(-\beta U(\mathbf{x})),$ (1) $$

as its invariant measure or stationary distribution, where $\beta$ is the inverse temperature.

One important application of MD is to compute expectations or observables (Olsson, ):

- 1.
Stationary observables:
$\displaystyle O_{a}=\mathbb{E}_{\mu}\left[a(\mathbf{x})\right].$
(2)
- 2.
Dynamic observables / Time correlation functions:
$\displaystyle O_{a(t),b(t+N\tau)}=\mathbb{E}_{\mathbf{x}_{t}\sim\mu}\left[
\mathbb{E}_{\mathbf{x}_{t+N\tau}\sim p_{\tau}(\mathbf{x}_{t+N\tau}|\mathbf{x}_
{t})}\left[a(\mathbf{x}_{t})\cdot b(\mathbf{x}_{t+N\tau})\right]\right],$
(3)

where $p_{\tau}(\mathbf{x}_{t+N\tau}|\mathbf{x}_{t})$ is the conditional probability density after $N$ simulation steps with time-step $\tau$. The maps $a$ and $b:\Omega\rightarrow\mathbb{R}^{L}$ serve as observable functions or ‘forward models’ characterizing microscopic observation processes, e.g. indicating whether a drug is bound or not, or an interatomic distance, leading to observables such binding affinities and off-rates of a drug to a target protein.

Unfortunately, stable MD simulations rely on using time discretizations, $\tau$, on the order of $10^{-15}\,\mathrm{s}$, whereas properties such as protein folding or ligand unbinding occur on much longer time scales, $\sim 10^{-3}-10^{-1}\,\mathrm{s}$. Due to the sequential nature of these simulations, an impractical number of simulation steps are necessary to predict these properties in an unbiased fashion. As a result, most MD data will be ‘off-equilibrium’, e.g. trajectories exploring one or a few of the modes of $\mu(\mathbf{x})$.

### 2.2 Boltzmann Generators

Boltzmann Generators (BG) (Noé et al., ) are a generative machine learning approach to draw samples i.i.d. from the Boltzmann distribution, $\mu(\mathbf{x})$, of physical many-body systems (eq. 1). BGs learn a diffeomorphic map, $\mathcal{F}_{\bm{\theta}}:\mathbb{R}^{N}\rightarrow\mathbb{R}^{N}$, from a latent space, equipped with a simple base density $p(\mathbf{z})$, to configurational space of a physical system, such that the push-forward density $\hat{p}(\mathbf{x})=\mathcal{F}_{\bm{\theta}}\#p(\mathbf{z})$ closely approximates the Boltzmann distribution (eq. 1). In practice, $\mathcal{F}_{\bm{\theta}}$ is implemented using an invertible neural network architecture with tractable Jacobian determinants to enable efficient sampling and exact sample likelihood computation (Chen et al., ; Papamakarios et al., ). BGs are trained either with approximately equilibrated simulation data or with biased simulations, from i.e. enhanced sampling, by employing appropriate reweighting (Ferrenberg & Swendsen, ; Shirts & Chodera, ). Unbiased samples can then be generated by importance re-sampling. Similarly, unbiased expectations can be computed using the importance weights $w(\mathbf{x})=e^{-\beta U(\mathbf{x})}/\hat{p}(\mathbf{x})$. As such, BGs are surrogates of the target Boltzmann distribution, but do not model time-correlation statistics.

### 2.3 Transfer operators

Transfer operators (Ruelle, ; Schütte et al., ) provide a framework to describe the evolution of probability densities over time. Let $p$ denote an initial probability density function on $\Omega$, and $\rho$ its $\mu$-weighted version, $p=\mu\rho$. The Markov operator, $T_{\Omega}$, is defined using a transition density $p(\mathbf{x}_{t+\tau}|\mathbf{x}_{t})$,

$$ $\displaystyle[T_{\Omega}(\tau)\circ\rho](\mathbf{x}_{t+\tau})\equiv\frac{1}{ \mu(\mathbf{x}_{t+\tau})}\int_{\Omega}\mu(\mathbf{x}_{t})\rho(\mathbf{x}_{t})p (\mathbf{x}_{t+\tau}|\mathbf{x}_{t})\,\mathrm{d}\mathbf{x}_{t}.$ (4) $$

This operator describes the $\mu$-weighted evolution of absolutely convergent probability density functions on $\Omega$ by a discrete-time increment $\tau$ given by the dynamics encoded in a transition density $p(\mathbf{x}_{t+\tau}|\mathbf{x}_{t})$. In the context of molecular dynamics, $T_{\Omega}$ is a $\mu$-weighted equivalent to the Markov operator discussed above (sec. 2.1). The spectral form of the transfer operator is expressed as

$$ $\displaystyle[T_{\Omega}(\tau)\circ\rho](\mathbf{x}_{t+\tau})=\sum_{i=1}^{ \infty}\lambda_{i}(\tau)\langle\rho|\phi_{i}\rangle\ \psi_{i}(\mathbf{x}_{t+ \tau}),$ (5) $$

where $\lambda_{i}(\tau)=\exp(-\tau\kappa_{i})$ are the eigenvalues, $\langle f|g\rangle=\int_{\Omega}f(\mathbf{x})g(\mathbf{x})\,\mathrm{d}\mathbf{x}$ and $\psi_{i}$ and $\phi_{i}$ are the corresponding right and left eigenfunctions, respectively. The eigenvalues $\lambda_{i}(\tau)$ depend on $\tau$, and $\kappa_{i}$ are characteristic relaxation rates. When the dynamics are reversible, all eigenvalues $\lambda_{i}$ are real-valued and lie in the interval $-1<\lambda_{i}\leq 1$. Notably, there is one eigenvalue $\lambda_{1}=1$ and its associated eigenfunction corresponds to the stationary distribution $\mu(\mathbf{x})$. Moreover, the left and right eigenfunction are related by a factor of the stationary density, $\phi_{i}(\mathbf{x})=\mu(\mathbf{x})\psi_{i}(\mathbf{x})$, and $\psi_{1}(\mathbf{x})=\mathbf{1}$.

### 2.4 Diffusion models

Diffusion models (DMs) (Ho et al., ; Song et al., ) are a popular generative modeling framework that approximates data densities $p(\mathbf{x}^{0})$ by learning to invert a noising process (‘forward diffusion’). The forward diffusion process is pre-specified and incrementally transforms the data distribution into a simple prior distribution, $p(\mathbf{x}^{T})$, through simulation of a time-inhomogenous Markov process represented by the SDE,

$$ $\mathrm{d}\mathbf{x}^{t}=f(\mathbf{x}^{t},t)\,\mathrm{d}t+g(t)\,\mathrm{d}W,$ (6) $$

where ($0<t<T$) is the diffusion time, $f$ and $g$ are chosen functions, and $\mathrm{d}W$ is a Wiener process. To generate samples from the data distribution, $p(x^{0})$, we can sample from $p(\mathbf{x}^{T})$ and solve the backward diffusion process (‘denoising process’) (Anderson, ),

$$ $\mathrm{d}\mathbf{x}^{t}=\left[f(\mathbf{x}^{t},t)-g^{2}(t)\nabla_{\mathbf{x}^ {t}}\log p(\mathbf{x}^{t}|t)\right]\mathrm{d}t+g(t)\,\mathrm{d}W.$ (7) $$

This backward process is approximated using a deep neural network $\nabla_{\mathbf{x}^{t}}\log p_{\bm{\theta}}(\mathbf{x}^{t}|t)=s_{\bm{\theta}}(
\mathbf{x}^{t},t)$. The learned backward SDE can be recast into a probability flow ODE (Maoutsa et al., ; Song et al., ), which in turn can be interpreted as a continuous normalizing flow (Chen et al., ) which facilitates efficient sampling and enables sample likelihood evaluation.

### 2.5 Implicit Transfer Operator Learning

Implicit Transfer Operator (ITO) Learning (Schreiner et al., ) is a framework for learning surrogate models of the transition density, $p(\mathbf{x}_{N\tau}|\mathbf{x}_{0})$, from MD data, where $N$ is an arbitrarily large integer. ITO leverages that the transfer operator framework allows us to express the transition density as

$$ $p(\mathbf{x}_{t+N\tau}|\mathbf{x}_{t})=\sum_{i=1}^{\infty}\lambda_{i}^{N}(\tau )\psi_{i}(\mathbf{x}_{t})\phi_{i}(\mathbf{x}_{t+N\tau}),$ (8) $$

where $\psi_{i}$ and $\phi_{i}$ are independent of $\tau$ and $N$. See Appendix A.1 for a complete derivation.

This decomposition inspired a strategy to learn a conditional generative model $\mathbf{x}_{t+N\tau}\sim p_{\bm{\theta}}(\mathbf{x}_{t+N\tau}|\mathbf{x}_{t},N)$ by sampling tuples $(\mathbf{x}_{t_{i}},\mathbf{x}_{t_{i}+N_{i}\tau},N_{i})$ and training a generative model to mimic the empirical transition density at multiple time-horizons, $N\tau$. Here, we approximate ITO models with a conditional denoising diffusion probabilistic model (cDDPM) of the form

$$ $\displaystyle p(\mathbf{x}^{0}_{t+N\tau}|\mathbf{x}_{t},N)\equiv\int p(\mathbf {x}^{0:T}_{t+N\tau}|\mathbf{x}_{t},N)\,\mathrm{d}\mathbf{x}^{1:T},$ (9) $$

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 denote diffusion time in Diffusion Models using superscripts, while physical time is represented using subscripts. The conditional sample likelihood is given by

$$ $\ell(\mathbf{I};{\bm{\theta}})\equiv\prod_{i\in\mathbf{I}}p_{\bm{\theta}}( \mathbf{x}^{0}_{t_{i}+N_{i}\tau}\mid\mathbf{x}_{t_{i}},N_{i})$ (10) $$

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 a MD trajectory of length $M\tau$, $\mathbf{x}=\{\mathbf{x}_{0},\mathbf{x}_{\tau},\dots,\mathbf{x}_{(M-1)\tau}\}$.

Training is performed by optimizing an approximation of the variational bound of the log-likelihood (Ho et al., ),

$$ $\mathcal{L}({\bm{\theta}})=\mathbb{E}_{i\sim\mathbf{I},\epsilon\sim\mathcal{N} (0,\mathbb{I}),t_{\mathrm{diff}}\sim\mathcal{U}(0,T)}\left[\lVert\epsilon-\hat {\epsilon}_{\bm{\theta}}(\widetilde{\mathbf{x}}^{t_{\mathrm{diff}}}_{t_{i}+N_{ i}\tau},\mathbf{x}_{t_{i}},N_{i},t_{\mathrm{diff}})\rVert_{2}\right],$ (11) $$

where $\widetilde{\mathbf{x}}^{t_{\mathrm{diff}}}_{t}=\sqrt{\bar{\alpha}^{t_{\mathrm{
diff}}}}\mathbf{x}_{t}+\sqrt{1-\bar{\alpha}^{t_{\mathrm{diff}}}}\epsilon$, with $\bar{\alpha}^{t_{\mathrm{diff}}}=\prod_{j}^{t_{\mathrm{diff}}}(1-\beta_{j})$ and $\beta_{j}$ is the variance of the forward diffusion process at diffusion time, $j$. $\hat{\epsilon}_{\bm{\theta}}(\cdot)$ is one of the two architectures presented by Schreiner et al. (). Following Arts et al. () we express the score as

$$ $s_{\bm{\theta}}(\mathbf{x}^{t_{\text{diff}}}_{t+N\tau},\mathbf{x}_{t},N_{i},t_ {\text{diff}})=-\frac{\hat{\epsilon}_{\bm{\theta}}(\widetilde{\mathbf{x}}^{t_{ \mathrm{diff}}}_{t_{i}+N_{i}\tau},\mathbf{x}_{t_{i}},N_{i},t_{\mathrm{diff}})} {\sqrt{1-\bar{\alpha}^{t_{\mathrm{diff}}}}}.$ (12) $$

## 3 Boltzmann Priors for Implicit Transfer Operator Learning

ITO training requires extensive, unbiased MD simulations to capture the statistical distribution of rare events. This data-intensive requirement can hinder ITO’s practical implementation. Furthermore, models trained on off-equilibrium simulations, which may exhibit non-representative statistics, can lead to inaccurate predictions, compromising their utility in downstream applications. In practice, however, both unbiased off-equilibrium simulations and information about the equilibrium distribution from potentially biased simulations (Hénin et al., ) can be cheaply generated, and both encode information about the dynamical behavior of molecular systems. We introduce Boltzmann Priors for Implicit Transfer Operator (BoPITO), as a learning paradigm for ITO models that leverage available information about the equilibrium distribution.

BoPITO uses pre-trained models of the equilibrium distribution to improve ITO learning in three ways. First, it helps ensure broad sampling of $\Omega$ proportional to the equilibrium distribution for subsequent MD simulations yielding information about the transition density $p(\mathbf{x}_{\tau}\mid\mathbf{x}_{0})$ from across $\Omega$. Second, we use it to fix the stationary part of the learned transition density, boosting the sample efficiency when learning models of molecular dynamics. Third, it imposes an inductive bias of long-time dynamical behavior allowing for recovery of unbiased Boltzmann distribution for long-time horizons. Fourth, using BoPITO, we introduce a novel tunable sampling protocol interpolating ITO models trained on off-equilibrium simulation data and an unbiased equilibrium distribution.

### 3.1 Efficient data generation

Training data for ITO models consists of several independent unbiased MD simulations. Following the adaptive sampling strategy (Bowman et al., ; Doerr & De Fabritiis, ; Viguera Diez et al., ; Betz & Dror, ), used extensively in the molecular dynamics simulation community, we use a pre-trained BG to generate initial conditions to simulations ensuring broad sampling across $\Omega$ proportional to the Boltzmann distribution. As long as these trajectories reach a ‘local equilibrium’ we can in principle recover an unbiased model of the molecular dynamics (Nüske et al., ), albeit without relying on running one or a few very long simulations to reach the global equilibrium.

### 3.2 Long-term dynamics inductive bias for ITO

One important application of ITO is to allow for one-step sampling of long-time-scale dynamics. However, real datasets often contain a limited number of effective samples for long time-scales, leading to potential biases in models. To mitigate this issue, we propose separating the equilibrium contribution from the time-dependent components:

$$ $p(\mathbf{x}_{t+N\tau}|\mathbf{x}_{t})=\mu(\mathbf{x}_{t+N\tau})+\sum_{i=2}^{ \infty}\lambda_{i}^{N}(\tau)\phi_{i}(\mathbf{x}_{t+N\tau})\psi_{i}(\mathbf{x}_ {t}),$ (13) $$

where we have used that $\lambda_{1}=1$ and its corresponding eigenfunctions are $\phi_{1}=\mu$ and $\psi_{1}=\mathbf{1}$. Additionally, we introduce a decay in the time-dependent component, creating an inductive bias that asymptotically samples from an available equilibrium model for long-term dynamics. Using the spectral decomposition of the transition density (eq. 13), we choose the score model as

$$ $\displaystyle s_{\bm{\theta}}(\mathbf{x}_{t+N\tau}^{t_{\text{diff}}},\mathbf{x }_{t},N,t_{\text{diff}})=s_{\mathrm{eq}}(\mathbf{x}_{t+N\tau}^{t_{\text{diff}} },t_{\text{diff}})+\hat{\lambda}^{N}s_{\mathrm{dyn}}(\mathbf{x}_{t+N\tau}^{t_{ \text{diff}}},\mathbf{x}_{t},N,t_{\text{diff}},{\bm{\theta}}),$ (14) $$

where $s_{\mathrm{eq}}(\mathbf{x}^{t}_{\text{diff}},t_{\text{diff}})$ is the score of a pre-trained surrogate of the equilibrium distribution model, $0<\hat{\lambda}<1$ is a hyper-parameter and $s_{\mathrm{dyn}}(\mathbf{x}_{t+N\tau}^{t_{\text{diff}}},\mathbf{x}_{t},N,t_{
\text{diff}},{\bm{\theta}})$ accounts for the time-dependent components. As $N\rightarrow\infty$, $s_{\mathrm{eq}}$ dominates and the model samples from the equilibrium model, see Appendix A.2 for an example. By interpreting the score field as a velocity field, we can reformulate the DM as a continuous normalizing flow which we can use for Metropolized sampling of the unbiased equilibrium distribution (Klein et al., ). The proposed factorization is principled, corresponding to a separation of the score of the transition density into a stationary and dynamic part, where the first part is considered known, see Appendix A.3 for details. In practice, we first train $s_{\mathrm{eq}}$ (if not provided) using equilibrium data. Then we train $s_{\mathrm{dyn}}$ with unbiased, possibly off-equilibrium, MD data while keeping $s_{\mathrm{eq}}$ fixed.

### 3.3 BoPITO interpolators

A significant issue in MD simulations is the unbiased sampling of transitions over high free-energy barriers, e.g. channels in $\mu(\mathbf{x})$ with very little probability mass, typically coinciding with mixing on $\Omega$. This difficulty leads to an inability to predict time-correlation statistics for large $N\tau$. With BoPITO we can define an interpolation between models trained on off-equilibrium simulations and the equilibrium distribution. In this manner, we can approximate dynamics not seen explicitly in the data. For a given maximum training lag, $N_{\text{max}}$, we define a BoPITO interpolator as a model with the following score function:

$$ $\displaystyle s_{\bm{\theta}}(\mathbf{x}^{t_{\text{diff}}},\mathbf{x}_{t},N_{ \text{int}},t_{\text{diff}})=s_{\mathrm{eq}}(\mathbf{x}^{t_{\text{diff}}},t_{ \text{diff}})+\hat{\lambda}^{N_{\text{int}}}s_{\mathrm{dyn}}(\mathbf{x}^{t_{ \text{diff}}},\mathbf{x}_{t},N_{\text{max}},t_{\text{diff}}),$ (15) $$

with $N_{\text{int}}\geq N_{\text{max}}$. This defines an interpolator because $N_{\text{int}}=N_{\text{max}}$ generates samples from the model distribution with maximum lag and $N_{\text{int}}\rightarrow\infty$ generates samples from the equilibrium model.

Inspired by methods in molecular biophysics (Kolloff & Olsson, ; Olsson et al., ; Bottaro & Lindorff-Larsen, ; Salvi et al., ), we propose to choose the interpolation parameter, $N_{\text{int}}$, as the most consistent with an unbiased dynamical observable, such as experimental data. That is, for a lag $N>N_{\text{max}}$ and an unbiased dynamic observable $O^{*}_{N}$, we choose

$$ $\displaystyle N_{\text{int},N}=\operatorname*{arg\,min}_{N_{\text{int}}}|O^{*} _{N}-O_{N_{\text{int}}}|,$ (16) $$

where $O_{N_{\text{int}}}$ is the dynamic observable estimated with samples from the BoPITO interpolator. This way BoPITO can integrate off-equilibrium and equilibrium MD with experimental data. In practice, we find the interpolator to generate some high-energy states. However, we can alleviate the high-energy structures by alternating long-lag interpolation steps with short-lag non-interpolation steps for local relaxation.

## 4 Results

For detailed parameters of the experiments below, we refer to Appendix A.4.

### 4.1 Systems

Prinz potential is a 1D potential commonly used for benchmarking MD sampling methods (Prinz et al., ). We set the observable functions, $a$ and $b$, eq. 3, to be the identity function for computing dynamic observables. For details, see Appendix A.5.

Alanine Dipeptide is a small peptide with 22 atoms. We use publicly available data from Dibak et al. (), containing $1\,\mu\mathrm{s}$ simulation time split in $20$ trajectories. Simulation is performed in an implicit solvent with $2\,\mathrm{fs}$ integration time-step, and data is saved every $1\,\mathrm{ps}$. We choose,

$$ $\displaystyle a(\mathbf{x})=b(\mathbf{x})=\begin{bmatrix}\sin{\phi}(\mathbf{x} )\\ \cos{\phi}(\mathbf{x})\end{bmatrix},$ (17) $$

where $\phi(\mathbf{x})$ is a torsion angle involved in the slowest transition in the system. For details, see Appendix A.6.

### 4.2 Boltzmann priors for training data generation

Modeling the transition density over $\Omega$ requires observing transitions across the state space. In the usual setting, only one or a few initial conditions are available when data collection starts. Here, we explore the case where a BG is available before data collection and compare it against the baseline, where only one initial condition is known.

We use a BG trained on equilibrium data of the Prinz potential to sample the initial conditions of our training trajectories. We also generate trajectories using a single starting point (crystal, $\mathbf{x}=0.75$) and compare their performance with long simulations (Figure 2). We compare the two different data generation strategies using the metric, $|\Delta\text{correlation}|$, which measures the absolute difference of the time-correlation function (dynamic observable) compared to the MD ground truth (Appendix A.7). We find that the performance of models using Boltzmann priors for data generation is superior for different lag times (Figure 2a) and number of trajectories (Figure 2b). The gap between BG and the crystal baseline increases with lag-time, and as expected decreases with the number of generated trajectories.

Figure: (a)
Refer to caption: x2.png

### 4.3 BoPITO efficiently samples long-term dynamics in a low-data context

Fixing the equilibrium contribution to the score field effectively reduces the number of parameters that need to be estimated. To test whether this prior information manifests as an improved sample efficiency we compared ITO and BoPITO models against each other with varying sizes of training data.

We find that the BoPITO models achieve a higher accuracy for long-term dynamics compared to ITO models when data is scarce, as the equilibrium distribution is known a priori and does not need to be learned from simulation data (Figure 3). The inductive bias in BoPITO models enables them to learn long-term dynamics, even in scenarios where ITO models fail. For the Prinz Potential, we find that while ITO suffers from poor performance modeling long-term dynamics when data is scarce (Figure 3a), BoPITO models (Figure 3b) accurately capture long-term dynamics without worsening the performance on short and medium timescales. The results for Alanine Dipeptide mirror those of the Prinz potential: we find that an order of magnitude more data is needed to train an ITO (Figure 3c) to the same accuracy as a BoPITO model trained on the same data (Figure 3d). Moreover, we show in Appendix A.8 that the energies of configurations sampled from our prior Boltzmann Generator used in our experiments match closely the energies of MD samples.

Figure: Figure 3: Absolute difference in correlation (lower is better, one-step sampling) between ITO (left) and BoPITO (right) split into short, medium, and long time-scales against the number of training trajectories for the Prinz potential (top) and Alanine Dipeptide (bottom). The shaded areas correspond to a $95\ \%$ confidence interval.
Refer to caption: x4.png

### 4.4 Interpolating between models trained on off-equilibrium data and the Boltzmann distribution with experimental data

In practical settings, our MD data will be off-equilibrium, e.g. having sampled only one or a few of the relevant modes in the Boltzmann distribution in a given trajectory. Consequently, unless extensive data across the domain $\Omega$ can be collected, models based on such data will be biased. An alternative to collecting more simulation data is to use a multi-modal strategy where experimental data is used to fill the gaps left by simulation data and bridge to long time-scale dynamics (Salvi et al., ; Kolloff & Olsson, ). In this section we explore the potential of BoPITO interpolators to integrate dynamic observables with off-equilibrium simulation data to recover unbiased long-term dynamics.

Figure: Figure 4: Fitting an interpolator of a biased BoPITO model. Both the ITO and BoPITO models are trained on biased MD simulations of Alanine Dipeptide. We fit a BoPITO interpolator selecting the most consistent ensemble with the unbiased correlation function specified in Equation 16. Numbers with blue arrows specify the interpolation parameter, $N_{\text{int}}$.
Refer to caption: x5.png

We showcase BoPITO interpolators on Alanine Dipeptide. By removing the transitions between the modes of the Boltzmann distribution corresponding to the slowest process in the system, e.g. $\phi$ crossing 0 0 or $2$, we generate a biased simulation data, resembling a realistic scenario. We then train ITO and BoPITO models using these biased trajectories. For lags $>100$, we sample the BoPITO interpolator estimated by matching to an unbiased dynamic observable defined by eq. 22, allowing us to overcome a systematic error in the correlation function observed in the biased MD data and for ITO (Figure 4).

Beyond reproducing the provided dynamic observable, the interpolator also demonstrates a remarkable ability to capture the underlying microscopic dynamics (Figure 5). Even with excellent agreement, the generated ensembles only slightly overestimates density in the transition state region. However, this effect could easily be alleviated by annealing a short MD simulation to the interpolation as was recently shown (Viguera Diez et al., ).

## 5 Related work

##### Sampling the Boltzmann distribution

Apart from Boltzmann Generator-based approaches, there are a number of traditional ways to draw statistical samples from the Boltzmann distribution of molecular systems. Prominent examples include, molecular dynamics or Markov Chain Monte Carlo simulations coupled with enhanced sampling strategies (Hénin et al., ; Kamenik et al., ), including replica-based approaches (Earl & Deem, ; Sidler et al., ; Pasarkar et al., ), conformational flooding (Grubmüller, ), meta-dynamics (Laio & Parrinello, ), and umbrella sampling (Torrie & Valleau, ), in particular when paired up with machine-learned collective variables (Chen & Ferguson, ; Wang et al., ; Herringer et al., ; M. Sultan & Pande, ). The success of these approaches relies on choosing the right mechanism to enhance sampling across high free energy barriers, or low probability regions, for any given case (Carter et al., ). Finding such mechanisms typically involves substantial manual engineering of collective variables and other hyper-parameters. Finally, there are transition path sampling approaches (Dellago et al., ; Bolhuis et al., ), which are particularly powerful when combined with reinforcement learning (Jung et al., ) or deep generative priors (Plainer et al., ).

##### Latent space simulators and Coarse graining

Classical (fine-grained) molecular dynamics simulation can be coarse-grained beyond the Born-Oppenheimer approximation, such that atomic nuclei are merged together into ‘beads.’(Noid, ) This strategy, in principle enables faster simulations due to the smaller number of particles and the acceleration of kinetics caused by the coarse-graining (CG) operation (Nüske et al., ). CG models can be estimated to closely approximate the thermodynamics of the corresponding fine-grained system through the optimization of two equivalent variational bounds (Noid et al., ; Lyubartsev & Laaksonen, ; Ercolessi & Adams, ). A relaxation of this bound was recently proposed to also allow for probabilistic reconstruction of fine-grained configurations (Chennakesavalu et al., ). These bounds have been used extensively to build CG force-field models (Husic et al., ; Wang et al., ; Majewski et al., ; Charron et al., ) including implicit solvation models (Chen et al., ; Katzberger & Riniker, ), using deep neural networks due to their ability to capture multibody terms (Wang et al., ) to ultimately accelerate the prediction of equilibrium properties of molecular systems. Similarly, the development of ‘latent space simulators’ where a learned, typically low-dimensional, latent space equipped either with a propagator (Sidky et al., ; Chennakesavalu et al., ; Wang et al., ), or not (Wang & Gómez-Bombarelli, ), is learned to enable efficient simulation. These approaches, in general aim to accelerate molecular simulations akin to BoPITO, yet due to the CG operation, the molecular dynamics (kinetics) will be accelerated, and detailed knowledge of the unbiased dynamics are needed to correct this (Nüske et al., ; Crommelin & Vanden-Eijnden, ).

Figure: Figure 5: BoPITO can incorporate unbiased dynamic observables to correct a model trained on biased data. Rows of increasing time-lag (from top to bottom). Contour plots correspond to a BoPITO interpolator. 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.
Refer to caption: x6.png

##### Transfer Operator surrogates

Analysis of MD data often involves building transfer operator surrogates such as Markov state models (MSM) (Schütte et al., ; Prinz et al., ; Swope et al., ; Husic & Pande, ), time-lagged independent components analysis (Molgedey & Schuster, ; Ziehe & Müller, ; Pérez-Hernández et al., ), Markov field models or dynamic graphical models, (Olsson & Noé, ; Mardt et al., ; Hempel et al., ), VAMPnets (Mardt et al., ; Wu & Noé, ), or observable operator models (Wu et al., ). Markov state models are time-space discrete approximations of the transfer operator and Deep Generative MSM (Wu et al., ) and VAMPnets (Mardt et al., ) learn the space discretization through deep neural networks. Dynamic graphical models or Markov field models represent a time-space discrete approximation of the transfer operator injecting a (conditional) independence assumption of molecular subsystems, when modeling the transition probability, leading to better scaling for systems with poor time-scale separation (Olsson & Noé, ). Apart from ITO (Schreiner et al., ) several other deep generative approaches for modeling the transition density of molecular dynamics have recently proposed. Timewarp where a normalizing flow is used to encode the transition density (Klein et al., ) with limited transferability and metropolized sampling is used to access asymptotically unbiased equilibrium samples (Hastings, ). Score dynamics use a DDPM to model the displacements of an initial configuration towards a time-lagged one, achieving picosecond time-steps simulation and limited transferability (Hsu et al., ). However, unlike in the context of MSMs (Trendelkamp-Schroer et al., ; Rosta & Hummer, ; Wu et al., ; ), there are no deep generative transition density surrogates available leveraging available information about the equilibrium distribution — BoPITO is one such method.

## 6 Limitations and future work

##### Choice of hyper-parameter λ ^ ^ 𝜆 \hat{\lambda} over^ start_ARG italic_λ end_ARG

The hyper-parameter $\hat{\lambda}$, defines a global relaxation or mixing time-scale of the dynamics after which the model is guaranteed to sample equilibrium. In Appendix A.9 we describe a protocol to determine a bound for this parameter. However, developing a similar protocol for a transferable BoPITO model would likely require modifications to accommodate the diversity of global relaxation time-scales across different systems.

##### No Chapman-Kolmogorov Guarantee

When training on biased or off-equilibrium data where we rely on establishing ergodicity through interpolation we cannot guarantee self-consistency of the dynamics in the Chapman-Kolmogorov sense.

##### Surrogate model

BoPITO inherits the current limitations of ITO, such as generalization over chemical space and thermodynamic variables, and scaling. Furthermore, current models cannot guarantee unbiased sampling dynamics for non-equilibrium ensembles, which would require closed-form expressions for the target path probabilities.

## 7 Conclusion

We introduce Boltzmann Priors for Implicit Transfer Operator Learning (BoPITO), a framework to enhance ITO learning in two ways. First, a broad sampling of configuration space is used to initialize short off-equilibrium MD simulations. Second, we parameterize the transition density as an interpolation towards a pre-trained Boltzmann Generator, improving sample efficiency by an order of magnitude. BoPITO is a principled approach to embedding prior knowledge of the stationary distribution of Markovian dynamics as an inductive bias for long-term dynamical behavior. The approach enables interpolation between models trained on off-equilibrium data and the equilibrium distribution, and we can recover accurate models of unseen dynamics when informed by unbiased observables. Consequently, BoPITO is the first method to allow for the integration of multiple sources of information into the generation of deep generative surrogates of molecular dynamics.

## Appendix A Appendix

Supporting information for Equilibrium priors and uncertainty estimation for deep generative models of molecular dynamics

Juan Viguera Diez^1,2, Mathias Schreiner^1, Ola Engkvist^1,2, and Simon Olsson^1∗

### A.1 Eigen decomposition of the Transfer Operator and transition density

The two-fold composition of the transfer operator, eq. 5, acting on an initial $\mu$-weighted density, $\rho$, is

$$ $\displaystyle[T^{2}_{\Omega}(\tau)\circ\rho](\mathbf{x}_{t+2\tau})$ $\displaystyle=\sum_{i=1}^{\infty}\lambda_{i}(\tau)\left\langle\sum_{j=1}^{ \infty}\lambda_{j}(\tau)\langle\rho|\phi_{j}\rangle\ \psi_{j}|\phi_{i}\right \rangle\ \psi_{i}(\mathbf{x}_{t+2\tau})$ $\displaystyle=\sum_{i=1}^{\infty}\lambda_{i}(\tau)\left(\sum_{j=1}^{\infty} \lambda_{j}(\tau)\langle\rho|\phi_{j}\rangle\ \langle\psi_{j}|\phi_{i}\rangle \right)\ \psi_{i}(\mathbf{x}_{t+2\tau})$ $\displaystyle=\sum_{i=1}^{\infty}\lambda^{2}_{i}(\tau)\langle\rho|\phi_{i} \rangle\ \psi_{i}(\mathbf{x}_{t+2\tau}),$ $$

where we have used the orthonormality of the eigenfunctions, that is $\langle\psi_{j}|\phi_{i}\rangle=\delta_{ij}$. Similarly,

$$ $\displaystyle[T^{N}_{\Omega}(\tau)\circ\rho](\mathbf{x}_{t+N\tau})=\sum_{i=1}^ {\infty}\lambda^{N}_{i}(\tau)\langle\rho|\phi_{i}\rangle\ \psi_{i}(\mathbf{x}_ {t+N\tau}).$ (18) $$

The transition density can be re-written in terms of the spectral decomposition of the transfer operator by choosing the initial density, $p$, as a Dirac delta function $\delta_{\mathbf{x}_{t}}(\mathbf{x})$, that is $\rho(\mathbf{x})=\frac{\delta_{\mathbf{x}_{t}}(\mathbf{x})}{\mu(\mathbf{x})}$. Then,

$$ $\displaystyle p(\mathbf{x}_{t+N\tau}|\mathbf{x}_{t})$ $\displaystyle=\mu(\mathbf{x}_{t+N\tau})\left[T^{N}_{\Omega}(\tau)\circ\frac{ \delta_{\mathbf{x}_{t}}}{\mu}\right](\mathbf{x}_{t+N\tau})$ $\displaystyle=\mu(\mathbf{x}_{t+N\tau})\sum_{i=1}^{\infty}\lambda^{N}_{i}(\tau )\left\langle\ \frac{\delta_{\mathbf{x}_{t}}}{\mu}\middle|\phi_{i}\right \rangle\ \psi_{i}(\mathbf{x}_{t+N\tau})$ $\displaystyle=\sum_{i=1}^{\infty}\lambda^{N}_{i}(\tau)\left\langle\ \delta_{ \mathbf{x}_{t}}(\mathbf{x})|\psi_{i}\right\rangle\ \psi_{i}(\mathbf{x}_{t+N \tau})\ \mu(\mathbf{x}_{t+N\tau})$ $\displaystyle=\sum_{i=1}^{\infty}\lambda_{i}^{N}(\tau)\psi_{i}(\mathbf{x}_{t}) \phi_{i}(\mathbf{x}_{t+N\tau}).$ $$

### A.2 Decay of time-dependent score term

In Figure 6, we show the average time-dependent component of the score,

$$ $\displaystyle s_{N}=\mathbb{E}_{i\sim\mathbf{I},t_{\mathrm{diff}}\sim\mathcal{ U}(0,T)}\left[\hat{\lambda}^{N}s_{\mathrm{dyn}}(\mathbf{x}_{t_{i}+N\tau}^{t_{ \text{diff}}},\mathbf{x}_{t_{i}},N,t_{\text{diff}},{\bm{\theta}})\right],$ (19) $$

of a trained model BoPITO model of the Prinz Potential with $\hat{\lambda}=0.994$. The model remains flexible for small lags but eventually decreases to 0 0, sampling from the equilibrium model.

Figure: Figure 6: Average time-dependent component of the score, eq. 19, of a trained model BoPITO model of the Prinz Potential.
Refer to caption: x7.png

### A.3 The score of the transition probability and BoPITO

The score of the transition probability can be written as

$$ $\displaystyle\begin{split}\nabla_{\mathbf{x}_{N\tau}}\log p(\mathbf{x}_{N\tau} |\mathbf{x}_{0},N)&=\sum_{i}^{\infty}\lambda_{i}(\tau)^{N}\frac{\psi_{i}( \mathbf{x}_{0})}{p(\mathbf{x}_{N\tau}|\mathbf{x}_{0},N)}\nabla_{\mathbf{x}_{N \tau}}\phi_{i}(\mathbf{x}_{N\tau})\\ &=\frac{\nabla_{\mathbf{x}_{N\tau}}\mu(\mathbf{x}_{N\tau})}{p(\mathbf{x}_{N \tau}|\mathbf{x}_{0},N)}+\sum_{i=2}^{\infty}\lambda_{i}(\tau)^{N}\frac{\psi_{i }(\mathbf{x}_{0})}{p(\mathbf{x}_{N\tau}|\mathbf{x}_{0},N)}\nabla_{\mathbf{x}_{ N\tau}}\phi_{i}(\mathbf{x}_{N\tau})\\ &=\frac{\nabla_{\mathbf{x}_{N\tau}}\mu(\mathbf{x}_{N\tau})}{\mu(\mathbf{x}_{N \tau})}\frac{\mu(\mathbf{x}_{N\tau})}{p(\mathbf{x}_{N\tau}|\mathbf{x}_{0},N)} \\ &+\hat{\lambda}^{N}\sum_{i=2}^{\infty}\left(\frac{\lambda_{i}(\tau)}{\hat{ \lambda}}\right)^{N}\frac{\psi_{i}(\mathbf{x}_{0})}{p(\mathbf{x}_{N\tau}| \mathbf{x}_{0},N)}\nabla_{\mathbf{x}_{N\tau}}\phi_{i}(\mathbf{x}_{N\tau}).\end {split}$ $$

$\frac{\nabla_{\mathbf{x}_{N\tau}}\mu(\mathbf{x}_{N\tau})}{\mu(\mathbf{x}_{N
\tau})}=\nabla_{\mathbf{x}_{N\tau}}\log\mu(\mathbf{x}_{N\tau})$ is the score of the Boltzmann distribution and can be modeled with the score of a Boltzmann Generator, $s_{\mathrm{eq}}(\mathbf{x}^{t_{\text{diff}}},t_{\text{diff}})$. $\frac{\mu(\mathbf{x}_{N\tau})}{p(\mathbf{x}_{N\tau}|\mathbf{x}_{0},N)}\rightarrow
1$ as $N\rightarrow\infty$ and can be modeled with a scalar neural network $1+\hat{\lambda}^{N}f_{\bm{\theta}}(\mathbf{x}_{N\tau}^{t_{\text{diff}}},
\mathbf{x}_{0},N,t_{\text{diff}})$. The last term can be modeled with $g_{\bm{\theta}}(\mathbf{x}_{N\tau}^{t_{\text{diff}}},\mathbf{x}_{0},N,t_{\text
{diff}})$. The corresponding score model is

$$ $\displaystyle s_{\bm{\theta}}(\mathbf{x}_{N\tau}^{t_{\text{diff}}},\mathbf{x}_ {0},N,t_{\text{diff}})=s_{\mathrm{eq}}(\mathbf{x}^{t_{\text{diff}}},t_{\text{ diff}})(1+\hat{\lambda}^{N}f_{\bm{\theta}}(\mathbf{x}_{N\tau}^{t_{\text{diff}} },\mathbf{x}_{0},N,t_{\text{diff}})+\hat{\lambda}^{N}g_{\bm{\theta}}(\mathbf{x }_{N\tau}^{t_{\text{diff}}},\mathbf{x}_{0},N,t_{\text{diff}})$ $\displaystyle=s_{\mathrm{eq}}(\mathbf{x}^{t_{\text{diff}}},t_{\text{diff}})+ \hat{\lambda}^{N}\left(s_{\mathrm{eq}}(\mathbf{x}^{t_{\text{diff}}},t_{\text{ diff}})f_{\bm{\theta}}(\mathbf{x}_{N\tau}^{t_{\text{diff}}},\mathbf{x}_{0},N,t _{\text{diff}})+g_{\bm{\theta}}(\mathbf{x}_{N\tau}^{t_{\text{diff}}},\mathbf{x }_{0},N,t_{\text{diff}})\right).$ $$

We can aggregate $f$ and $g$ to a single neural network component, $s_{\mathrm{dyn}}(\mathbf{x}_{N\tau}^{t_{\text{diff}}},\mathbf{x}_{0},N,t_{
\text{diff}})$, to get

$$ $\displaystyle s_{\bm{\theta}}(\mathbf{x}_{N\tau}^{t_{\text{diff}}},\mathbf{x}_ {0},N,t_{\text{diff}})=s_{\mathrm{eq}}(\mathbf{x}^{t_{\text{diff}}},t_{\text{ diff}})+\hat{\lambda}^{N}s_{\mathrm{dyn}}(\mathbf{x}_{N\tau}^{t_{\text{diff}}} ,\mathbf{x}_{0},N,t_{\text{diff}})$ (20) $$

### A.4 Experimental parameters

#### A.4.1 Boltzmann Priors for dataset generation

For different number of trajectories, $n$, we train $5000/n$ ITO models on $n$ Prinz potential trajectories of length 150. Trajectories do not overlap among different trainings. The maximum model lag is 100 steps.

#### A.4.2 BoPITO samples efficiently long-term dynamics in data-sparse scenarios

##### Prinz potential

10 models are trained on non-overlapping trajectory sets for different numbers of trajectories. The trajectory length is $10,000$ and the maximum model lag is $1,000$. We consider lags of $10,25\text{, and }50$ as short, $100,200\text{, and }300$ as medium, and $500,750\text{, and }10,00$ as long time-scales in our experiments, see Figure 8 (b) for a visual reference. $\hat{\lambda}$ is set $0.982$ for BoPITO experiments with $1$, $5$ and $10$ trajectories and to $0.994$ for the rest.

##### Alanine Dipeptide

10 models are trained on potentially overlapping random trajectory sets for different numbers of trajectories. The trajectory length is $12,500$ and the maximum model lag is $10,000$. We consider lags of $5,10\text{, and }50$ as short, $100,500\text{, and }1000,$ as medium, and $2500,7500\text{, and }10000$ as long time-scales in our experiments, see Figure 9 (b). $\hat{\lambda}$ is set $0.9996$ for all BoPITO models.

#### A.4.3 Interpolating between models trained on off equilibrium data and the Boltzmann distribution

We remove the transitions resembling the slowest process in the system, $\phi$ crossing 0 0 or $2$, to generate biased simulation data. When a transition occurs, we remove one frame before and after the transition and split the trajectory. We train both biased ITO and BoPITO models with a maximum model lag of $100$ on the resulting biased dataset. For lags $>100$, we sample the BoPITO interpolator fitted on the unbiased dynamic observable defined by eq. 22, see Figure 4. We perform interpolation by sampling with the score model in eq. 15 followed by one round of non-interpolation sampling with lag $100$ steps for local relaxation.

#### A.4.4 Architectural and training details

We use MB-ITO as the architecture of models for the Prinz potential, and SE3-ITO, an $SE(3)-$equivariant neural network, for Alanine Dipeptide, both introduced in Schreiner et al. (). We report architectural and training hyper-parameters in Table 1. Models are trained until convergence in the log-log loss plot.

**Table 1: Architectural and training parameters of MB-ITO models (a) and SE3-ITO (b).**
| Diffusion steps | $500$ |
| --- | --- |
| Noise schedule | Sigmoidal |
| Batch size | $2,097,152$ |
| Learning rate | $0.001$ |
| Layers | $3$ |
| Embedding dimension | $256$ |
| Net dimension | $256$ |
| Optimizer | Adam |

### A.5 Prinz potential

The Prinz potential is a 1D potential commonly used for benchmarking MD sampling methods. The potential is defined as

$$ $\displaystyle U(x)=4\left(x^{8}+0.8e^{-80x^{2}}+0.2e^{-80(x-0.5)^{2}}+0.5e^{-4 0(x+0.5)^{2}}\right).$ (21) $$

We generate trajectories using an Euler-Maruyama integrator using the library Deeptime (Hoffmann et al., ). We set the integrator time-step to $1\cdot 10^{-5}$, and the temperature, mass, and damping factor to 1. In Figure 7 we show histograms of the position of a particle after $N$ steps, starting from $0.75$, and in Figure 8 (a) we report an implied time-scales plot. We use the identity function to define a dynamic observable/correlation function of this system and show it in Figure 8 (b).

Figure: Figure 7: Conditional density of a simulation of the Prinz potential starting at $x=0.75$ for different lags (steps). Long-term dynamics approaches the Boltzmann distribution.
Refer to caption: x8.png

Figure: (a)
Refer to caption: x9.png

### A.6 Alanine Dipeptide

Alanine Dipeptide is a small peptide with 22 atoms. We use publicly available data from Dibak et al. (), containing a total of $1\ \mu s$ simulation time split in $20$ trajectories. Simulation is performed in an implicit solvent with $2\ fs$ integration time-step and data is saved every $1\ ps$. We choose,

$$ $\displaystyle a(x)=b(x)=\begin{bmatrix}\sin{\phi}\\ \cos{\phi}\end{bmatrix},$ (22) $$

to define a dynamic observable of this system. The torsion angle $\phi$ is involved in the slowest transition observed in the simulation, see Figure 10. We combine these vectors taking the inner product for computing dynamic observables. We show this correlation function in Figure 9 (a), and implied time-scales plot in Figure 9 (b) and a histogram of the torsions $\phi$ and $\psi$ aggregating all simulation data in Figure 10.

Figure: (a)
Refer to caption: x11.png

Figure: Figure 10: Histogram of the torsional angles $\phi$ and $\psi$ of Alanine Dipeptide (insert). Data is aggregated among all trajectories in the dataset introduced in Dibak et al. ().
Refer to caption: x13.png

### A.7 Metrics

We evaluate models computing differences in the dynamic observables introduced in section 4.1 w.r.t. long MD simulations,

$$ $|\Delta\text{correlation}|_{N}=|O^{*}_{N}-O^{\text{model}}_{N}|,$ (23) $$

where $O^{*}_{N}$ is the normalized observable predicted by MD for lag $N$ and $O^{\text{model}}_{N}$ is the model’s prediction. See Appendix A.10 for details on normalization. We report the average difference over different training runs.

### A.8 Energies for Alanine Dipeptide

In Figure 11 we observe a remarkable agreement in energy densities of samples generated with our Boltzmann Generator and MD. GBSAOBForce is the implicit solvent model energy.

Figure: Figure 11: Energy and energy components of samples generated with a Boltzmann Generator (BG) and MD.
Refer to caption: x14.png

### A.9 Fitting λ ^ ^ 𝜆 \hat{\lambda} over^ start_ARG italic_λ end_ARG

The hyperparameter $\hat{\lambda}$ controls the time-scale at which the BoPITO model transitions to sampling the equilibrium model. If $\hat{\lambda}$ is too small, the model may prematurely relax to equilibrium, limiting its ability to capture non-equilibrium dynamics. Conversely, if $\hat{\lambda}$ is too large, the benefits of the BoPITO framework could diminish. Therefore, careful tuning of $\hat{\lambda}$ can be crucial for optimal performance.

As illustrated in Figure A.9, grid-search hyperparameter tuning can select an appropriate $\hat{\lambda}$ for a single-system model. Too small values of $\hat{\lambda}$ can hinder the model’s ability to capture long-time-scale dynamics, leading to increased loss. We recommend choosing the smallest $\hat{\lambda}$ that yields a plateau in the $\hat{\lambda}$-loss curve (elbow rule), as demonstrated in Figure A.9 (b). However,
practitioners should be cautious about sampling longer lags than the implied time-scale defined by $\hat{\lambda}$ if they cannot guarantee the system relaxes to equilibrium for those time-scales.

Figure: (a)
Refer to caption: x15.png

### A.10 Correlation function normalization

We subtract the mean and divide by correlation at lag 0 0 to normalize our dynamic observables. Subtracting the mean guarantees that the correlation function asymptotically decays to 0 0, and we divide by the correlation at lag 0 0 so that correlation can be read as a fraction of correlation at time 0 0. The resulting normalized correlation function is

$$ $\displaystyle\frac{E[\left(f(\mathbf{x}_{t})-E[f(\mathbf{x}_{t})]\right)\left( g(\mathbf{x}_{t+\Delta t})-E[g(\mathbf{x}_{t+\Delta t})]\right)]}{E[\left(f( \mathbf{x}_{t})-E[f(\mathbf{x}_{t})]\right)\left(g(\mathbf{x}_{t})-E[g(\mathbf {x}_{t})]\right)]}.$ (24) $$

When comparing different methods, we compute $E[f(\mathbf{x}_{t})]$, $E[g(\mathbf{x}_{t+\Delta t})]$ and $E[\left(f(\mathbf{x}_{t})-E[f(\mathbf{x}_{t})]\right)\left(g(\mathbf{x}_{t})-E
[g(\mathbf{x}_{t})]\right)]$ using long MD simulations and use these same normalizing factors for all methods.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - I do not fully follow the advantage of BoPITO recovering the unbiased Boltzmann distribution over long time horizons: isn’t this something that only applies to equilibrium distributions? But at the same time, it seems like the authors are saying they can incorporate off-equilibrium data and/or estimate something about off-equilibrium simulations?

- Can the authors compare their method to other machine learning methods for molecular dynamics simulation?

- Is energy being conserved? How can you guarantee physical dynamics?

- Can the authors show that this method has potential to be broadly applicable by seeing if it works for harder systems, such as larger proteins, or polymers?
soundness: 2
strengths: - This method could enable running MD simulations over much longer timescales
- By incorporating Boltzmann priors into the training process, this could allow more data efficiency and also ensure recovering the Boltzmann distribution over time
summary: This work looks at predicting molecular dynamics simulations quickly over time, using longer time steps. The authors introduce a method using diffusion models to train on both unbiased and biased MD data. The diffusion model is used to estimate the transition density, with the goal of simulating MD more quickly. BoPITO, using Boltzmann priors, is supposed to be more sample efficient than ITO (Implicit Transfer Operator). In BoPITO, the authors also use pre-trained models of the equilibrium distribution to improve the ITO method. The authors test the method on the 1D Prinz potential and alanine dipeptide.
weaknesses: - There are many ways to do molecular dynamics simulations. A common way is through methods like machine learning potentials. While the time steps that are taken there are smaller, these potentials are much more accurate and have been used for more diverse systems. It would be helpful if the authors provided some comparison to this method, as well as other methods that also try to directly predict molecular dynamics.

- The authors have only run this on two very toy systems: the 1D Prinz potential and alanine dipeptide. These systems are very far off from getting closer towards practical utility of this method. It would be more convincing if there were other systems studied, such as larger proteins or systems like the polymers or electrolytes system from Fu et al. (2023) TMLR. 


- This method does not seem to guarantee important properties such as energy conservation.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: 1. Definition regarding bias and off-equilibrium

The papers uses the term biased & unbiased, off-equilibrium. As far as I understood, the definitions could be summarized as the following:
- Off-equilibrium data: simulation results (trajectories) that did not explore some modes of Boltzmann distribution
- Biased simulation & samples: simulation done & sampled achieved by biased molecular dynamics, e.g. enhanced sampling, a modified version of the naive molecular dynamics
- Unbiased simulations & samples: simulation done & sampled achieved by naive molecular dynamics, possibly may not find some modes if the time horizon is too short

To sum up, is the following a correct summarization for the method in section 3.2 and section 4?
- BG are trained with unbiased and equilibrium data for a score model
- Then the time-dependent components score are then trained with unbiased data, possibly off-equilibrium 
- Section 4.3 presents that BoPITO can approximate long-term dynamics in the above situation
- Section 4.4 shows that the training data for BG are biased, and BoPITO can recover it?

2. Boltzmann generators in prior

Boltzmann generators [1] were trained by combining two modes: training by energy and training by example.
- Training by energy: Generate a distribution in the configuration space from a Gaussian prior using the model, and minimize the difference between Boltzmann distribution using the energy term as weights.
- Training by example: Since training by energy alone tends to focus on sampling on the most stable meta-stable state, they additionally use training by some valid configurations, e.g., trajectories from naive MD simulations.
Now, my question is about section 4.2 of the paper, where “BG is available before data collection” and “BG trained on equilibrium data” seems to be a contradiction. I may have misunderstood something, could the authors please clarify this part?

3. Derivation of parameterization

In Appendix A.3, the authors derive the score of the transition probability. When deriving equation, is $s_{dyn}(\ldots) = s_{eq}(x^{t_{diff}}, t_{diff}) f_\theta(\ldots) + g_\theta(\ldots)$ ?

4. Comparison to ITO

From Figure 3 to Figure 4, the authors compare ITO and BoPITO. As far as I understood, the difference between these two methods is the parameterized epsilon, i.e. score function. They have separated the equilibrium part and time-dependent components, and have additionally trained the score function by fixing the equilibrium part and only updating the time-dependent term. Are all the other components, i.e., model and data, identical?

[1] Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning, Science 2019
soundness: 3
strengths: The paper applies Boltzmann generators in ITO models to capture modes that have very small probability mass in the Boltzmann distributions.
summary: The authors point out that a prior work Implicit Transfer Operator (ITO) relies on extensive unbiased simulation, and suggest Boltzmann priors to improve the ability to approximate dynamics. To be specific, the authors first train a score model using the equilibrium data, afterwards fix the score model and only train the time-dependent components with unbiased and possibly off-equilibrium MD data.
weaknesses: 1. Pre-trained Boltzmann generator (BG)

My main concern with this paper is that it assumes a pre-trained BG. Since BG is trained using the energy function and data samples (more detail in Q2), I am a bit confused about what the authors intended. Is BG only used to generate initial conditions for simulations, for various sampling across the configuration space? 

2. Lack of detail of less extensive use of data compared to ITO

One of the important limitations the authors pointed out was the need for extensive,  unbiased molecular dynamics data. However, I do not see any results from the experiment nor analysis that BoPITO succeeds in generating results similar to ITO with less data, especially for Alanine Dipeptide.

3. Limited baselines and molecular systems
    - While the authors claim that ITO can be improved with BG, however, I don’t see enough experimental results to conclude it.
    - Comparison with ITO for molecular systems is done for only Alanine Dipeptide, while ITO has presented qualitative results and quantitative results for the folding process of Chignolin.
    - Comparison of quantitative stationary and dynamic observables to naive molecular dynamics has not been presented, while it has been for ITO for four molecular systems (Chignolin, Trp-Cage, BBA, Villin). To support the claim that BopITO truly predicts time-correlation statistics (dynamic observable), this seems to be needed.
    
    Overall, it seems that the paper does not present enough experiment results to conclude that BoPITO improves compared to ITO [1].
    
4. Missing ground truth transition probability for Alanine Dipeptide

One main and important result is the condition transition probability density for Alanine Dipeptide, figure 5. Though the figure shows that BoPITO and ITO (unbiased) show similar results, there are no ground truth data given to compare with. Though I understand it might be impractical to compute the transition densities for alanine dipeptide for 15,000 trajectories of length up to 1 nano-second compared to trajectories of length 500ps as in ITO[1], it would be good if the authors clarified this.

5. Requiring a well-trained Boltzmann generator

Since the proposed method uses Boltzmann generators as priors, a well-trained Boltzmann generator is essential. Although it has rather easily been done for Alanine Dipeptide, it seems that extension to complex molecular systems such as Chignolin is not flexible. 

[1] Implicit Transfer Operator Learning: Multiple time-resolution surrogates for Molecular dynamics, NeurIPS 2023

Minor

- Typos
    - Page 4 last paragraph, improve ITO learning in three ways ⇒ fourth ways
- Some notations are not clarified & used
    - $x$: molecular configuration first in page 1, it is the 3D coordinates of molecular system being $x \in \mathbb{R}^{3 \times M}$ where $M$ is the number of particles?
    - Diffeomorphic map $\mathcal{F}_\theta: \mathbb{R}^N \rightarrow \mathbb{R}^N$ , are the dimension of the latent space and configuration space necessarily identical?
- section 2.3 - transfer operators
    - The part explaining equation 5, the eigenvalues, and right and left eigenfunctions sounds awkward.
- Figure 3
    - There are no (a), (b), (c), (d) in the figure, as it is pointed out in section 4.3.
    - The authors ought to show that using BG as priors is better in terms of different time scales and number of trajectories, so it might have been better to plot the columns and legends switched, e.g., ITO / BoPITO and different time scales.
- Figure 5
    - I can’t find the black cross in the first column indicating the initial condition. Could you increase the size of the mark or change the color to something like red?
- Figure 8, 9
    - Explanation on the colors is missing. Could you add a legend for it?
- Equation 15
    - Theta seems to be missing in ${s_{dyn}}$
    - Subscripts for ${x^{t_{diff}}}$ seems to be missing?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: - At the bottom of page 2 the authors write: "As a result, most MD data will be ‘off-equilibrium’, e.g. trajectories exploring one or a few of the modes of $\mu(x)$". Is it correct to interpret 'off-equilibirum' in this case as 'non-ergodic'?
- In line 59-60: " [...] BoPITO, by construction, guarantees asymptotically unbiased equilibrium statistics.". Where do asymptotic guarantees for unbiasedness come from? I presume that reweighting still needs to be enforced in order to have an asymptotically unbiased sampler [1] though it is not immediately clear to me where this occurs. 
- in line 116-117: "BGs are trained either with approximately equilibrated simulation data or with biased simulations, from i.e. enhanced sampling, by employing appropriate reweighting[...]". I think this is not entirely correct, strictly speaking. BGs can also be trained in a data-free manner and a combination of data-informed (biased/unbiased) and data-free losses (see [2]). Did the authors solely rely on negative log-like hood training or also on data-free losses (e.g., Reverse KL)?
- I believe it would be useful to have a reference around line 120 about sampling with reweighting to enforce asymptotic unbiasedness, see e.g., [1]. 
- What do the authors mean by **importance re-sampling**?
- When looking at the results in Fig. 2 I am not sure what is the crystal baseline. Is it some SOTA or how is this obtained?
- As a follow-up, question: are there any other ML-based models that can be used in this case in order to prove the goodness of the proposed method? e.g., some other baselines that are worth comparing here such as naive ITOs, or some sort of BGs able to handle MD at different time steps [3]. 
- Is there an intuitive explanation for having in the top right corner of Fig. 3 the long time scale difference in correlation to be smaller than medium and short?
- What is an intuitive understanding of $N_{\textrm{int}}$? Is that the resolution in the interpolation process between equilibrium and off-equilibrium distribution? In other words, some sort of annealing parameter?
- Are there any baselines to compare against in Fig.5? See also the weakness stated above in terms of lack of comparison with existing methods. For instance, does FAB [4] allow the computation of the same type of dynamics? If yes, how does it compare to BoPITO? If not, why?
- This is more of a general detour question due to my limited knowledge of MD simulations. In line 475 onwards, the authors write: "These approaches, in general, aim to accelerate molecular simulations akin to BoPITO, yet due to the CG operation, the molecular dynamics (kinetics) will be accelerated, and detailed knowledge of the unbiased dynamics are needed to correct this [...]". I am wondering how this can be possible. When coarse-graining a system and running MD at a coarser level, you inevitably lose information on the degrees of freedom at a smaller scale. I am naively wondering if there are ways to rigorously correct for this and still retain asymptotic guarantees in this case. 

- Lastly, I am wondering if the following intuition behind this work is correct: I train a core component of the equilibrium distribution using BGs. Then, I use some biased MD data to refine my backbone model in order to capture rare events, making it an 'ergodic sampler'. The idea is to use the ITOs paradigm to have some smooth interpolation between the equilibrium-off-equilibrium regimes encoded in the model in different ways (BG+biased MD data).


## References
- [1] [Nicoli, Kim A., et al. "Asymptotically unbiased estimation of physical observables with neural samplers." Physical Review E 101.2 (2020): 023304.](https://link.aps.org/accepted/10.1103/PhysRevE.101.023304)
- [2] [Noé, Frank, et al. "Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning." Science 365.6457 (2019): eaaw1147.](https://www.science.org/doi/10.1126/science.aaw1147)
- [3] [Klein, Leon, et al. "Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics." Advances in Neural Information Processing Systems 36 (2024).](https://proceedings.neurips.cc/paper_files/paper/2023/file/a598c367280f9054434fdcc227ce4d38-Paper-Conference.pdf)
- [4] [Midgley, Laurence Illing, et al. "Flow annealed importance sampling bootstrap." arXiv preprint arXiv:2208.01893 (2022).](https://arxiv.org/pdf/2208.01893)
soundness: 3
strengths: - The paper is clearly written and easy to follow. Overall, it represents a pleasant read. 
- The authors made a great effort to make the paper self-contained, including a lot of background sections covering the basics of MD, BGs, ITOs, and diffusion models. 
- Numerical results are convincing in showing advantages compared to ITOs (PRO), although the approach is not compared to other baselines/methods in the literature (CON). This allows us to only partially appreciate the advantages of the algorithm. 
- The authors address the limitations of the current methods and envision possible avenues for future work. 
- The code attached as additional material looks good, properly formatted, and seems to run out-of-the-box
summary: This paper proposes to combine Boltzmann Generators (BG) with Implicit Transfer Operators (ITOs). ITOs are limited because a lot of unbiased molecular dynamics (MD) data is needed for training. This framework combines ITOs with a BG before being able to sample efficiently and be asymptotically unbiased from the equilibrium distribution. 

From a more practical standpoint, the score function of the Score-based Diffusion Model is separated into an equilibrium contribution (from the BG prior) and the time-dependent component (to tackle off-equilibrium simulations). BoPITO defines an interpolation between models trained on off-equilibrium simulations and the equilibrium distribution, thus making it a more efficient sampler of long-time MD simulations.
weaknesses: - I find the main weakness of this work being the lack of comparison with other ML-based approaches that can deal with sampling at different time scales (e.g., [3,4]). Moreover, it would also be useful to show a ground truth in the plots when possible, e.g., Fig. 5. Adding more baselines would certainly improve the manuscript. 
- As a suggestion, rather than a weakness, I feel it would be better to move the discussion to the related work at the beginning of the paper. This would help people to associate some works they may already know with the present work in case they are familiar with the literature. This would represent a useful prior to have before getting into the core of the paper. 
- At the bottom of page 4, the authors say that BoPITO improves ITO learning in three different ways. However, in what follows, there are four points. I also believe it would be helpful to list these contributions visually with bullet points.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: * The formalism decomposes the operator in a sum of left and right eigenfunctions. Are these eigenfunctions used anywhere in the training objectives directly, or are they solely used to justify/derive the form of the score function ansatz? 
* In line 139 and 142 it says $\lambda_i = exp(-\tau \kappa_i)$ but also that $-1 < \lambda_i < 1$. How is it possible that $\lambda_i = exp(-\tau \kappa_i) < 0$ though?
* I am having some trouble understanding the interpolators section. N_max is the maximum training lag, but then we sample with N_int > N_max. Doesn't that cause out-of-distribution problems? Maybe you could add a bit more explanation in the appendix as to why the given model implies an interpolator. 
* line 60: How exactly do you define "off-equilibrium" and "biased"? 
* line 191 and following: what is $t_{diff}$?
* Looking at Figure 3, it seems like ITO and BoPITO have very similar performance on short and medium time scales, and BoPITO mostly outperforms ITO on long timescales, presumably due to the enforcement of the Boltzman distribution for long time horizons. I am wondering how often it is precisely the medium time scales that we are interested in though? 

## Suggestions
* Include a section that explains the ITO framework in more detail and includes pseudo-code
* line 107 and other places, I find formulations like "few of the modes of μ(x)" are easier to read if written as  "few of the modes of the *Boltzmann distribution* μ(x)"
* line 209 says "improve ITO learning in three ways" but then four points are listed
soundness: 3
strengths: * Temporal coarse-graining of molecular dynamics is of very high relevance to simulate long-horizon observables more efficiently
* Making use of cheaply obtainable off-equilibrium data is a very promising approach to train more scalable models
* Method seems sound and well-founded in physical/statistical principles 
* The results look promising, although somewhat limited due to the small scale experiments
summary: The paper uses a pretrained Boltzmann Generator to include un-biased equilibrium information into Implicit Transfer Operators (ITOs).
ITOs learn "time coarse grained" transition densities from molecular dynamics (MD) data.
The challenge is that ITOs require (a) unbiased (b) large quanities of MD data to train.
Boltzmann Generators (BG) can approximate the true Boltzmann distribution even if trained with biased data when combined with importance resampling.


The paper suggest three ways to improve ITOs with Boltzmann generators:
1. Sampling of conformations proportional to the equilibrium distribution from the Boltzmann Generator as starting points for trajectories
2. Inductive bias: Separate learning of the transition density into a fixed stationary part (the Boltzmann generator) and learned time-dependent component to improve sample efficiency.
(Based on the spectral decomposition of the transfer operator, where the equilibrium distribution is an eigenfunction with an eigenvalue of one.) An explicit temporal decay of the time-dependent component ensures convergence to the unbiased Boltzmann distribution for long-time horizons.
3. Approximate rare dynamics outside the data, by interpolating between ITOs trained on off-equilibrium data (time-dependent component) and the unbiased equilibrium distribution (Boltzmann generator).
weaknesses: * The paper would benefit from clearer writing and more extensive background sections. It seems to assume that readers are already familiar with the implicit transfer operators framework, which will often not be the case. More explanations and pseudocode (even if it is just in the appendix) would make the paper much easier to read.  
* The experiments are conducted only on small-scale systems, and it is unclear how well this approach would scale to more realistic scenarios.

Meta-review: As positives for this paper, reviewers have mentioned the importance of the topic of this paper, focusing on enabling temporal coarse grained molecular dynamics. Making use of cheaply obtainable off-equilibrium data was seen as a promising approach to train more scalable models. The proposed method was also lauded for its soundness and being well-founded in statistical principles. Finally, the code that the authors attached was found to be well written and running out of the box. On the initial smaller scale experiments, the results were found to be promising. 

As areas for improvement, a common point raised by multiple reviewers was the limited evaluation on the 1D Prinz potential and Alanine dipeptide. Another common point was the request to add comparisons against other baseline methods, with reviewer r47p asking for comparisons against MD simulations with ML potentials, and reviewer hu2r asking for comparison against other ML methods that can sample at different time scales. Furthermore, several reviewers asked for clarifications, and made suggestions to make the writing clearer and more accessible. 

The authors have added experiments on Chignolin, which has addressed one of the main concerns raised by the reviewers. The authors have explained why they have omitted comparisons against MD simulations with ML potentials, explaining this is an orthogonal direction and that it is important to compare methods using the same underlying potential energy surface. Furthermore, the authors have explained why it’s hard to compare against other time-lagged ML methods for the point the paper wants to make. The requests for clarifications have been answered well by the authors and they have incorporated suggestions such as inclusion of pseudocode into their revised paper.

During the rebuttal phase 3 out of 4 reviewers have indicated that most of their concerns were addressed in the rebuttal. Reviewer r47p remained of the opinion that comparison against MD simulations with ML potentials and  demonstrations of larger scale applicability beyond the already added chignolin could improve the paper, but does not think this should block acceptance of the paper. Given the explanations of the authors regarding comparison with MD simulations with ML potentials, which seem reasonable to me, and the already added chignolin results, I conclude that the authors have addressed all major concerns raised by the reviewers. Given the unanimous recommendation to accept this paper among reviewers, I therefore recommend to accept this paper.

Decision: Accept (Poster)