Atomic Trajectory Modeling with State Space Models for Biomolecular Dynamics

TL;DR

This paper introduces ATMOS, a novel state space model-based generative framework that efficiently produces atom-level biomolecular trajectories, capturing long-range temporal dependencies and outperforming existing methods in accuracy and applicability.

Contribution

The paper presents ATMOS, the first SSM-based generative model capable of modeling long-range temporal dependencies in biomolecular trajectories for both monomers and complexes.

Findings

Achieves state-of-the-art trajectory generation accuracy

Successfully models long-range temporal dependencies

Effective for both monomeric and complex biomolecular systems

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.