Generative Modeling of Molecular Dynamics Trajectories (2409.17808v1)

Abstract: 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.

• The paper introduces MDGen, a deep generative framework that models full MD trajectories using tokenization and a SiT architecture, achieving up to 1000x speedup.
• It employs stochastic interpolants and SE(3)-invariant tokenization to capture spatial and temporal dynamics, enabling tasks like forward simulation and interpolation.
• The framework’s versatility enhances molecular mechanism hypothesis generation and design, offering advanced capabilities for computational chemistry.

Generative Modeling of Molecular Dynamics Trajectories: An Expert Overview

The paper "Generative Modeling of Molecular Dynamics Trajectories" presents an innovative approach to surrogate modeling in molecular dynamics (MD) using deep generative models. The authors propose MDGen, a framework that leverages generative modeling to simulate entire molecular dynamics trajectories. This extends beyond traditional methods that emulate either only the transition densities or the equilibrium distribution.

Core Contributions

MDGen introduces several novel methodologies, all centered around the modeling of full trajectories:

1. Simulation Tasks: By conditioning on certain frames, MDGen can perform forward simulation, interpolation (transition path sampling), upsampling, and inpainting. This positions the framework as flexible and applicable across a broad range of molecular simulation tasks.
2. Tokenization and Model Architecture: The model tokenizes molecular trajectories into $SE(3)$-invariant units, utilizing a Scalable Interpolant Transformer (SiT) framework to handle the time-series data. By doing so, it efficiently handles both the spatial and temporal dimensions inherent in molecular simulations.
3. Use of Stochastic Interpolants: Integrating stochastic interpolants allows for end-to-end training of flow-based models, facilitating the generation of realistic, diversified trajectory data.

Results and Evaluations

MDGen's performance is evaluated on several fronts:

• Forward Simulation: The model can accurately reproduce free energy surfaces and dynamic qualities like torsional relaxation, with speedups of 10x-1000x compared to traditional MD simulations.
• Interpolation: MDGen demonstrates proficiency in generating transition paths that align well with those identified in Markov State Models (MSMs). The model's generated paths exhibit higher likelihoods and improved validation over traditional MSM-sampled paths from short simulations.
• Upsampling: The method effectively reconstructs fast dynamics from trajectories sampled with longer time intervals, recovering information otherwise lost.
• Inpainting for Design: MDGen shows improved sequence recovery over baselines when generating inner residue designs consistent with observed dynamics.

Implications and Future Directions

The potential applications of MDGen are substantial:

• Mechanism Hypothesis Generation: The ability to interpolate trajectories can help propose molecular mechanisms for given end states, useful for studying rare events in molecular processes.
• Molecular Design: Inpainting offers possibilities in designing molecules with desired dynamic properties, for example, enhancing enzymatic action or specific protein-protein interactions.
• Scale and Generalization: The architecture, particularly with enhancements like Hyena operators for long trajectories, indicates promising scope for scaling to larger biological systems, including multi-chain proteins or protein complexes, thereby pushing the boundaries of protein simulation.

Conclusion

By broadening the application scope of generative models to full molecular trajectories, MDGen introduces a versatile tool in computational chemistry and molecular biology. Its ability to handle both forward and inverse simulation tasks, alongside significant computational efficiency, makes it a promising framework for accelerating discoveries in molecular sciences. Future work will likely focus on architectural improvements and extending methodologies to diverse, more complex systems. This paper's contributions mark a significant step towards integrating deep learning frameworks into practical molecular dynamics applications.