Atomic Trajectory Modeling with State Space Models for Biomolecular Dynamics

2026-03-18

AI Summary

Existing methods struggle to efficiently generate atomistic biomolecular dynamics trajectories with long-range temporal dependencies, particularly underperforming on protein–ligand complex systems. This work proposes ATMOS, a novel framework that introduces state space models to atomic trajectory generation for the first time. By integrating a Pairformer-based state transition mechanism with a diffusion decoding module, ATMOS unifies the autoregressive modeling of both monomeric proteins and protein–ligand complexes. Trained on large-scale molecular dynamics trajectory datasets (e.g., mdCATH and MISATO) as well as PDB structures, the method effectively captures long-term temporal dependencies, significantly improving both accuracy and efficiency in trajectory generation. ATMOS achieves state-of-the-art performance across both system types, demonstrating its robustness and generalizability.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Planning, Routing, and Scheduling: Planning with Markov Models (MDPs, POMDPs)
• Natural Language Processing: Generation

Application Category

• Social Networks and Social Media: Generative AI / large language models and their impact on social systems
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• biomolecular dynamics
• atomic trajectory modeling
• temporal dependencies
• protein-ligand systems
• conformation generation

Innovation

• State Space Models
• Atomic Trajectory Generation
• Biomolecular Dynamics
• Diffusion-based Autoregressive Decoding
• Pairformer