Atomic Trajectory Modeling with State Space Models for Biomolecular Dynamics

Liang Shi; Jiarui Lu; Junqi Liu; Chence Shi; Zhi Yang; Jian Tang

q-bio.BMMar 18, 2026

Atomic Trajectory Modeling with State Space Models for Biomolecular Dynamics

Authors: Liang Shi, Jiarui Lu, Junqi Liu, Chence Shi, Zhi Yang, Jian Tang

Organizations: School of Computer Science, Peking University · BioGeometry · Mila - Québec AI Institute · Université de Montréal ·

Abstract

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.

Figures & tables

Figure 1: Overview of the Proposed Atomic Trajectory Modeling Framework.

		Conformation Models		Trajectory Models				Reference MD traj.
		AlphaFlow-MD	ESMFlow-MD	MDGen	TEMPO	ConfRover	ATMOS	Reference MD traj.
Predicting flexibility	Pairwise RMSD $r$ $\uparrow$	0.61	0.46	0.65	0.73	0.69	0.92	0.96
	Global RMSF $r$ $\uparrow$	0.57	0.47	0.67	0.61	0.62	0.90	0.95
	Per-target RMSF $r$ $\uparrow$	0.85	0.80	0.77	0.52	0.81	0.89	0.95
Distributional accuracy	Root mean $\mathcal{W}_{2}$ -dist. $\downarrow$	3.40	5.00	3.86	4.98	4.52	2.70	1.89
	MD PCA $\mathcal{W}_{2}$ -dist. $\downarrow$	2.26	2.16	2.75	2.38	3.04	1.79	1.67

Table 1: Evaluation on mdCATH . The best scores are highlighted in bold and the second-best scores are underlined .

r

: Pearson correlation;

J

: Jaccard similarity;

\mathcal{W}_{2}

: 2-Wasserstein distance.

Figure 2: Case Study on 1s79a00 , a target from the test set of mdCATH. (A) the visualization of ensemble from MD simulation (reference) or ATMOS (sampled). (B) The visualization of snapshot at specific frame index. (C) The RMSF versus the residue index. (D) The moving RMSD to the initial structure (frame 0) along the trajectory frames.

Model	Inference Time (Hour)	Peak Mem (GB)
MD Simulation	$>$ 10	$<$ 1
ConfRover	0.467	22.53
ATMOS	0.174	3.445
ATMOS-Parallel	0.044	6.117

Table 2: Efficiency evaluation on the 1gnla01 protein (139 AA). MD simulation is performed with Amber24 on GPU ( Salomon-Ferrer et al., 2013 ) , employing an integration time step of 4 fs.

	Method	per-target RMSF $r$ $\uparrow$	Root mean $\mathcal{W}_{2}$ -dist $\downarrow$	Steric Clashes $\downarrow$	t-RMSD Error $\downarrow$
Ligand	DenoisingLD ( Wu and Li, 2023 )	0.028	0.791	1.825	34.01
	GNNMD ( Siebenmorgen et al., 2024 )	0.135	0.870	81.69	0.672
	NeuralMD-SDE ( Liu et al., 2025 )	0.574	0.698	7.667	0.698
	VerletMD ( Liu et al., 2025 )	0.456	36.83	14.67	43.78
	DynamicBind ( Lu et al., 2024b )	0.678	1.470	7.113	-
	Boltz-2 ( Passaro et al., 2025 )	0.613	1.237	0.210	-

Table 3: Evaluation on MISATO . The best scores are highlighted in bold and the second-best scores are underlined .

r

: Pearson correlation;

\mathcal{W}_{2}

: 2-Wasserstein distance. Boltz‑2 is trained on the entire MISATO dataset, which may lead to data leakage; we include it as a strong but potentially optimistic baseline.

Metric	ATMOS	w/o bi-level sampling	ODE
Pairwise RMSD $r$ $\uparrow$	0.92	0.88	0.73
Global RMSF $r$ $\uparrow$	0.90	0.84	0.64
Per-target RMSF $r$ $\uparrow$	0.89	0.85	0.63
Root mean $\mathcal{W}_{2}$ -dist. $\downarrow$	2.70	3.06	3.89
MD PCA $\mathcal{W}_{2}$ -dist. $\downarrow$	1.79	2.09	2.73
Joint PCA $\mathcal{W}_{2}$ -dist. $\downarrow$	2.32	2.61	3.28

Table 4: Ablation study on sampling protocol .

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 3: Detailed architecture of our framework. BIF: Bidirectional Information Flow.

Hyperparameter	Value
Number of Pairformer blocks	4
Single representation dimension	384
Pair representation dimension	128
Triangle multiplication hidden dimension	128
Number of triangle attention heads	4
Transition layer expanding factor	4

Appendix

Table 5: Architectural hyperparameters for the modified Pairformer.

Molecular Dynamics Simulation

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

See all 55 papers

Align Your Structures: Generating Trajectories with Structure Pretraining for Molecular Dynamics
Apr 5, 2026Aniketh Iyengar, Jiaqi Han, Pengwei Sun +3Molecular Dynamics Simulation 3D Molecular Generation
Stanford University · Lambda, Inc
ProtScape: A molecular structure and energy-aware representation for protein conformation generation
Oct 27, 2024Siddharth Viswanath, Xingzhi Sun, Lucas Lee +9Molecular Dynamics Simulation Molecular Representation Learning
Department of Computer Science, Yale University · Department of Genetics, Yale University · Program in Computing, Boise State University +2
Equivariant Asynchronous Diffusion: An Adaptive Denoising Schedule for Accelerated Molecular Conformation Generation
Mar 10, 2026Junyi An, Chao Qu, Yun-Fei Shi +33D Molecular Generation Molecular Generation
Shanghai Academy of Artificial Intelligence for Science, SAIS · Artificial Intelligence Innovation and Incubation (AI3) Institute, Fudan University

Apr 5, 2026J/K move · Enter open · S save

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

Aniketh Iyengar, Jiaqi Han, Pengwei Sun, Mingjian Jiang+2

Stanford University · Lambda, Inc

Molecular Dynamics Simulation 3D Molecular Generation Diffusion Models

Open paper AskPDF arXiv

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.

Abstract

Molecular Dynamics Simulation

Explore similar work

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

ProtScape: A molecular structure and energy-aware representation for protein conformation generation

Equivariant Asynchronous Diffusion: An Adaptive Denoising Schedule for Accelerated Molecular Conformation Generation