EquiJump: Protein Dynamics Simulation via SO(3)-Equivariant Stochastic Interpolants

## Abstract

Abstract Mapping the conformational dynamics of proteins is crucial for elucidating their functional mechanisms. While Molecular Dynamics (MD) simulation enables detailed time evolution of protein motion, its computational toll hinders its use in practice. To address this challenge, multiple deep learning models for reproducing and accelerating MD have been proposed drawing on transport-based generative methods. However, existing work focuses on generation through transport of samples from prior distributions, that can often be distant from the data manifold. The recently proposed framework of stochastic interpolants, instead, enables transport between arbitrary distribution endpoints. Building upon this work, we introduce EquiJump, a transferable SO(3)-equivariant model that bridges all-atom protein dynamics simulation time steps directly. Our approach unifies diverse sampling methods and is benchmarked against existing models on trajectory data of fast folding proteins. EquiJump achieves state-of-the-art results on dynamics simulation with a transferable model on all of the fast folding proteins.

## 1 Introduction

Proteins are the workhorses of the cell, and simulating their dynamics is critical to biological discovery and drug design (Karplus and Kuriyan,, ). Molecular Dynamics (MD) simulation is an important tool that leverages physics for time evolution, enabling precise exploration of the conformational space of proteins (Hollingsworth and Dror,, ). However, sampling with physically-accurate molecular potentials requires small integration time steps, often making the simulation of phenomena at relevant biological timescales prohibitive (Lane et al.,, ).

To tackle this challenge, several studies have adopted deep learning models to capture surrogates of MD potentials and dynamics (Noé et al.,, ; Durumeric et al.,, ; Arts et al.,, ). More recent works (Schreiner et al.,, ; Li et al.,, ; Jing et al.,, ) have proposed to use deep learning-based simulators trained on long-interval snapshots of MD trajectories to predict future states given some starting configuration. These models draw from neural transport models (Ho et al.,, ; Lipman et al.,, ), learning a conditional or guided bridge between a prior distribution ($\rho_{0}=\mathcal{N}$) and the target data manifold of simulation steps ($\rho_{1}=\rho_{\textrm{data}}$). In contrast to this, the recent paradigm of Stochastic Interpolants (; Albergo and Vanden-Eijnden,, ) provides a method for directly bridging distinct arbitrary distributions.

In this work, we utilize this framework and introduce EquiJump, a Two-Sided Stochastic Interpolant model which bridges between long-interval timesteps of protein simulation directly (Figure 1). EquiJump is SO(3)-equivariant and simulates all heavy atoms directly in 3D. We train a transferable model on 12 fast-folding proteins (Majewski et al.,, ; Lindorff-Larsen et al.,, ) and successfully recover their dynamics.

Our main contributions are as follows:

- •
We extend the Two-Sided Stochastic Interpolants framework to simulate the dynamics of three-dimensional representations. By training on trajectory data, our approach directly leverages the close relationship between consecutive timesteps.
- •
We introduce a novel four-track SO(3)-equivariant neural architecture to implement a generative transport operator, which we apply to model the dynamics of all heavy atoms in protein monomers.
- •
We learn a transferable simulator for the dynamics of 12 fast-folding proteins of different sizes, demonstrating its effectiveness in reproducing long-term dynamics and ensemble distributions.

Figure: Figure 1: Direct bridging of 3D all-atom simulation time steps: EquiJump runs an stochastic interpolants-based transport process on coordinates and 3D geometric representations to generate future time frames from an initial state. Gray boxes depict transport across the learned latent space, which takes Gaussian noise perturbations and uses noise ($\boldsymbol{\hat{\eta}}$) and drift ($\boldsymbol{\hat{b}}$) predictions to directly transform all-atom proteins across time and 3D space.
Refer to caption: extracted/5922300/figures/interpolants.png

## 2 Related Work

Recent advancements in protein modeling through deep learning have led to the development of several models capable of replicating molecular dynamics (MD) trajectories. (Wang et al.,, ; Husic et al.,, ) introduced supervised models trained through direct force matching, demonstrating their ability to transfer simulations across different proteins. (Majewski et al.,, ) represents builds a unified transferable model for multiple proteins. Their approach relies on force matching, which approximates atomistic forces. In contrast, EquiJump eliminates the need for force data by using a two-sided stochastic transport framework that directly generates next steps of configurations, enabling time stepping that is orders of magnitude larger than those required by this method. We show that our model outperforms force-matching approaches by maintaining kinetic consistency and replicating dynamical observables more precisely.

(Fu et al.,, ) learns
to predict accelerated, coarse grained dynamics of polymers with GNNs. (Köhler et al.,, ) utilizes Normalizing Flows (Gabrié et al.,, ) for coarse grained force-matching, while (Arts et al.,, ) builds upon Denoising Diffusion Probabilistic Models (DDPM) (Ho et al.,, ) through Graph Transformers (Shi et al.,, ; Costa,, ).
These approaches focus heavily on coarse-grained representations, which limit their ability to simulate the full complexity of protein dynamics at the all-atom level. In contrast, EquiJump operates directly at the all-atom scale, achieving efficiency through SO(3)-equivariant neural networks and residue representation (Costa et al.,, ). This allows EquiJump to provide a more accurate and comprehensive simulation of protein dynamics, combining the precision of all-atom modeling with the computational efficiency typically seen in coarse-grained approaches.

Recent models have proposed generating samples from a prior distribution while conditioning on an initial configuration. Timewarp (Klein et al.,, ) enhances MCMC sampling with conditional normalizing flows, while ITO (Schreiner et al.,, ) uses a PaiNN-based network (Schütt et al.,, ) to learn a conditional diffusion model for next-step prediction. Similarly, F^3low (Li et al.,, ) employs FramePred (Yim et al.,, ) and Optimal Transport Guided Flow Matching (Zheng et al.,, ). Finally, (Jing et al.,, ) applies one-sided stochastic interpolants (Ma et al.,, ) to interpolate or extrapolate molecular configurations. These approaches rely on transforming Gaussian priors via stochastic (SDE) or ordinary differential equations (ODE), where the prior often lies far from the true data distribution. Instead, EquiJump uses two-sided stochastic interpolants to directly bridge trajectory snapshots, leveraging the configuration proximity of consecutive timesteps and enabling a transport that stays close to physical states (Figure 2; Appendix A.1).

## 3 Methods

### 3.1 Stochastic Interpolants

Figure: Figure 2: Neural Transport of Tensor Clouds. (a) DDPM defines an SDE for denoising samples from a Gaussian prior, while standard (b) Flow Matching traces a velocity field-based ODE for moving the Gaussian samples. (c) Two-Sided Stochastic Interpolants instead enable transporting through a local, normally-perturbed latent space that remains close to the manifold of the data endpoints.
Refer to caption: extracted/5922300/figures/transport.png

Neural transport methods have demonstrated outstanding performance in generative tasks (Ma et al.,, ; Liu et al.,, ; Lipman et al.,, ; Ho et al.,, ). Stochastic Interpolants (; Albergo and Vanden-Eijnden,, ) are a class of generative models that generalize flow matching and diffusion models, reaching state-of-the-art results in image generation (Ma et al.,, ; ). In particular, two-sided stochastic interpolants enable learning the transport between two arbitrary probability distributions (Figure 2).
Given two distributions $p_{0}$ and $p_{1}$, we define the stochastic process $\{\mathbf{X}_{\tau}\}$:

$$ $\displaystyle\mathbf{X}_{\tau}=I(\tau,\mathbf{X}_{0},\mathbf{X}_{1})+\gamma( \tau)\mathbf{Z}$ (1) $$

where $\tau\in[0,1]$ is the time parameterization, the endpoints are distributed as $\mathbf{X}_{0}\sim p_{0}$, $\mathbf{X}_{1}\sim p_{1}$, and the latent variable is distributed as $\mathbf{Z}\sim\mathcal{N}(0,\mathbb{I})$. To ensure boundary conditions, the interpolant $I$ must satisfy the following: $I(0,\mathbf{X}_{0},\mathbf{X}_{1})=\mathbf{X}_{0}$ and $I(1,\mathbf{X}_{0},\mathbf{X}_{1})=\mathbf{X}_{1}$, and $\gamma(0)=\gamma(1)=0$. The probability $p(\tau,\mathbf{X})$ of a stochastic interpolant satisfies the transport equation:

$$ $\displaystyle\partial_{\tau}p(\tau,\mathbf{X})+\nabla\cdot(b(\tau,\mathbf{X})p (\tau,\mathbf{X}))=0$ (2) $$

and the boundary conditions $p(0,\mathbf{X})=p_{0}$ and $p(1,\mathbf{X})=p_{1}$. Here, $b(\tau,\mathbf{X})$ is the expected velocity:

$$ $\displaystyle b(\tau,\mathbf{X})=\mathbb{E}\left[\partial_{\tau}\mathbf{X}_{ \tau}\ |\ \mathbf{X}_{\tau}=\mathbf{X}\right]=\mathbb{E}\left[\partial_{\tau}I (\tau,\mathbf{X}_{0},\mathbf{X}_{1})+\partial_{\tau}\gamma(\tau)\mathbf{Z}\ | \ \mathbf{X}_{\tau}=\mathbf{X}\right]$ (3) $$

We can similarly define the noise term $\eta(\tau,\mathbf{X})$ as:

$$ $\displaystyle\eta(\tau,\mathbf{X})=\mathbb{E}[\mathbf{Z}\ |\ \mathbf{X}_{\tau} =\mathbf{X}]$ (4) $$

In practice, the exact forms of $b$ and $\eta$ are not known for arbitrary distributions $p_{0}$, $p_{1}$, and are thus parameterized by neural networks. () shows that we can learn the functions $\hat{b}\approx b$ and $\hat{\eta}\approx\eta$ by optimizing:

$$ $\displaystyle\min_{\hat{b}}\int_{0}^{1}\mathbb{E}\left[\frac{1}{2}\left\lVert \hat{b}(\tau,\mathbf{X}_{\tau})\right\rVert^{2}-\left(\partial_{\tau}I(\tau, \mathbf{X}_{0},\mathbf{X}_{1})+\partial_{\tau}\gamma(\tau)\mathbf{Z}\right) \cdot\hat{b}(\tau,\mathbf{X}_{\tau})\right]d\tau$ (5) $$

$$ $\displaystyle\min_{\hat{s}}\int_{0}^{1}\mathbb{E}\left[\frac{1}{2}\left\lVert \hat{\eta}(\tau,\mathbf{X}_{\tau})\right\rVert^{2}+\mathbf{Z}\cdot\hat{\eta}( \tau,\mathbf{X}_{\tau})\right]d\tau$ (6) $$

We can then sample $\mathbf{X}_{\tau=1}\sim p(\tau=1,\mathbf{X}_{1})$ through an ordinary differential equation (ODE), or a stochastic differential equation (SDE):

$$ $\displaystyle d\mathbf{X}_{\tau}=\hat{b}(\tau,\mathbf{X}_{\tau})d\tau$ (7) $$

$$ $\displaystyle d\mathbf{X}_{\tau}=\left(\hat{b}(\tau,\mathbf{X}_{\tau})-\frac{ \epsilon(\tau)}{\gamma(\tau)}\hat{\eta}(\tau,\mathbf{X}_{\tau})\right)d\tau+ \sqrt{2\epsilon(\tau)}dW_{\tau}$ (8) $$

where $W_{\tau}$ is the Weiner process.
Once we have learned the expected velocity $\hat{b}$ and noise $\hat{\eta}$, the above equations can be integrated numerically starting from $(\tau=0,\mathbf{X}_{0}\sim p_{0})$ to $(\tau=1,\mathbf{X}_{1}\sim p_{1})$. Furthermore, following from eqs. 7 and 8 the probability $p(\tau,\mathbf{X}_{\tau})$ is SO(3)-equivariant when $\hat{b}$ and $\hat{\eta}$ are SO(3)-equivariant and $dW_{\tau}$ is isotropic. We provide more details on interpolant parameterization in Appendix A.2.

### 3.2 Two-Sided Stochastic Interpolants for Dynamics Simulation

We extend the two-sided stochastic interpolant framework to learn a time evolution operator from trajectory data $[\mathbf{X}^{t}]_{t=1}^{L}$. Given a source time step $\mathbf{X}^{t}$ and its consecutive target step $\mathbf{X}^{t+1}$, we define the distribution boundaries of our interpolant as $\rho_{0}=\rho(\mathbf{X}^{t})$ and $\rho_{1}=\rho(\mathbf{X}^{t+1}\mid\mathbf{X}^{t})$. The conditional nature of the target distribution requires that our predictions for drift $\hat{b}$ and noise $\hat{\eta}$ are explicitly conditioned on the source step $\mathbf{X}^{t}$. We apply this approach to simulating all-atom protein dynamics, as depicted in Figure 1. In this context, $\mathbf{X}^{t}$ represents a 3D all-atom protein conformation at time $t$, which is provided as input to our model. We frame it as the source distribution, and set $\mathbf{X}_{\tau=0}=\mathbf{X}^{t}$. We then employ an iterative process governed by the integration of eqs. 7 and 8 from $\tau=0$ to $\tau=1$. This produces a sample $\mathbf{X}_{\tau=1}$, which follows the distribution $\mathbf{X}_{\tau=1}\sim\rho_{1}$, generating a next step in the simulation $\mathbf{X}^{t+1}$.

### 3.3 Multimodal Interpolants of Geometric Representations

We treat data $\mathbf{X}$ represented as geometric features positioned in three-dimensional space, $\mathbf{X}=[(\mathbf{V}_{i},\mathbf{P}_{i})]_{i=1}^{N}$, which we refer to as the Tensor Cloud representation (Figure 2). In this formulation, each $\mathbf{V}_{i}$ is a tensor of irreducible representations (irreps) of O(3) or SO(3), associated with a 3D coordinate $\mathbf{P}_{i}\in\mathbb{R}^{3}$. The feature representations $\mathbf{V}$ are arrays of irreps up to order $l_{max}$. For each $l\in[0,l_{max}]$, the tensor $\mathbf{V}^{l}$ represents geometric features with dimensions $\mathbf{V}^{l}\in\mathbb{R}^{H\times(2l+1)}$, where $H$ denotes the feature multiplicity.

We extend interpolant eqs. 8 and 7 to the multi-modal type $\mathbf{X}_{i}=(\mathbf{V}_{i},\mathbf{P}_{i})$ by integrating geometric features and coordinate components as $d\mathbf{X}_{i}^{\tau}=(d\mathbf{V}_{i}^{\tau},d\mathbf{P}_{i}^{\tau})$. For computing the losses eqs. 5 and 6, we define the Tensor Cloud dot product as $\mathbf{X}_{i}\cdot\mathbf{X}_{j}=\mathbf{V}_{i}\cdot\mathbf{V}_{j}+\mathbf{P}
_{i}\cdot\mathbf{P}_{j}$. In general, treating the feature and coordinate components independently allows for different parameterizations of the interpolant. In this work, we use the same interpolant form for both components, and only adjusting the scaling $\gamma(t)$ of the variable $\mathbf{Z}$. For further details see Appendix A.2.

#### 3.3.1 Protein Structure Representation

We represent a protein monomer $(\mathbf{R},\mathbf{X})$ as a sequence $\mathbf{R}$ and a Tensor Cloud $\mathbf{X}$. Our model is designed to update $\mathbf{X}$ while being conditioned on $\mathbf{R}$. Each residue $i$ consists of three components: a residue label $\mathbf{R}_{i}\in\mathcal{R}=\{\textrm{ALA, GLY},\dots\}$, the Cα 3D coordinate $\mathbf{P}^{\alpha}_{i}\in\mathbb{R}^{3}$, and a geometric feature of order $l=1$ with multiplicity 13, $\mathbf{V}^{A}_{i}\in\mathbb{R}^{13\times 3}$. This feature encodes the relative 3D vector from the Cα to all other heavy atoms in the residue, following a canonical ordering. For residues with fewer than 13 non-Cα heavy atoms, we pad the atom vectors. This modeling approach, based on (King and Koes,, ; Costa et al.,, ), allows for the direct representation of all heavy atoms in 3D, while maintaining a coarse-grained representation anchored on the Cα.

Figure: Figure 3: EquiJump Architecture: (a) The Self-Interaction Layer updates geometric features independently, mixing $\mathbf{V}^{l}$ of different degrees into new features through a Tensor Square operation. (b) The Spatial Convolution layer updates representations by aggregating the tensor product of neighbors messages with the spherical harmonics embedding of the relative 3D vector between the positions of those neighbors. (c) We stack the above modules to form a block, and build a base network out of $L$ blocks for making predictions. (d) A shared conditioner and 4 headers are built from the base network. The conditioner processes sequence and the current simulation step, producing latent embeddings that are fed to the prediction headers. The headers independently predict features and coordinates updates for drift and noise components of the two-sided stochastic interpolant process.
Refer to caption: extracted/5922300/figures/model.png

#### 3.3.2 Neural Network Architecture and Training

To efficiently process 3D data, we utilize Euclidean-equivariant neural networks (Geiger and Smidt,, ; Miller et al.,, ). We design a neural network to predict the drift $\hat{b}=(\hat{b}_{\mathbf{V}},\hat{b}_{\mathbf{P}})$ and noise $\hat{\eta}=(\hat{\eta}_{\mathbf{V}},\hat{\eta}_{\mathbf{P}})$ terms, conditioned on the sequence $\mathbf{R}$, the source structure $\mathbf{X}^{t}$, the latent transport structure $\mathbf{X}^{t}_{\tau}$, and the latent time $\tau$ (Figure 3).

The EquiJump layer is built from two SO(3)-equivariant modules: the Self-Interaction (Figure 3.a) module for updating features $\mathbf{V}^{l}$ indepedently from coordinates, based on (Costa et al.,, ; Batatia et al.,, ); and the Spatial Convolution (Figure 3.b) module for sharing information between neighbors, based on Tensor Field Networks (Thomas et al.,, ). At each layer, we also employ residual connections and the SO(3)-equivariant layer norm from (Liao and Smidt,, ). We build a deep neural network (DNN) by stacking $L$ times the above blocks (Figure 3.c). Refer to Appendix A.2 for additional details.

We use 5 DNNs in our model (Figure 3.d): 1 conditioner network $\mathbf{f}_{\textrm{cond}}$ and 4 header networks for predicting each of $\hat{b}_{\mathbf{V}}$, $\hat{b}_{\mathbf{P}}$, $\hat{\eta}_{\mathbf{V}}$, $\hat{\eta}_{\mathbf{P}}$ independently. Given a configuration $\mathbf{X}^{t}$,
we first prepare a hidden representation $\tilde{\mathbf{X}}^{t}=\mathbf{f}_{\textrm{cond}}(\mathbf{R},\mathbf{X}^{t})$. The 4 headers take $(\tilde{\mathbf{X}}^{t},\mathbf{X}^{t}_{\tau},\tau)$ to produce predictions of each component of the drift $\hat{b}$ and the noise $\hat{\eta}$. For efficiency, the embedding $\tilde{\mathbf{X}}^{t}$ is made independent of $\tau$, and only the prediction headers are used in the integration loop of the latent transport. We train and sample these networks following Algorithms 1 and 2:

Figure: Algorithm 1 EquiJump Training

Figure: Algorithm 2 EquiJump Sampling

## 4 Experiments and Results

### 4.1 Fast-Folding Proteins

To evaluate the capability of our model to reproduce protein dynamics, we leverage the dataset of 12 fast-folding proteins produced by (Majewski et al.,, ), and originally investigated in (Lindorff-Larsen et al.,, ). The dataset consists of 100 ps spaced snapshots of MD for 12 proteins ranging from 10 to 80 residues. The trajectories are made up of several NVT runs (20 to 100 ns) at $T$=350 K from different starting configurations sampling the phase space, for an aggregated simulation time of hundreds to thousands of $\mu$s per protein. We refer to the original work for more details on the dataset.

While training on the whole dataset is effective in producing configurations within the original phase space, it is not optimal to correctly reconstruct the potential energy surface.
This is because when performing MD the relative probability of different configurations is extremely sensitive to the correct description of the dynamics of the transition states. However, for the slowest modes of the system these states are by their own nature very high in free energy and therefore heavily under-represented in our training dataset. For this reason, taking inspiration by classical sampling methods such as umbrella sampling (Torrie and Valleau,, ) and metadynamics (Laio and Parrinello,, ), we propose a reweighing of the training set that will enable our model to better learn the complex dynamics of transition states. To achieve this, we define collective variables and add a bias to the sampling probability in training, to compensate the free energy of the system.

We first define the relevant degrees of freedom through TICA analysis (Pérez-Hernández et al.,, ) (see below) and then construct a small number of clusters through k-means in this simplified space.
Since we are interested in the long-time dynamics of the system, this should be a reasonable basis to consider for reweighing the trajectories.
If our data came from a long NVT dynamics, the free energy $F_{i}$ of cluster $i$ would be proportional to the number of samples: $N_{i}\propto e^{-\beta F_{i}}$ ($\beta$ being the inverse temperature). In the discrete basis, the reweighing with $e^{\beta F_{i}}$ would thus naturally correspond to a uniform distribution over the clusters, as desired.
Since our dataset was obtained through a biased procedure to enhance variety (Majewski et al.,, ), the relation between probability and free energy does not hold true before reweighing. However, we still propose to sample uniformly on the clusters, as this is a simple procedure that does not depend on the population of each clusters, if not for its small effect on the definition of the clusters themselves. We show the cluster centers and distributions of population sizes for each protein in Figure 7. Refer to Appendix A.3 for more details.

Figure: Figure 4: EquiJump Samples: (a) We visualize the distribution in 3D of 1500 backbone random samples of EquiJump trajectories. We align samples to the crystal backbone (shown in black) and verify that our model stays close to the native state basin. (b) We show Ramachandran plots and (c) mean pairwise distance matrices of EquiJump samples against reference trajectory data.
Refer to caption: extracted/5922300/figures/samples.png

### 4.2 Model Training and Sampling

We train a model for time-evolving any of the 12 fast-folding proteins following Algorithm 1. We use the Adam optimizer (Kingma and Ba,, ) with linearly decreasing learning rate from $1\times 10^{-2}$ to $1\times 10^{-3}$ over 150k steps, with batch size of 512. We perform all our experiments on NVIDIA A100 machines with varying number of GPUs. We investigate diverse model hyperparameters and choose the model that achieves lowest loss.

We sample the dynamics by starting from a configuration and iteratively applying the model. We employ 500 steps of integration to obtain the next configuration. While each evaluation is clearly slower than the original model used for MD, our timesteps are effectively 25000 times longer. Moreover, simulations can be very efficiently parallelized on modern hardware. We further verify that our model can be iterated for more than 50000 steps without leaving the original simulation basin. To analyze the quality of our inference, we perform 500 simulations of 500 steps starting from points taken from the training set.

### 4.3 Configuration Space and Free Energy Surface Analysis

In Figure 4, we present samples and compare the densities of dihedral angles and pairwise $\mathbf{C}_{\alpha}$ distances between EquiJump and reference trajectory data. By accurately recovering these metrics across the sampled conformations, we demonstrate that EquiJump remains within the configuration space of the original dataset and effectively avoids sampling nonphysical states.

Figure: Figure 5: Free Energy on TICA components for the 12 Fast-Folding Proteins. We compare the free energy of EquiJump against that of the reference and that of available model CG-MLFF. The free energy for each plot is set to 0 at the minimum and the color map is in units of $k_{\mathrm{B}}T$ at the MD temperature of 350 K. EquiJump succesfully recovers the dynamics of the proteins, covering the phase space and stabilizing the basin of most.
Refer to caption: extracted/5922300/figures/free-energy-comparison.png

It is useful to represent the phase space explored by the protein in a reduced dimensional space that highlights the slow macroscopic modes of the system. A common projection for this task is represented by the Time-lagged Independent Component Analysis (TICA) (Pérez-Hernández et al.,, ). We obtain reference TICA components from the original trajectories by considering a similarity based on the Euclidean distance between Cα and a lagtime of 2ns. In order to estimate long-time probabilities, we reweigh the configurations. We first cluster the configuration using K-means in the first 4 TIC dimensions with 100 clusters. We then build a Markov State Model (MSM) on the basis these clusters by estimating the transition matrix at long time-lag (45 to 95ns). Finally, from the MSM largest eigenvectors we obtain the steady state probability of each cluster, which we use to reweigh them and approximate the Boltzmann distribution. Note that this reweighing is extremely sensitive to the correct description of the transition states, as a wrong sampling of the transition probability will exponentially affect the relative probability of different basins. As such, the correct description of these probability distributions is a good indicator of the faithfulness of the long-term dynamics. The density plot in the first two TIC dimensions for these reweighed samples is shown in Figure 5. In these plots, different regions represent different conformational states and their proximity indicates time autocorrelation. The logarithmic color scale represents the free energy in units of $k_{B}T$, and it is shifted to the minimum of each plot. We additionally estimate the free energy against $\mathbf{C}_{\alpha}$-RMSD, which we show in 6. We find that EquiJump sucessfully captures the dynamics of the 12 fast-folding proteins. While performing large steps of 100ps, our model is able to find most basins, covering their phase spaces and accurately reproducing the dynamical profiles of most proteins.

Figure: Figure 6: Free Energy on $\mathbf{C}_{\alpha}$-RMSD for the 12 Fast-Folding Proteins. We align trajectory samples to the reference crystal, and measure $\mathbf{C}_{\alpha}$-RMSDs (x-axis). Using hidden markov model weights based on our TICA-based clusters, we reweight $\mathbf{C}_{\alpha}$-RMSD counts to obtain free energy estimates (y-axis). We find that EquiJump successfully approximates the free energy curves of reference trajectories.
Refer to caption: extracted/5922300/figures/free-energy-rmsd.png

### 4.4 Model Comparison

In order to assess the quality of EquiJump, we compare it to available transferable model CG-MLFF (Majewski et al.,, ), which is based on force fields and employs Langevin sampling for dynamical generation. To the best of our knowledge, this is the only other multi-protein model that covers the 12 fast-folding proteins. We compare the learned free energy of EquiJump to that of CG-MLFF in Figures 5 and 6.

In Figure 5, we observe that EquiJump outperforms the existing model in terms of coverage and the extent of the free energy profiles for most proteins. This indicates that our model more accurately reconstructs the slow components and rare events of the dynamics, which are particularly challenging to capture due to their low frequency. This performance demonstrates the data efficiency of our approach and its ability to capture finer details of protein conformational transitions.

Notably, in some proteins (ntl9 and protein b), both models produce stationary distributions that misrepresent the location of the free energy basin. However, in both cases, EquiJump more effectively reconstructs the extent and profile of the landscape, accessing a significantly larger portion of the ground-truth basins in its dynamics. In other proteins (a3D and lambda), we observe that while CG-MLFF remains stable within the native basin, EquiJump shows some bias towards less likely states, highlighting the need for further model tuning and training to improve accuracy.

Furthermore, in Figure 6, we observe that our model outperforms CG-MLFF in replicating the free energy curves based on $\mathbf{C}\alpha$-RMSD. These curves represent the distribution of $\mathbf{C}\alpha$-RMSD after the system reaches equilibrium and are highly sensitive to the accurate estimation of conformational transitions, making them a robust evaluation metric for dynamics models. In these plots, it can be seen that our model correctly associates basin states as well as the curve profile across a wide range of RMSD values. For example, in trpcage and protein b, CG-MLFF shows large inaccuracies for states far from the basin, while our approach stays close to the ground truth across the whole profile.

## 5 Conclusion

In this work, we extended Two-Sided Stochastic Interpolants for three-dimensional dynamics simulation and applied this method to all-atom protein dynamics. To enable efficient and scalable simulations, we introduced a novel four-track SO(3)-equivariant architecture for operating on 3D all-atom protein structures. We validated our approach through a series of experiments on fast-folding proteins, demonstrating that our model accurately reproduces complex dynamics across different systems using a single model. Additionally, it outperformed other methods in capturing both short- and long-term conformational movements. Our results indicate that with further scaling and training, greater accuracy and generalizability could be achieved. Future work will focus on improving transferability and generalization to larger datasets, as well as extending EquiJump to handle more complex systems, such as ligand interactions and multimeric proteins.

## Appendix A Appendix / supplemental material

### A.1 Brownian Dynamics in Protein-Solvent Systems and Connection to Stochastic Interpolants

In the study of molecular dynamics of proteins immersed in solvents, it is crucial to account for the interactions between the proteins and the surrounding fluid. Proteins in a solvent experience random collisions with solvent molecules, leading to stochastic behavior that can be effectively modeled using Brownian dynamics (Ermak and McCammon,, ). This approach captures the random motion arising from thermal fluctuations and solvent effects, providing a realistic depiction of protein behavior in biological environments.

Generalized frictional interactions among the particles can be incorporated into the Langevin equation through a friction tensor $\mathbf{R}$ (Schlick,, ). This tensor accounts for the action of the solvent on the particles and modifies the Langevin equation to:

$$ $\mathbf{M}\ddot{\mathbf{X}}(t)=-\nabla E(\mathbf{X}(t))-\mathbf{R}\dot{\mathbf {X}}(t)+\mathbf{W}(t),$ (9) $$

where $\mathbf{M}$ is the mass matrix, $\mathbf{X}(t)$ represents the particle positions at time $t$, $E(\mathbf{X}(t))$ is the potential energy, and $\mathbf{W}(t)$ is a random force representing thermal fluctuations from the solvent. The mean and covariance of the random force $\mathbf{W}(t)$ are given by:

$$ $\langle\mathbf{W}(t)\rangle=0,\quad\langle\mathbf{W}(t)\mathbf{W}(t^{\prime})^ {T}\rangle=2k_{B}T\mathbf{R}\delta(t-t^{\prime}),$ (10) $$

where $k_{B}$ is Boltzmann’s constant, $T$ is the temperature, and $\delta(t-t^{\prime})$ is the Dirac delta function. This relation is based on the fluctuation-dissipation theorem, a fundamental result that connects the friction experienced by a particle to the fluctuations of the random force acting upon it, assuming the particle is undergoing random motion around thermal equilibrium.

This description ensures that the ensemble of trajectories generated from Eq. 9 is governed by the Fokker-Planck equation, a partial differential equation that describes the evolution of the probability density function of a particle’s position and momentum in phase space during diffusive motion.

In this context, the random force $\mathbf{W}(t)$ is modeled as white noise with no intrinsic timescale. The inertial relaxation times, given by the inverses of the eigenvalues of the matrix $\mathbf{M}^{-1}\mathbf{R}$, define the characteristic timescale of the thermal motion. When these inertial relaxation times are short compared to the timescale of interest, it is appropriate to neglect inertia in the governing equation, effectively discarding the acceleration term by assuming $\mathbf{M}\ddot{\mathbf{X}}(t)\approx 0$. Under this approximation, Eq. 9 simplifies to the Brownian dynamics form:

$$ $\dot{\mathbf{X}}(t)=-\mathbf{R}^{-1}\nabla E(\mathbf{X}(t))+\mathbf{R}^{-1} \mathbf{W}(t).$ (11) $$

This simplification reflects that solvent effects are sufficiently strong to render inertial forces negligible, resulting in motion that is predominantly Brownian and stochastic in nature. This description is particularly effective for modeling very large, dense systems whose conformations in solution are continuously and significantly altered by the fluid flow in their environment.

To stably evolve Brownian dynamics eq. 11 over time, small integration steps are usually required due to the stiffness of the physical manifold. Molecular systems exhibit a wide range of timescales: fast atomic motions such as bond vibrations and thermal fluctuations occur on the order of femtoseconds, while slower conformational shifts and larger-scale rearrangements may take place over much longer periods. These necessitate small time steps to accurately capture the system’s rapid changes without numerical instability. In contrast, Stochastic Interpolants eq. 8, while also following the form of eq. 11, enable smoothing of the data manifold by convolution with small Gaussian perturbations. This leads to a latent representation that is robust to noise, allowing for larger integration steps. The smoother manifold helps overcome local energy barriers and navigate the broader conformational landscape more efficiently, making it possible to simulate molecular dynamics on extended timescales without losing stability.

### A.2 Architecture Details

Algorithms 3, 4, 5 describe the components of EquiJump. The Tensor Square operation in Alg. 3 (line 1) is applied independently within each channel. The residual sum of Alg. 5 (line 5) is only performed on the geometric features, since positions are fixed.

Our best performing model uses hidden dimensionality 64x0e + 64x1e for the irreps, conditional number of layers $L_{\textrm{cond}}=6$, and header number of layers $L_{\textrm{header}}=4$. This parameterization is the largest we have tested, our experimentation indicates further scaling is a promising direction of research.

For interpolant parameterization we use $I(\tau,\mathbf{X}_{0},\mathbf{X}_{1})=(1-\tau)\mathbf{X}_{0}+\tau\mathbf{X}_{1}$, $\gamma(\tau)=\sigma\cdot\tau\cdot(1-\tau)$ and fixed time dependent diffusion coefficient $\epsilon(\tau)=1.0$ in sampling. Where $\sigma=3.0,1.0$ for the coordinates and geometric features, respectively. In future work we will investigate how different interpolant parameterizations affect the performance of our model.

Figure: Algorithm 3 Self-Interaction

Figure: Algorithm 4 Spatial Convolution

Figure: Algorithm 5 EquiJump Deep Network

### A.3 Dataset Details

We adapt the dataset produced by (Majewski et al.,, ). The dataset consists of trajectories of 500 steps at intervals of 100ps. Refer to the table below for the number of trajectories curated. To include all relevant residues, in addition to the standard vocabulary of residues, we also include a canonical form of Norleucine (NLE). To choose clusters that are diverse and dynamically relevant for enhanced sampling in training, we perform K-means clustering on 2D TICA components and find 200 clusters for each protein. Figure 7 visualizes cluster centers and the distribution of population sizes. Our enhanced dataset first samples a cluster, then from the cluster a configuration and its transition.

| Protein | Residues | Trajectories |
| --- | --- | --- |
| Chignolin | 10 | 3744 |
| Trp-Cage | 20 | 3940 |
| BBA | 28 | 7297 |
| Villin | 34 | 17103 |
| WW domain | 35 | 2347 |
| NTL9 | 39 | 7651 |
| BBL | 47 | 18033 |
| Protein B | 47 | 6094 |
| Homeodomain | 54 | 1991 |
| Protein G | 56 | 11272 |
| a3D | 73 | 7113 |
| x<br>$\lambda$-repressor | 80 | 15697 |

Figure: Figure 7: Cluster Centers and Distribution of Cluster Population Sizes.
Refer to caption: extracted/5922300/figures/clusters.png

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: Is it possible to check if the simulated trajectory stays on a closed manifold in the phase space and evaluate its (shadow) Hamiltonian? Would it make sense?  Please try to estimate the phase space volumes and/or energy fluctuations over time.

Please provide ProCheck/Molprobity Ramachandran plot statistics, bond lengths/angles, clash scores, and overall statistics. Please compare these to the values computed on the MD trajectories.

Please provide mathematical proofs or empirical demonstrations of equivariance for the key components of their architecture that are not directly taken from prior work, otherwise please explain the key equivariance components of the prior work, such as equivariant linear layers, etc.
soundness: 3
strengths: The network generalizes on multiple protein families and outperforms the SOTA results.
summary: The authors present a novel equivariant network that predicts  states from an MD trajectory. Presented results look convincing, but unfortunately the State of Reproducibility is not written and the model/code are not available for evaluation.
weaknesses: It will be useful to assess the stereochemical quality of the produced intermediate protein structures using standard metrics, ProCheck and/or Molprobity. 

A minor point -- It will be also useful to formally prove SO(3) equivariance of the architecture blocks (without reading the original papers).

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: See the weaknesses section above.
soundness: 3
strengths: 1. The presentation of the method is very clear and easy to understand.

2. To the best of my knowledge, the modification of the stochastic interpolants framework for generating protein dynamics is novel.

3. Employing stochastic interpolants to generate trajectories is a well-founded and rational approach.
summary: In the paper, the authors introduce the EquiJump framework to generate protein dynamics simulation trajectories. Based on the stochastic interpolants framework, the authors capture the conditional probability distribution at different time to generate the whole trajectory. Additionally, the paper extend the interpolants on geometric features and design new network architecture. Experiments on 12 fast-folding proteins show the empirical performance of the framework.
weaknesses: 1. I believe it is essential to conduct comparative analyses with other baseline methodologies to establish the efficacy of your approach. For instance, utilizing a Machine-Learned Force Field (MLFF) that incorporates your network architecture could serve as a viable baseline. Such a comparison could demonstrate that your interpolants method outperforms the MLFF, independent of the network design, thereby highlighting the strengths of your technique.

2. I have observed that the sampling process within the current framework is rather inefficient and time-consuming. Therefore, it would be beneficial to include a comprehensive comparison of computational costs between the EquiJump framework and preceding MLFF methods. Such an analysis would be instrumental in assessing the practicality and efficiency of the EquiJump framework in relation to established techniques.

3. It is also crucial to engage in a detailed discussion on contemporary methods of trajectory generation. References [1-4] represent concurrent studies in trajectory generation. Providing a clear delineation of the distinctions and relationships between your work and these existing studies will help elucidate the unique contributions and advancements your research offers to the field.

[1]. Du, Yuanqi, et al. "Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling." arXiv:2410.07974 (2024).

[2]. Zhang, Xi, et al. "Trajectory Flow Matching with Applications to Clinical Time Series Modeling." arXiv:2410.21154 (2024).

[3]. Luo, Shengjie et al. “Bridging Geometric States via Geometric Diffusion Bridge.” arXiv:2410.24220 (2024).

[4]. Han, Jiaqi, et al. "Geometric Trajectory Diffusion Models." arXiv:2410.13027 (2024).

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Clarify the training-test-validation splits.
Comment on performance comparison approach more quantitatively (not needing new experiments).
Model size comparison.
Comment on what the specific rotation representation is, and how does it actually help. Also, some insight on how the network archietcture for EquiJump can be motivated.
soundness: 2
strengths: A substantive assessment of the strengths of the paper, touching on each of the following dimensions: originality, quality, clarity, and significance. We encourage reviewers to be broad in their definitions of originality and significance. For example, originality may arise from a new definition or problem formulation, creative combinations of existing ideas, application to a new domain, or removing limitations from prior results. You can incorporate Markdown and Latex into your review. See https://openreview.net/faq.

The proposed architectural novelty is high with clear descriptions of the new elements. Although I suggest in figure 4, where the architecture is drawn to use consistent naming across the various panes, e.g. what is called the ‘deep network’ in panel (c) becomes the ‘conditioner’ in panel (d) …. Also two panel c’s appear in the figure etc.

The experiments are done on a dataset of 12 folding proteins, which appear to a standard point of comparison. 

I appreciated the representation of molecular data using geometric features, although it was unclear what the irreducible representation is here (Lie algebraic coordinates?).
summary: This paper proposes new methods for speeding up protein dynamics simulation using neural networks. The core novelty lies in the development of novel architectures termed EquiJump. This network is shown to recover protein dynamics as measured by trajectory error in 3D as well the evolution of free energy metrics over time.
weaknesses: Although the architecture is novel, it is less clear how it is motivated, and why this particular combination works. This also ties in to comparisons, as only one primary comparison point is used – CGMLFF.
Rotation data representation was unclear – ‘irrep’ would benefit by a clearer definition. Does using special approaches like this, does it reduce the parameterization of the rest of the network?


It would be also important to consider model size when comparing to other models, e.g.  total number of parameters of the network for CGMLFF is noted as 294,565 in (Majewski 2023). Would also be useful to see training curves for the models. 

Training-validation-test splits are not clearly specified from what I can tell. (Majewski 2023) states the data was randomly split between training (85%), validation (5%), and testing (10%). 

Improvements in performance are mostly left as visual assessment, e.g. figure 6 and figure 5. It is unclear how one would conclude anything just by looking at these figures. Multiple quantitative ways to compare seem feasible, e.g. see supplementary docs in (Majewski 2023).

https://static-content.springer.com/esm/art%3A10.1038%2Fs41467-023-41343-1/MediaObjects/41467_2023_41343_MOESM1_ESM.pdf

---

confidence: 5
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. Can you provide more explanations about one-side and two-side stochastic interpolants? It is difficult to understand the contribution of your model if one has not much knowledge about it.
2. How many orders do you use for the feature representation $V$ ? And how do you initialize $V^0$ ($l=0$)?
3. Are there some trainable parameters of $I_\tau$?  From algorithm 1,  it seems that $I_\tau$ is just a linear interpolant, without anything to train.
soundness: 3
strengths: 1. The proposal of  two-sided stochastic interpolants is inspiring and interesting. It enables the model to make sequence generation ($X^t\rightarrow X^{t+1}$), addressing the limitation of common (one-sided) stochastic interpolants, which are restricted to generating samples only from a specified prior distribution (e.g., a Normal Distribution).
2. The paper provides theoretical guarantee of EquiJump, which is verified by  the experimental results on fast-folding proteins. It successfully generates the  trajectories with longer timesteps. Moreover, the representation and visualization of the results are wonderful.
summary: This paper proposes EquiJump for protein dynamics simulation. It utilizes two-sided stochastic interpolants to solve the issue of mismatch between  prior distribution and data distribution. The model is designed to meet SO(3) equivariance to capture the dynamics physically, and is evaluated on 12 fast-folding proteins, achieving state-of-the-art performance compared with CG-MLFF.
weaknesses: 1. There are some typos:
  + In Eq. (6), the target to minimize should be $\hat{\eta}$，not $\hat{s}$
  + In line 94, "represents" and "builds" are repeated. Maybe you want to choose either of them?
  + In figure(3),  there are two subfigures (c). And in the final subfigure, the initial step is supposed to be $X^t$， not $X^\tau$.
2. The experiments can be more comprehensive:
  + More baselines. EquiJump is only compared with CG-MLFF. It will be more convincing to compare with predictive models, such as EGNN [1], EqMotion [2] and ESTAG[3]. They are recent equivariant graph neural networks (GNNs) and have been used  to predict molecular dynamics. 
  + More representations. The experimental results are presented by plots and curves. Detailed numerical results (presented by table) can help us understand the priority of your model better.  For example, you can present the RMSD between your generated conformations and the ground truth structures.
  + More metrics. Jensen-Shannon divergence (JSD) is also used as a metric for trajectory prediction.
  + More datasets. I am curious about the performance of your model on the MDAnalysis [4] dataset, too.  It is also a prevailing protein dynamics dataset used for evaluating ML model.
  + Ablation on steps. Currently, you set the integration steps to 500. It is interesting to see the performance of fewer steps (e.g., steps=200/300/400).

[1] Satorras, Vıctor Garcia, Emiel Hoogeboom, and Max Welling. "E (n) equivariant graph neural networks." International conference on machine learning. PMLR, 2021.

[2] Xu, Chenxin, et al. "Eqmotion: Equivariant multi-agent motion prediction with invariant interaction reasoning." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[3] Wu, Liming, et al. "Equivariant spatio-temporal attentive graph networks to simulate physical dynamics." Advances in Neural Information Processing Systems 36 (2024).

[4] Michaud‐Agrawal, Naveen, et al. "MDAnalysis: a toolkit for the analysis of molecular dynamics simulations." Journal of computational chemistry 32.10 (2011): 2319-2327.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 4
questions: 1. Please give more explanations of the neccesity to propose a novel $\mathrm{SO}(3)$-equivariant architecture rather than using existing architecutres like Equiformer. Ablation results will be better.
2. To prepare the training dataset, the resampling technique was used to uniformly sample states with different free energies. I wonder that while the dataset was biased from the Boltzmann distribution, the minimizer of the training objective might also be biased from the real dynamics. Please correct me if I am wrong.
3. Please provide more experimental results of stronger baselines by training from scratch if possible. It would be more appreciated if the results could be reported through statistical metrics.
soundness: 3
strengths: 1. A good try to leverage two-sided stochastic interpolants for MD, which can be considered as a Markov process and therefore serves as a suitable scenario for bridging the arbitrary data endpoints.
2. The work proposes the tensor cloud representation and a corresponding $\mathrm{SO}(3)$-equivariant module, obtaining a more detailed representation of 3D coordinates.
3. Inspired by classical MD methods like umbrella sampling, the work utilizes the resampling technique to uniformly sample different states from MD trajectories for training.
summary: This work focuses on enhancing molecular dynamics of proteins and proposes EquiJump, a two-sided stochastic interpolant framework tailored for arbitrary data endpoints with $\mathrm{SO}(3)$-equivarient modules. Experiments on 12 fast-folding proteins verify its superiority by accurately replicating the dynamics.
weaknesses: 1. This work is slightly lacking in innovation: the concept and conclusions of the two-sided stochastic interpolant were elucidated in [1], and the $SO(3)$-equivariant architecture of EquiJump is similar to other existing architectures as well, such as Equiformer [2]. Further explanations of the neccesity to use such architectures are warranted.
2. EquiJump was trained and evaluated both on the same dataset (fast-folding proteins). Although the dynamics are restored accurately, there's no proof for the transferability to other protein systems of the model.
3. EquiJump was only compared to CG-MLFF in this work, where CG-MLFF is performed in a corase-grained fashion, leading to a relatively unfair comparison. It would be more convicing to compare with more advanced models that work on all-atom systems, such as Timewarp [3]. Meanwhile, providing statistical results will be better.

**Reference**

[1] Albergo, M. S., Boffi, N. M., & Vanden-Eijnden, E. (2023). Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797.

[2] Liao, Y. L., & Smidt, T. (2022). Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206.11990.

[3] Klein, L., Foong, A., Fjelde, T., Mlodozeniec, B., Brockschmidt, M., Nowozin, S., ... & Tomioka, R. (2024). Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36.

Meta-review: The paper is on using machine learning, particularly generative modeling, for molecular dynamics. It uses Albergo et al’s (ICML 2024) stochastic interpolants on SO3 equivariant models to simulate 3D conformation state transfer of an MD for 12 proteins with relatively long time steps. 

The reviewers found the presentation generally good despite some typos and unclarities, the proposed method reasonable for the stated purpose, the results significant compared to one included MLFF baseline.

On the other hand, they were concerned about the quality of the produced trajectory, lack of enough quantitative scalar measures, lack of a comprehensive conceptual and empirical comparison with other recent works on modeling MD trajectories including recent MLFF works or other generative modeling approaches, and they requested more diverse set of benchmark and evaluations including for out-of-distribution proteins, as well as additional side and ablation studies including a separation of backbone model/representation and the specific generative model.

The authors provided a thorough rebuttal and major update of the paper which improved several aspects including more empirical evaluation of the chemical validity of the generated individual states and full trajectory, a formal result of equivariance of the used model,  a more detailed and comprehensive conceptual comparison with other recent works, tables with quantitative single scalar metrics for the quality of the generated trajectory, and some additional basic baselines such as DDPM. 	

The AC believes the paper had clear merits in its initial form and has further considerably improved during the rebuttal. However, the AC recommends rejection for two main reasons: (1) the paper’s empirical evidence needs to be strengthened: since the main contribution of the paper is the use of a more proper modeling with stochastic interpolants for a specific application of long-term prediction of MD trajectories, it needs to (a) find a way to compare with other well-designed generative models to show the importance of the arbitrary prior and (b) evaluate on more benchmarks including OoD data to enable a solid progress in the new line they open for ML modeling of long steps with large molecules. (2) the paper requires another round of review: the paper has been majorly updated during the rebuttal and needs some more additions, therefore it requires a fresh round of review of the details of all new textual and formal discussions, evaluations, and empirical results.

Decision: Reject