Align Your Structures: Generating Trajectories with Structure Pretraining for Molecular Dynamics

TL;DR

This paper introduces a structure pretraining framework for molecular dynamics trajectory generation, combining a diffusion-based structural model with an interpolator to improve realism and temporal consistency.

Contribution

It proposes a novel two-stage approach leveraging structural pretraining and an interpolator to enhance MD trajectory generation from limited data.

Findings

Outperforms existing methods in geometric accuracy

Achieves realistic dynamical and energetic measurements

Effective across small molecules, peptides, and proteins

Abstract

Generating molecular dynamics (MD) trajectories using deep generative models has attracted increasing attention, yet remains inherently challenging due to the limited availability of MD data and the complexities involved in modeling high-dimensional MD distributions. To overcome these challenges, we propose a novel framework that leverages structure pretraining for MD trajectory generation. Specifically, we first train a diffusion-based structure generation model on a large-scale conformer dataset, on top of which we introduce an interpolator module trained on MD trajectory data, designed to enforce temporal consistency among generated structures. Our approach effectively harnesses abundant structural data to mitigate the scarcity of MD trajectory data and effectively decomposes the intricate MD modeling task into two manageable subproblems: structural generation and temporal alignment. We comprehensively evaluate our method on the QM9 and DRUGS small-molecule datasets across unconditional generation, forward simulation, and interpolation tasks, and further extend our framework and analysis to tetrapeptide and protein monomer systems. Experimental results confirm that our approach excels in generating chemically realistic MD trajectories, as evidenced by remarkable improvements of accuracy in geometric, dynamical, and energetic measurements.