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

TL;DR

This work tackles the challenge of simulating long-horizon protein dynamics beyond the reach of standard molecular dynamics by introducing STAR-MD, a scalable SE(3)-equivariant autoregressive diffusion model with joint spatio-temporal attention. By combining a causal diffusion transformer, continuous-time conditioning, and contextual noise, STAR-MD captures non-separable spatio-temporal dependencies while avoiding cubic spatial complexity, enabling stable microsecond trajectories on large proteins. The authors provide theoretical backing via Mori–Zwanzig formalism to justify history-dependent modeling and memory enrichment, and demonstrate state-of-the-art performance on the ATLAS benchmark at 100 ns, with robust extrapolation to 240 ns and 1 µs, outperforming existing baselines in conformational coverage, structural validity, and dynamic fidelity. The approach promises accelerated exploration of protein function by delivering physically plausible dynamics at biologically relevant timescales with scalable computation.

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.

Figures (29)

• Figure 1: Overview of STAR-MD generation. Input contains protein sequence and a starting conformation. In autoregressive diffusion generation, structural information of previously generated conformations and current noisy conformations are encoded into single and pair representations. A joint spatio-temporal attention block is employed to capture context information to update the single representation of the current frame. The main model block iterated to diffuse a clean conformation for the current frame and added to history for generating next frames.
• Figure 2: Kinetic fidelity and conformational coverage on the ATLAS 100ns benchmark. (a) Comparing C$\alpha$ coordinate RMSD (top) and autocorrelation (bottom) at varying lagtime for different models. STAR-MD better captures the overall trend and characteristic magnitudes, similar to the MD reference runs (dashed lines). Shaded bands represent $\pm1$ standard deviation. The small size of the shaded bands demonstrate the robustness of this metric. (b) Conformational coverage comparison for 6XB3-H for all models and 3 MD simulations. Generated trajectories are projected onto the first two principal components (PCs) of the reference MD simulation (gray contours). Only structurally valid frames are considered for this plot. Recall and Validity are reported for each run. Baseline methods (MDGen, ConfRover) exhibit limited diversity, becoming confined to a small region of the conformational landscape. AlphaFolding's generated trajectories consist of all structurally implausible frames, with a validity of 0%. In contrast, STAR-MD demonstrates significantly broader exploration (with a recall value of 0.65), visiting two of the major modes observed in the MD reference, matching the diversity seen in independent MD runs.
• Figure 2: Results on molecular dynamics trajectory generation at 240ns and 1µs timescales. STAR-MD demonstrates competitive performance across different temporal scales, with particularly strong quality metrics at both timescales.
• Figure 3: Long-horizon stability and error accumulation across different time scales. We plot the structural validity percentage over time for trajectories generated by STAR-MD and baseline models, evaluated at 100ns (left), 240ns (middle), and 1µs (right) horizons. Shaded bands represent $\pm 1$ standard deviation across 5 different repeats. While most models exhibit clear error accumulation over simulation time, STAR-MD maintains high structural validity, regardless of the simulation time scale.
• Figure 4: Stability of STAR-MD across different temporal strides for 1µs generation. We plot the structural validity over time for two 1µs trajectories generated with different strides: 2.5 ns/frame (400 steps) and a more challenging 1.2 ns/frame ($\sim$833 steps). Solid lines show mean validity, while shaded bands represent $\pm 1$ standard deviation across test proteins. Our models remain stable and maintain high structural quality even when generating much longer sequences of frames than seen in training.

Theorems & Definitions (4)

• Proposition 1
• proof
• Remark 1: Memory Enrichment
• Remark 2: Spatio-Temporal Coupling