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

Authors: Allan Dos Santos Costa, Ilan Mitnikov, Franco Pellegrini, Ameya Daigavane, Mario Geiger, Zhonglin Cao, Karsten Kreis, Tess Smidt, Emine Kucukbenli, JOSEPH JACOBSON
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: biomolecules, proteins, molecular dynamics, generative models, neural transport, stochastic interpolants

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.

Decision: reject
Reviewer scores: 6, 6, 3, 6, 5

Reviews and discussion:
<thread>
creation date:
1729577514174

thread_type:
official_review

<thread_object>
creator: reviewer_zvv7

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729588323247

thread_type:
official_review

<thread_object>
creator: reviewer_sfuf

<content>
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.

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.

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.

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1729992100619

thread_type:
official_review

<thread_object>
creator: reviewer_hnmf

<content>
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.

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?).

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

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730545603435

thread_type:
official_review

<thread_object>
creator: reviewer_jvpj

<content>
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.

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.

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).

questions:
See the weaknesses section above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730670163502

thread_type:
official_review

<thread_object>
creator: reviewer_hezm

<content>
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.

strengths:
The network generalizes on multiple protein families and outperforms the SOTA results.

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).

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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732910869731

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revision Changes

comment:
We express our sincere gratitude to all reviewers for their valuable feedback. We have diligently addressed concerns and implemented the suggested improvements, significantly enhancing the manuscript. We summarize key changes in the revised manuscript below:

- **Comparison of Generative Transport Methods**: in new Section 4.3 we compare our model’s generative framework to other generative transport methods (Figure 2) [1, 2, 3]. We quantitatively evaluate the performance of these transport models across noise scales (Table 1) and qualitatively compare their free energies (Figure 4). Our results demonstrate the efficacy of our proposed approach to MD snapshot data. 
- **Model Size Ablation and Quantitative Comparisons**: in updated Section 4.4 and Appendices 8 and 9 we now ablate our model performance across model capacities and against CG-MLFF [4] with extensive quantitative measures. These experiments demonstrate the notable precision of our model in large-capacity regimes.
- **Analysis of Chemistry**: in Appendix 6 we evaluate the performance of EquiJump in generating chemically valid data across different proteins. We show the distribution of measures for chemical validity (bond length, bond angles, atomic clashes) for model samples and reference data, concluding that our model performs chemically stable simulation.
- **Performance Estimates**: in Section 4.4.1 we estimate simulation speedups relative to the original MD method that created the training data, while discussing performance of comparable models. We study the behavior of our approach across capacities, and find that EquiJump enables simulation acceleration with small loss of long-term dynamics accuracy.
- **Sampling Parameters Ablation**: in Appendix 11 we now investigate the impact in sampling of changing transport hyperparameters (number of steps, noise schedule). Our results are promising for further reduction of the number of sampling steps. 
- **Structural Metrics**: following revision suggestions, we added diverse metrics (RMSD, GDT-TS, Radius of Gyration, Fraction of Native Contacts) for quantifying the performance of our model in relevant structural properties, and now provide strong statistical comparisons (Jensen-Shannon Divergence) for different model benchmarks. The summarized results described above make extensive use of those scores through tables 1, 2, 3, 4, and 6.
- **Expanded Related Work**: we further developed Section 2 with updated elaboration about the discussed similar work [5-11], highlighting their contributions and primary differences to our approach.
- **Expanded Stochastic Interpolants**: we added further description on One-Sided Interpolants in Section 3.1, and added it to the generative comparison of Section 4.3.
- **Additional Free Energy Curves**:  we provide additional estimated free energy curves for Radius of Gyration and Fraction of Native Contact in Appendix 10, Figures 13 and 14.  

We thank the reviews again for the constructive discussion, and hope the new updates and results can demonstrate the utility of our work to the field. 

> [1] Lipman, Yaron, et al. "Flow matching for generative modeling." arXiv preprint arXiv:2210.02747 (2022).
>
> [2] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in neural information processing systems 33 (2020): 6840-6851
>
> [3] Albergo, Michael S., Nicholas M. Boffi, and Eric Vanden-Eijnden. "Stochastic interpolants: A unifying framework for flows and diffusions." arXiv preprint arXiv:2303.08797 (2023).
>
> [4] Majewski, Maciej, et al. "Machine learning coarse-grained potentials of protein thermodynamics." Nature communications 14.1 (2023): 5739.
>
> [5]. Du, Yuanqi, et al. "Doob's Lagrangian: A Sample-Efficient Variational Approach to Transition Path Sampling." arXiv:2410.07974 (2024).
>
> [6]. Zhang, Xi, et al. "Trajectory Flow Matching with Applications to Clinical Time Series Modeling." arXiv:2410.21154 (2024).
>
> [7]. Luo, Shengjie et al. “Bridging Geometric States via Geometric Diffusion Bridge.” arXiv:2410.24220 (2024).
>
> [8]. Han, Jiaqi, et al. "Geometric Trajectory Diffusion Models." arXiv:2410.13027 (2024).
>
> [9] Satorras, Vıctor Garcia, Emiel Hoogeboom, and Max Welling. "E (n) equivariant graph neural networks." International conference on machine learning. PMLR, 2021.
>
> [10] 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.
>
> [11] Wu, Liming, et al. "Equivariant spatio-temporal attentive graph networks to simulate physical dynamics." Advances in Neural Information Processing Systems 36 (2024).
</content>
</thread_object>

</thread>


Metareview:
metareview:
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.