Generative Modeling of Molecular Dynamics Trajectories

Molecular dynamics (MD) is a powerful technique for studying microscopic phenomena, but its computational cost has driven significant interest in the development of deep learning-based surrogate models. We introduce generative modeling of molecular trajectories as a paradigm for learning flexible multi-task surrogate models of MD from data. By conditioning on appropriately chosen frames of the trajectory, we show such generative models can be adapted to diverse tasks such as forward simulation, transition path sampling, and trajectory upsampling. By alternatively conditioning on part of the molecular system and inpainting the rest, we also demonstrate the first steps towards dynamics-conditioned molecular design. We validate the full set of these capabilities on tetrapeptide simulations and show that our model can produce reasonable ensembles of protein monomers. Altogether, our work illustrates how generative modeling can unlock value from MD data towards diverse downstream tasks that are not straightforward to address with existing methods or even MD itself. Code is available at https://github.com/bjing2016/mdgen.

Paper

References (53)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer b8Jp7/10 · confidence 2/52024-07-10

Summary

The paper suggests a flow-based generative framework on molecular trajectories, with various downstream tasks such as forward simulation and transition path sampling. Additionally, the model is trained in a transferable setting, across tetrapeptides.

Strengths

1. Extensive experiments over various downstream tasks 2. Transferable settings for tetrapeptides

Weaknesses

1. Experiment baselines The baselines of experiments are mostly the Markov State Models. I think it would also be good if there were some comparison between other models, though I understand that many prior works targeted Alanine dipeptide not tetrapeptides. - Forward simulation: ITO$^{[1]}$, Timewarp$^{[2]}$ - Interpolation (Transition path sampling): PIPS$^{[3]}$ 2. (Minor) Necessity of additional tasks The necessity of additional tasks relatively seems weak compared to tasks such as forward simulation, TPS, specially the inpainting design. Rather than additional tasks, ablations for stability might be a better? One can obviously see that scaling to long trajectories shows the stability against the time scale, and protein simulation shows the stability against space complexity. **Minor typos, suggestions** - Definition of S-MPNN only exists in the Appendix. It would great to point out that more details are in the appendix, in the first paragraph of section 4.4 - Figure 6 is not referenced in the main paper, only the appendix - Figure 2F, reference of that blue indicates the side chains and orange indicates the backbones seems to be missing [1] Implicit transfer operator learning: Multiple time-resolution surrogates for molecular dynamics, NIPS 2023 [2] Timewarp: transferable acceleration of molecular dynamics by learning time-coarsened dynamics, NIPS 2023 [3] Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths, NIPS 2023

Questions

1. Difference between downstream tasks Could the upsampling task be seen as a superset of interpolation? Since upsampling with two given frames would the same as interpolation. 2. Training on one tetrapeptide Just curious, though the authors has presented a transferable setting, are there any results when the model is trained for a specific tetrapeptide and tested on downstream tasks?

Rating

7

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

1. Unconditional generation As the authors have mentioned, unconditional generation is impossible since the model relies on key frames.

Reviewer hVZ77/10 · confidence 4/52024-07-10

Summary

The authors propose MDGen -- a generative model to sample molecular dynamics trajectory conditioned on key frames. This is a direct application of video generation techniques to solve domain challenges in protein modeling. Specifically, SiT and flow matching models are used to sample SE(3)-invariant representation of all-atom protein representations. This work demonstrates the effectiveness of MDGen primarily on tetrapeptide systems, where the authors showcase four downstream application tasks including forward simulation, interpolation, upsampling, and dynamics-conditioned inpainting. In general, I find this manuscript well-written and easy to follow. The model performance looks reasonable on tetrapeptides, yet the results are proof-of-concept in nature and generalization to larger proteins remain challenging. However, it is one of the pioneering work in AI protein modeling to directly emulate MD simulation trajectories using data-driven approaches. To that end, I think it would be beneficial for this work to gain visibility across the research community to inspire future studies.

Strengths

- It is one of the pioneering work to adopt video generation techniques for MD trajectory generation. Although conceptually straightforward, good practices to generate time-coherent MD trajectories across different protein systems remain underexplored. - The authors demonstrated a variety of downstream tasks using the same model architecture. The underlying modeling framework seems versatile and transferrable across different applications.. - Performance benchmark and analysis on tetrapeptides are comprehensive and provides insights to modeling these peptide systems. - I think it is a good idea to model residue offsets relative to the key frames in order to bypass the need to learn sequence-to-structure mapping. MD simulations always start from a seed structure, so I do not think this is a key limitation as mentioned in L#310-312.

Weaknesses

- Benchmark and evaluation results on tetrapeptides, although comprehensive, are proof-of-concept in nature. It may not be sufficient to demonstrate transferability to general protein systems. - Performance on ATLAS (i.e., larger proteins instead of short peptides) does not seem promising. MDGen performance is worse than AlphaFlow in Table 4. I wonder if the main bottleneck is training data quality/availability, or model architecture?

Questions

- L#141, when $K > 1$, how to ensure roto-translation prediction consistency across $K$ key frames and obtain a final $\hat{g}_j$? - Table 2, with 100 ns being the ground truth, the non-zero JSD in the last column originates from subsampling the simulation trajectory? - Figure 2F. My understanding is that sidechains exhibit faster dynamics while backbone motions are slower. The low correlation for backbone suggests that MDGen is not good at learning slower dynamics, which are typically more interesting to researchers? - Temporal coherence between generated protein conformations is mainly evaluated using auto-correlation in this work. Is it possible to show other metrics to capture detailed structural quality and variation during time evolution? - Why is MDGen more effective at sequence recovery than MPNN? More explanation and analysis would be helpful here. - Would it be possible to emulate MD simulation trajectory of the 12 fast folding proteins from [Shaw 2009](https://dl.acm.org/doi/abs/10.1145/1654059.1654126)? They are smaller than ATLAS proteins and longer than tetrapeptides, with much longer simulation time and rich dynamics. - It would be nice to see if MDGen could infer a trajectory given an [apo/holo pair](https://arxiv.org/abs/2304.02198).

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See weaknesses and questions.

Reviewer zaG34/10 · confidence 4/52024-07-12

Summary

The paper presents a new framework for generating trajectory of molecular geometries, ie, generative modeling for molecular dynamics. The paper proposes tokenization methods to tokenize the trajectory and learn flow models on the data. Experiments demonstrate the effectiveness of several tasks including forward sampling, interpolation, and up sampling.

Strengths

1. The paper tackles a new problem in molecular dynamics generation, which has not been explored in existing literature. 2. The paper is in good structure and easy to follow. 3. The paper provides a detailed analysis of several domain tasks on interested molecular structures, which demonstrate the critical usage in some scenarios.

Weaknesses

1. Limited ML technical contribution, as all components exist in previous molecular generative models. 2. The experiment is comprehensive from a domain perspective. However, I feel the experiments lack some benchmarking comparison with state-of-the-art molecular generative models for related tasks. See my question below.

Questions

I think existing methods can also tackle several tasks. For example, for the forward sampling task, previous generative MD models like Timewarp (Klein et al., 2024) and ITO (Schreiner et al., 2024) can also be used for the task. A numerical comparison with these baselines can help to justify the effectiveness of the proposed method.

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Limitation is nicely discussed in the paper.

Reviewer b3SL6/10 · confidence 3/52024-07-14

Summary

In this work, the authors proposed MDGen, a new framework that aims to model molecular dynamics trajectories via generative modeling techniques. By properly encoding the Protein MD trajectories according to the characteristics of key frames, MDGen adopts flow matching techniques (both continuous and discrete flow matching) to generatively model MD trajectories. As a unified framework, MDGen is able to perform diverse tasks including forward simulation, interpolation, upsampling and inpaiting. Extensive experiments are conducted to demonstrate the effectiveness of MDGen.

Strengths

1. The problem this work aims to tackle is of great significance in scientific domains lie computational biology. 2. The formulation of molecular (protein) trajectories by using key frame references is reasonable and compact for reducing the modeling difficulties. 3. The experiments are comprehensive. 4. The paper is well-written and easy to follow.

Weaknesses

1. Lack of discussion on related works. This work does not discuss related works on the same topic. Some works are mentioned in the Introduction section, but I still recommend that there should be an independent Related Works section for comprehensive discussion. Here are also several works that are worth discussing: (1) EGNO, which uses neural operator learning approach to also model the trajectory dynamics of molecules; (2) DiffMD, which uses diffusion models to simulate molecular dynamics. The quality of this work should be further improved if the authors could carefully discuss the differences between MDGen and these works and the strengths of MDGen compared to these works. 2. Lack of ablation studies. MDGen is composed of several parts, including the design of the backbone model, the design choices of flow matching framework, and the adoption of Hyena architecture for efficiency consideration. In addition to the aimed tasks, it would further improve the quality of this work if the authors could conduct ablation studies on these aspects to help readers know what the influence of each part of MDGen is.

Questions

N/A

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors carefully discuss the limitations of this work.

Reviewer TTLp5/10 · confidence 3/52024-07-20

Summary

The paper presents a novel generative model for molecular dynamics (MD) trajectories called MDGEN. This model aims to serve as a flexible surrogate for MD simulations by generating entire trajectories conditioned on initial frames. It addresses tasks such as forward simulation, transition path sampling, trajectory upsampling, and dynamics-conditioned molecular design. The model is evaluated on tetrapeptide simulations and demonstrates its capability to generate reasonable ensembles of protein monomers.

Strengths

Novelty and Scope: The approach introduces a novel paradigm for surrogate modeling of MD, extending the capabilities of existing models to handle a variety of tasks that are not straightforward with current methods. Generative Framework: The use of generative modeling for entire MD trajectories is a significant advancement, as it allows for a broader range of applications including forward and inverse problems. Comprehensive Evaluation: The paper evaluates MDGEN on several tasks, demonstrating its effectiveness in forward simulation, interpolation, upsampling, and inpainting. The results show promising performance in terms of distributional similarity, dynamical content, and computational efficiency. Technical Implementation: The detailed description of the tokenization process and the flow model architecture provides a clear understanding of how the model operates. The use of SE(3)-invariant tokens and the scalable interpolant transformer (SiT) backbone are well-motivated choices.

Weaknesses

Complexity and Accessibility: The model’s complexity might pose challenges for reproducibility and accessibility for researchers who are not deeply familiar with both molecular dynamics and advanced generative modeling techniques. Evaluation on Larger Systems: While the paper provides proof-of-concept evaluations on proteins, the primary focus remains on smaller tetrapeptides. The model's scalability and effectiveness on larger and more complex molecular systems need further exploration. Dependence on Key Frames: The reliance on key frames for conditional generation limits the model’s ability to perform unconditional generation or inpainting of residue roto-translations, which could be a significant limitation in certain applications. Computational Resources: The paper lacks detailed information on the computational resources required for training and inference, which is crucial for understanding the practical implications of using MDGEN in various research settings.

Questions

How can the model be adapted or improved to reduce its reliance on key frames? Exploring techniques for unconditional generation or alternative ways to handle the roto-translations without predefined key frames could enhance the model's flexibility and applicability. What architectural changes or enhancements could improve the model's performance on larger molecular systems such as proteins? Investigating more scalable architectures or hybrid approaches that combine the current method with other techniques tailored for large systems could address this limitation. How does the computational cost of training the model compare to traditional MD simulations, and what are the implications for its practical use? Providing detailed information on computational requirements and potential optimizations could help in assessing the model's feasibility for widespread use. What alternative tokenization strategies could be explored to extend the model's applicability to a wider range of molecular systems? Research into tokenization methods that can handle diverse molecular structures and dynamics could broaden the model's utility. How can additional conditioning types (e.g., textual descriptions, experimental data) be incorporated into the model, and what benefits might they provide? Experimenting with and integrating various forms of conditioning could enhance the model's ability to generate more accurate and contextually relevant trajectories. What are the potential impacts of data quality and availability on the model's performance, and how can these challenges be mitigated? Addressing data-related challenges through techniques like data augmentation, transfer learning, or synthetic data generation could improve the model's robustness and applicability. Can additional evaluation metrics be developed to provide a more comprehensive assessment of the generated trajectories' quality? Identifying and implementing new evaluation criteria could offer deeper insights into the strengths and limitations of the model's output.

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

Reliance on Key Frames: The model relies on key frames for parameterizing the roto-translations, which means it cannot perform unconditional generation or inpainting of residue roto-translations. This dependency might limit its applicability to scenarios where key frames are not easily obtainable or where full trajectories need to be generated from scratch. Scalability to Larger Systems: The architecture shows weaker performance on larger systems such as proteins compared to smaller peptides. This suggests that the current model and architecture might not be well-suited for handling the complex motions and larger size of protein structures without further modifications or enhancements. Computational Resources: While the paper mentions significant speedups compared to traditional MD simulations, the computational resources required for training the model (e.g., GPU hours) are not explicitly discussed. This information is crucial for understanding the practicality and scalability of the approach. Generalization to Diverse Systems: The current tokenization and modeling strategies are tailored to peptides and proteins. For more diverse molecular systems such as organic ligands, materials, or explicit solvent systems, alternative tokenization strategies might be necessary. This limits the immediate applicability of the model to a broader range of molecular simulations. Limited Exploration of Additional Conditioning: The paper primarily explores conditioning on initial frames and residue identities. Other types of conditioning, such as textual or experimental descriptors, are not explored but could open up further applications and improve the model's utility. Data Availability and Quality: The success of the model heavily depends on the availability of high-quality MD trajectory data. For many complex systems, obtaining such data can be challenging, which could limit the model's applicability and performance. Evaluation Metrics: While the paper uses several rigorous metrics for evaluation, the choice of metrics may not fully capture all aspects of the generated trajectories' quality. Additional metrics or more diverse evaluation criteria could provide a more comprehensive assessment.

Reviewer QAiG7/10 · confidence 4/52024-07-24

Summary

The authors introduce MDGen as a novel approach for modeling MD trajectories. They demonstrate the capabilities of this method in tasks such as interpolation, upsampling, and inpainting of small peptides. The accuracy as well as speed of the new approach compared to the ground truth baseline is quantitatively evaluated. Initial experiments toward upscaling to small proteins are shown.

Strengths

The idea of MDGen is novel and very well presented in this manuscript. The results are convincing and interesting.

Weaknesses

1. Parts of Sections 3.1 and 3.2 are very condensed and hard to follow. A more detailed description in the SI would be helpful, where the most important aspects of the cited work is also repeated. 2. The suitability of the chosen representation for longer amino acid chains is questionable. This is also mentioned in the manuscript, but nonetheless, proteins are mentioned many times (more than 30) in the manuscript, while almost all experiments are actually performed on very small peptides. It should be stated in a more prominent place that upscaling to proteins is not trivial. 3. The representation limits the model to learn MD trajectories of natural amino acids, as no all-atom representation is used directly. This should be made clearer in the manuscript. Minor points: A lot of figures have no proper axis labels (e.g. Fig 3, 4, 5, 6). This should be fixed. The best models in Table 4 should be indicated in bold.

Questions

1. How often do clashes and other high-energy structures occur in the generated trajectories? 2. When comparing to other methods and approaches in the experimental section - do all of them use a similar reduced representation or do the other methods generate all-atom representations?

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

The approach is limited to peptides. Transfer to any other molecules is questionable due to a lack of suitable representation/tokenization.

Reviewer b8Jp2024-08-08

I thank the authors for the detailed response and additional experiments. Additional experiments on forward simulation shows better performance of MDGen over Timewarp and ITO, where prior works fail in longer simulation time. All questions have been resolved, and I have raised my score accordingly.

Reviewer QAiG2024-08-08

Thank you for the clarifications. This answers all my questions and I encourage that all updates announced here will be integrated in the final version. As the changes do not change the main message and claims, I will keep my score as it is, so I continue to support acceptance of the paper.

Reviewer hVZ72024-08-12

Response to Rebuttal

Dear authors, Thanks for your response. Most of my concerns have been properly addressed. I think this pioneering work on generative modeling for protein/peptide dynamics should be shared with the community to inspire follow-up work. I have raised my scores accordingly.

Reviewer TTLp2024-08-12

Thanks for the clarifications

Some of my concerns have been addressed. I raised my score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC