Scalable Spatio-Temporal SE(3) Diffusion for Long-Horizon Protein Dynamics

TL;DR

STAR-MD introduces a scalable, SE(3)-equivariant diffusion model with joint spatio-temporal attention, enabling long-horizon, microsecond-scale protein trajectory generation with high fidelity and structural validity.

Contribution

The paper presents STAR-MD, a novel causal diffusion transformer that efficiently models complex spatio-temporal dependencies for long-horizon protein dynamics simulation.

Findings

Achieves state-of-the-art performance on ATLAS benchmark.

Successfully generates stable microsecond protein trajectories.

Outperforms previous methods in conformational coverage and structural validity.

Abstract

Molecular dynamics (MD) simulations remain the gold standard for studying protein dynamics, but their computational cost limits access to biologically relevant timescales. Recent generative models have shown promise in accelerating simulations, yet they struggle with long-horizon generation due to architectural constraints, error accumulation, and inadequate modeling of spatio-temporal dynamics. We present STAR-MD (Spatio-Temporal Autoregressive Rollout for Molecular Dynamics), a scalable SE(3)-equivariant diffusion model that generates physically plausible protein trajectories over microsecond timescales. Our key innovation is a causal diffusion transformer with joint spatio-temporal attention that efficiently captures complex space-time dependencies while avoiding the memory bottlenecks of existing methods. On the standard ATLAS benchmark, STAR-MD achieves state-of-the-art performance across all metrics--substantially improving conformational coverage, structural validity, and dynamic fidelity compared to previous methods. STAR-MD successfully extrapolates to generate stable microsecond-scale trajectories where baseline methods fail catastrophically, maintaining high structural quality throughout the extended rollout. Our comprehensive evaluation reveals severe limitations in current models for long-horizon generation, while demonstrating that STAR-MD's joint spatio-temporal modeling enables robust dynamics simulation at biologically relevant timescales, paving the way for accelerated exploration of protein function.