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

Aniketh Iyengar; Jiaqi Han; Pengwei Sun; Mingjian Jiang; Jianwen Xie; Stefano Ermon

cs.LGApr 5, 2026

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

Authors: Aniketh Iyengar, Jiaqi Han, Pengwei Sun, Mingjian Jiang, Jianwen Xie, Stefano Ermon

Organizations: Stanford University · Lambda, Inc

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.

Figures & tables

Figure 1: The overall two-stage framework of EGInterpolator . Structure pretraining: We first pretrain a conformer model $\bm{\epsilon}_{\theta}$ on a large-scale conformer dataset. MD fine-tuning: The model is then combined with additional temporal interpolator ${\mathbf{s}}_{\phi}^{\mathrm{tp}}$ to approach the MD distribution $p^{\mathrm{md}}({\mathbf{x}}^{[T]})$ .

Figure 2: Cascaded temporal interpolator block.

Panel of Figure 3 — Figure 3: ( A ) reports performance of BasicES with borrowed numbers from Xu et al. (2022) on SOTA baselines; ( B ) Example conformers from BasicES on both QM9 & Drugs

	Method	JSD (Mean — Median) ( $\downarrow$ )
	Method	Bond Angle		Bond Length		Torsion		TICA_0		TICA_0,1
QM9	MD Oracle	0.042	0.028	0.032	0.031	0.192	0.134	0.318	0.291	0.413	0.394
	AR + EGNN	0.702	0.677	0.770	0.780	0.702	0.761	0.770	0.788	0.820	0.824
	AR + ET	0.705	0.746	0.680	0.721	0.553	0.586	0.568	0.562	0.783	0.786
	AR + GeoTDM	0.752	0.746	0.699	0.694	0.466	0.506	0.456	0.463	0.714	0.719
	GeoTDM	0.691	0.690	0.676	0.670	0.489	0.527	0.449	0.453	0.691	0.694

Table 1: Performance Comparison on QM9 Unconditional Generation and Drugs Forward Simulation.

Figure 4: ( A ) Bond length and ( B ) torsion angle distributions from reference ( red ), our generations ( green ), and GeoTDM ( blue ). MSM occupancies from reference versus ( C ) our generations and ( D ) MD oracles. Autocorrelations of torsion angles for an example molecule from ( E ) reference, ( F ) our generations, and ( G ) GeoTDM. Gray dashed line marks the 1/e decorrelation threshold.

Figure 5: ( A ) Reference free energy surface along the top two TICA components. ( B ) Generated interpolation trajectory projected onto the reference surface (red = start, orange = end). Surface is colored by metastate assignment. ( C ) Key frames from intermediate metastates. ( D ) Statistics comparing JSD, valid path rate, average path probability, and valid path probability for generated trajectories and replicate MD oracles.

	Method	JSD (Mean — Median) (↓)
	Method	Bond Angle		Bond Length		Torsion		TICA_0,1
		Mean	Median	Mean	Median	Mean	Median	Mean	Median
QM9	EGInterpolator-N	0.538	0.538	0.583	0.580	0.441	0.494	0.680	0.685
QM9	EGInterpolator	0.305	0.292	0.210	0.188	0.363	0.380	0.636	0.642
Drugs	EGInterpolator-S	0.325	0.330	0.330	0.321	0.414	0.419	0.673	0.672
Drugs	EGInterpolator-N	0.332	0.332	0.386	0.383	0.455	0.466	0.698	0.703

Table 2: Ablation on QM9 Unconditional Generation and Drugs Forward Simulation.

Figure 6: ( A ) Free energy surface along top two TICA components for a reference ( top ), ours ( top ), and GeoTDM ( bottom ) tetrapeptide trajectory. ( B ) Torsion auto-correlations from ours reference ( left ), ours ( middle ), and GeoTDM ( right ) ( C ) Key collective variable distribution JSD metrics. ( D ) C- $\alpha$ RMSF analysis and visualization for selected protein dynamics generation.

Appendix figures & tables26 assets

Supplementary material from the paper’s appendix.

Appendix

Metrics (JSD)	EGInterpolator	MDGen
Bond Angle (↓)	0.092	N/A
Bond Length (↓)	0.056	N/A
Torsion BB (↓)	0.378	0.130
Torsion SC (↓)	0.189	0.093
Torsion All (↓)	0.265	0.109
TICA 0 (↓)	0.409	0.230

Appendix

Table 3: Results on MDGen Tetrapeptide Dataset

Checkpoint	COV-R (%) $\uparrow$		MAT-R (Å) $\downarrow$		COV-P (%) $\uparrow$		MAT-P (Å) $\downarrow$
	Mean	Med.	Mean	Med.	Mean	Med.	Mean	Med.
99	90.18	94.59	0.2969	0.3049	55.23	51.36	0.4932	0.4823
539	87.62	92.03	0.2574	0.2613	58.12	53.24	0.4451	0.4445

Appendix

Table 4: Conformer metrics on QM9 compared between two checkpoints.

Figure 7: Statistics evaluating the JSD with the reference trajectories, valid path rate, average path probability, and valid path probability of our generated trajectories and replicate MD oracles.

Method	Bond Angle	Bond Length	Torsion	TICA 0	TICA 0,1
GeoTDM	0.640 0.645	0.643 0.645	0.498 0.503	0.531 0.550	0.712 0.720
EGInterpolator-Simple	0.208 0.192	0.258 0.244	0.385 0.399	0.462 0.465	0.660 0.662
EGInterpolator-Casc	0.173 0.153	0.142 0.112	0.377 0.388	0.454 0.441	0.650 0.644
EGInterpolator-Casc-U	0.220 0.202	0.195 0.168	0.414 0.429	0.499 0.496	0.689 0.697

Appendix

Table 5: Performance comparison on Drugs Forward Simulation versus Unconditional Generation. Reported values are JSD (Mean — Median)

\downarrow

	Method	JSD (Mean — Median) ( $\downarrow$ )
	Method	Bond Angle		Bond Length		Torsion		TICA 0		TICA 0,1
QM9	MD Oracle	0.042	0.028	0.032	0.031	0.192	0.134	0.318	0.291	0.413	0.394
	GeoTDM	0.691	0.690	0.676	0.670	0.489	0.527	0.449	0.453	0.691	0.694
	EGInterpolator	0.305	0.292	0.210	0.188	0.363	0.380	0.417	0.406	0.636	0.642
	EGInterpolator-Both	0.231	0.219	0.168	0.158	0.348	0.367	0.393	0.390	0.623	0.631
Drugs	MD Oracle	0.036	0.023	0.030	0.028	0.215	0.131	0.484	0.494	0.610	0.630

Appendix

Table 6: Additional Performance Comparisons on QM9 Unconditional Generation and Drugs Forward Simulation.

Dataset	EGInterpolator vs GT W1 $\downarrow$	GeoTDM vs GT W1 $\downarrow$
QM9	0.8127	2.9201
Drugs	0.7728	12.7664

Appendix

Table 7: Top: Average Wasserstein-1 (W1) distance between predicted and ground-truth (GT) energy profiles for EGInterpolator-Casc and GeoTDM across dataset test sets. Bottom: Per-block W1 analysis in forward simulation roll-outs for Drugs.

Metric	EGInterpolator vs GT W1 $\downarrow$
Overall Energy W1	0.3806
GeoTDM Energy W1	12.8636

Appendix

Table 8: Top: Average Wasserstein-1 (W1) distance between predicted and ground-truth (GT) energy profiles for EGInterpolator and GeoTDM across tetrapeptides. Bottom: Per-block W1 analysis of forward simulation roll-outs.

Block	Frame Range	Mean W1 $(\downarrow)$
1	1–250	0.2303
2	251–500	0.2778
3	501–750	0.3693
4	751–1000	0.2769
5	1001–1250	0.4541
6	1251–1500	0.3962

Appendix

Table 9: Block-wise Wasserstein-1 Progression. Mean W1 error across 250-frame blocks during forward simulation.

Figure 8: JSD metrics computed for Bond Angles, Bond Lengths, Torsions, and Decorrelation Times. Compared between EGInterpolator (green) and EGInterpolator- F (purple).

Train $\rightarrow$ Test	Method	Bond Angle		Bond Length		Torsion		TICA_0		TICA_0,1
		Mean	Med.	Mean	Med.	Mean	Med.	Mean	Med.	Mean	Med.
Standard $\rightarrow$ Enlarged	GeoTDM	0.690	0.690	0.674	0.668	0.488	0.529	0.452	0.451	0.695	0.699
	EGInterpolator-N	0.539	0.538	0.584	0.582	0.447	0.492	0.438	0.440	0.678	0.685
	EGInterpolator	0.307	0.293	0.214	0.194	0.361	0.385	0.416	0.409	0.633	0.639
Enlarged $\rightarrow$ Standard	GeoTDM	0.757	0.757	0.782	0.793	0.488	0.533	0.454	0.453	0.697	0.703
	EGInterpolator-N	0.470	0.460	0.540	0.544	0.433	0.481	0.443	0.440	0.681	0.691

Appendix

Table 10: JSD Metric (

\downarrow

) for QM9 Unconditional Generation. Top: Trained on Standard Train, evaluated on Enlarged Test. Bottom: Trained on Enlarged Train, evaluated on Standard Test.

Model	Bond Angle (↓)	Bond Length (↓)	Torsion (↓)
EGInterpolator Normal (QM9)	0.305 / 0.292	0.210 / 0.188	0.363 / 0.380
EGInterpolator $\alpha=1$ (QM9)	0.398 / 0.391	0.618 / 0.613	0.358 / 0.372
EGInterpolator Normal (DRUGS)	0.173 / 0.153	0.142 / 0.112	0.377 / 0.388
EGInterpolator $\alpha=1$ (DRUGS)	0.435 / 0.445	0.580 / 0.091	0.378 / 0.382

Appendix

Table 11: JSD metrics for bond angle, bond length, and torsion across QM9 and DRUGS datasets with and without temporal layers.

Dataset	GT MD	Original Gen.	$\alpha=1$ Gen.	Shuffled GT	Shuffled Gen.
QM9 Test	101.0	13.59	5.2	5.2	5.2
Drugs Test	130.1	185.64	5.2	5.2	5.2

Appendix

Table 12: Mean torsional decorrelation times (ps) across test sets, comparing GT MD data, our original generations, i.i.d. conformer generations (

\alpha=1

), and shuffled variants. Shuffled data collapse to the frame rate of 5.2 ps, reflecting a lack of temporal structure.

Figure 9: Top: Logits of

\alpha

for each spatial and temporal layer after convergence on the QM9 unconditional generation task. Bottom: Logits of

\alpha

for each spatial and temporal layer after convergence on the DRUGS forward simulation task. Both: Results obtained with EGInterpolator-Casc .

$\lambda$	COV-R (%) $\uparrow$		MAT-R (Å) $\downarrow$		COV-P (%) $\uparrow$		MAT-P (Å) $\downarrow$
	Mean	Med.	Mean	Med.	Mean	Med.	Mean	Med.
Casc
0.000	87.99	91.98	0.2539	0.2600	58.30	53.62	0.4430	0.4397
0.025	87.82	92.98	0.2570	0.2598	57.70	53.17	0.4470	0.4396
0.050	88.34	92.84	0.2568	0.2556	58.29	53.78	0.4460	0.4439
0.075	88.16	93.61	0.2577	0.2610	57.38	52.86	0.4490	0.4467

Appendix

Table 13: QM9 results across

\lambda

for Casc and Simple .

Embedding Type	Input	Dimension
Atom Embedding	Atomic Number	30
Bond Embedding	No Bond, Bond Type, 2nd/3rd-order edge	4

Appendix

Table 14: Atom and bond embedding specifications.

Atom Features
Indices	Description	Options	Type
0–1	Aromaticity	true, false	One-hot
2–7	Hybridization	$sp$ , $sp^{2}$ , $sp^{3}$ , $sp^{3}d$ , $sp^{3}d^{2}$ , other	One-hot
8	Partial charge	$\mathbb{R}$	Value
9–16	Implicit valence	0, 1, 2, 3, 4, 5, 6, other	One-hot
17–24	Degree	0, 1, 2, 3, 4, 5, 6, other	One-hot

Appendix

Table 15: Node feature vector based on atom-level properties.

Figure 10: Distributions computed from reference conformers shown in red, Checkpoint 539 in green, and Checkpoint 99 in purple. We see that 539 aligns more closely with reference distributions across all collective variables and shows improved discretization of torsional states.

Dataset & Duration	OpenMM MD	4x Block Diffusion	Full Diffusion
Drugs (5.2 ns)	584.52	201.70	161.08
QM9 (2.6 ns)	151.38	-	60.08

Appendix

Table 16: Average time (s) taken to generate trajectory

Figure 11: Distributions computed from reference QM9 trajectories (red), EGInterpolator (green), and GeoTDM (purple). Across all examples, our framework more closely matches the reference distributions across all collective variables and better captures torsional state discretizations than GeoTDM.

Figure 12: Autocorrelations of individual torsion angles for an example molecule, comparing reference trajectories with generations from EGInterpolator and GeoTDM. For the challenging task of capturing temporal de-correlation behavior, EGInterpolator closely follows the reference dynamics, whereas GeoTDM fails to model frame-to-frame correlations effectively.

Figure 13: Distributions computed from reference Drugs trajectories (red), EGInterpolator (green), and GeoTDM (purple). Across all examples, our framework aligns closely with reference distributions across all collective variables and exhibits improved torsional state discretization compared to GeoTDM.

SMILES	GT	EGInterpolator	GeoTDM
CC(CO)(CO)CC#N	$-440.287\pm 0.005$	$-440.249\pm 0.036$	$-438.206\pm 1.813$
COC[C@@]1(CO)N[C@H]1C	$-441.428\pm 0.008$	$-441.364\pm 0.075$	$-439.407\pm 3.932$
C#CCCC@HOCC	$-388.299\pm 0.006$	$-388.225\pm 0.117$	$-385.743\pm 2.259$
CCOCCCN1CC1	$-405.525\pm 0.007$	$-405.387\pm 0.426$	$-402.751\pm 3.424$
CC(=O)C@HCCO	$-460.165\pm 0.007$	$-460.140\pm 0.021$	$-458.121\pm 1.225$
CCCC@@(CC)OC	$-390.780\pm 0.006$	$-390.753\pm 0.026$	$-387.954\pm 3.434$

Appendix

Table 17: Top: Mean and standard deviation (Hartrees) of energies for selected QM9 test molecules, comparing ground-truth (GT), EGInterpolator , and GeoTDM . Bottom: Block-wise energy means and standard deviations for selected Drugs test molecules, showing how EGInterpolator tracks GT distributions across successive diffusion blocks.

Figure 14: Generated QM9 interpolation trajectories from EGInterpolator , projected on the reference surface. The red point denotes the start frame, and the orange point denotes the end frame. The reference surface is colored by metastate assignment. Each row corresponds to a different molecule, and each column shows a generated interpolation. These examples illustrate the model’s ability to generate efficient and meaningful transition paths.

Figure 15: Generated Drug interpolation trajectories from EGInterpolator , projected onto the reference surface. The red point indicates the start frame, and the orange point indicates the end frame. The reference surface is colored by metastate assignment. Each row corresponds to a different molecule, and each column shows a generated interpolation. These examples highlight the model’s ability to generate efficient and meaningful transition paths.

Figure 16: Top: Logits of

\alpha

for each spatial and temporal layer after convergence on QM9. Bottom: Logits of

\alpha

for each spatial and temporal layer after convergence on DRUGS. Both: Results obtained with EGInterpolator-Casc for the interpolation task.

Panel of Figure 17 — Figure 17: Results obtained with EGInterpolator-Simple .

Molecular Dynamics Simulation

Molecular dynamics simulation models atomic motion to study molecular structure, material behavior, reactions, and thermodynamic properties.

See all 55 papers

Atomic Trajectory Modeling with State Space Models for Biomolecular Dynamics
Mar 18, 2026Liang Shi, Jiarui Lu, Junqi Liu +3Molecular Dynamics Simulation State Space Models
School of Computer Science, Peking University · BioGeometry · Mila - Québec AI Institute +3
Fine-Tuning Diffusion Models for Molecular Generation via Reinforcement Learning and Fast Sampling
May 31, 2026Guang Lin, Shikui Tu, Lei XuMolecular Optimization Diffusion Model Fine-Tuning
Department of Computer Science and Engineering, Shanghai Jiao Tong University, Shanghai, China · Guangdong Laboratory of Artificial Intelligence and Digital Economy (SZ), Shenzhen, Guangdong, China
FlashMol: High-Quality Molecule Generation in as Few as Four Steps
May 7, 2026Xinyuan Wei, Zian Li, Shaoheng Yan +2Diffusion Distillation Few-Step Diffusion Sampling
Institute for Artificial Intelligence, Peking University · Yuanpei College, Peking University · School of Intelligence Science and Technology, Peking University +2

Mar 18, 2026J/K move · Enter open · S save

Atomic Trajectory Modeling with State Space Models for Biomolecular Dynamics

Liang Shi, Jiarui Lu, Junqi Liu, Chence Shi+2

School of Computer Science, Peking University · BioGeometry · Mila - Québec AI Institute+3

Molecular Dynamics Simulation State Space Models

Open paper AskPDF arXiv

Understanding the dynamic behavior of biomolecules is fundamental to elucidating biological function and facilitating drug discovery. While Molecular Dynamics (MD) simulations provide a rigorous physical basis for studying these dynamics, they remain computationally expensive for long timescales. Conversely, recent deep generative models accelerate conformation generation but are typically either failing to model temporal relationship or built only for monomeric proteins. To bridge this gap, we introduce ATMOS, a novel generative framework based on State Space Models (SSM) designed to generate atom-level MD trajectories for biomolecular systems. ATMOS integrates a Pairformer-based state transition mechanism to capture long-range temporal dependencies, with a diffusion-based module to decode trajectory frames in an autoregressive manner. ATMOS is trained across crystal structures from PDB and conformation trajectory from large-scale MD simulation datasets including mdCATH and MISATO. We demonstrate that ATMOS achieves state-of-the-art performance in generating conformation trajectories for both protein monomers and complex protein-ligand systems. By enabling efficient inference of atomic trajectory of motions, this work establishes a promising foundation for modeling biomolecular dynamics.

Abstract

Molecular Dynamics Simulation

Explore similar work

Atomic Trajectory Modeling with State Space Models for Biomolecular Dynamics

Fine-Tuning Diffusion Models for Molecular Generation via Reinforcement Learning and Fast Sampling

FlashMol: High-Quality Molecule Generation in as Few as Four Steps