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

AI summary

Overview

Research area: Generative machine learning for computational structural biology — specifically, diffusion models for simulating protein molecular dynamics (MD) trajectories.

Technical level: Advanced. The paper assumes familiarity with diffusion models, SE(3) equivariance, transformer attention, and molecular dynamics concepts such as RMSD, tICA, and VAMP-2 scores.

Scope: The paper introduces STAR-MD (Spatio-Temporal Autoregressive Rollout for Molecular Dynamics), an SE(3)-equivariant autoregressive diffusion transformer that generates protein trajectories over microsecond timescales using joint spatio-temporal attention, and benchmarks it against MDGen, AlphaFolding, and ConfRover on the ATLAS dataset plus newly generated 250 ns and 1 µs reference simulations.

What This Paper Is About

Classical MD simulation is the gold standard for studying how proteins move, but its femtosecond integration steps make biologically relevant timescales (microseconds to milliseconds) extremely expensive to reach. Existing generative models accelerate this by learning dynamics from data, but they are typically limited to nanoseconds, restricted to small proteins, and degrade over long rollouts because they either cap memory to a fixed window or model space and time in separate, interleaved modules. STAR-MD's goal is to generate physically plausible, stable trajectories across microsecond horizons and large proteins by modeling space and time jointly inside a single attention mechanism.

Key Contributions

1. STAR-MD architecture. A causal diffusion transformer for autoregressive protein trajectory generation that uses joint spatio-temporal (S×T) attention over residue-frame tokens, rather than the interleaved "space-then-time" modules used in prior work, with 2D Rotary Position Embedding encoding residue and frame indices.

2. Training and inference techniques for stable rollouts. Historical context noise addition during both training and inference (τ ∼ U[0, 0.1]), block-diffusion-style causal training that concatenates clean and noisy frames for a single forward pass, and continuous-time conditioning on physical stride (Δt ∼ LogUniform[10⁻², 10¹] ns) via adaptive layernorm.

3. A Mori-Zwanzig theoretical argument. The paper uses the Mori-Zwanzig formalism to argue that coarse-grained dynamics are non-Markovian and thus require temporal history, and it presents a "Memory Inflation" proposition claiming that removing explicit pairwise features inflates the memory kernel into a non-separable spatio-temporal form — motivating joint rather than factorized attention.

4. Extensive long-horizon evaluation. Benchmarking on the standard 100 ns ATLAS task plus new 240 ns (32 proteins) and 1 µs (8 proteins) reference simulations generated with ATLAS protocols, covering structural validity, conformational coverage, and dynamic fidelity.

Main Findings

• State-of-the-art on the 100 ns ATLAS benchmark. In Table 1, STAR-MD reports JSD 0.43 ± 0.01, recall 0.54 ± 0.01, tICA 0.17 ± 0.00, RMSD deviation 0.07 ± 0.02, autocorrelation deviation 0.02 ± 0.00, VAMP-2 deviation 0.10 ± 0.02, Cα validity 86.81% ± 0.64, all-atom validity 98.18% ± 0.05, and combined validity 85.29% ± 0.62. The running text states 0.58 recall and 86.36% validity, which differs slightly from Table 1's values; the MD oracle is given as JSD 0.31, recall 0.67, tICA 0.17, and 96.43% combined validity.

• Baselines produce many physically invalid frames. On the 100 ns task, MDGen and ConfRover reach only 71.83% and 56.94% Cα validity and 68.31% and 52.06% combined validity respectively, while AlphaFolding produces almost no valid structures (0.82% all-atom validity, 0.47% combined); AlphaFolding also failed to run on the four largest test proteins due to out-of-memory errors and was evaluated on 78/82 proteins.

• Stable extrapolation to microsecond timescales. Trained only on 100 ns ATLAS data with no fine-tuning, STAR-MD keeps combined validity at 83.15% ± 1.99 on 240 ns and 79.93% ± 1.04 on 1 µs, while MDGen drops to 56.60% and 24.81%, AlphaFolding to 0.63% and 0.06%, and ConfRover-W to 36.51% and 36.91%.

• Error accumulation is the key failure mode of baselines. Figure 3 shows structural validity declining sharply with simulation time for baselines across 100 ns, 240 ns, and 1 µs horizons, while STAR-MD stays close to the MD oracle; the authors attribute this largely to the historical-context noise mechanism.

• ConfRover cannot scale to long horizons without modification. Its pairwise KV cache exceeded hardware limits (1869 GB CPU RAM plus 8× H100 GPUs with 80 GB VRAM each), forcing the authors to evaluate a windowed-attention variant with attention sink tokens (ConfRover-W). STAR-MD needs only O(NL) KV-cache memory for single-residue features versus O(N²L) for ConfRover.

• Favorable complexity trade-off. Joint S×T attention costs O(N²L²) compute, while Pairformer + pairwise temporal attention costs O(N³L + N²L²) and Pairformer + single temporal attention costs O(N³L + NL²); STAR-MD avoids the cubic spatial term.

• Robustness to temporal resolution. The 1 µs task used a 2.5 ns stride (400 frames); regenerating at 1.2 ns stride (the text says roughly 850 frames, while Figure 4's caption says approximately 833 steps) without retraining maintained high structural quality.

• Qualitative coverage on an example protein. For 6XB3-H, Figure 2(b) reports STAR-MD recall of 0.65 and visits two major modes of the reference MD landscape, whereas MDGen and ConfRover remain confined to a small region and AlphaFolding shows 0% validity.

• Ablations confirm design choices matter. In Table 3 (100 ns), removing contextual noise drops combined validity to 76.12%, separate spatial/temporal attention to 86.70%, and placing S×T attention outside the diffusion block to 82.56%, versus 85.18% for the full model in that table (note the table's STAR-MD row reports JSD 0.42 and recall 0.57, slightly different from Table 1).

Methodology in Plain English

STAR-MD generates a trajectory one frame at a time. Each frame is a set of per-residue rigid bodies (translations plus rotations), modeled in SE(3). A denoising diffusion model learns to predict the noise added to these rigid bodies, so sampling the reverse process yields a clean conformation.

The twist is conditioning. Instead of feeding the model a fixed window of past frames, or compressing history into a static condition, STAR-MD uses a causal transformer where each token represents a (residue, frame) pair. Attention is computed jointly across both residues and frames (with 2D rotary embeddings for residue and frame indices), so the model can directly learn couplings such as how motion at one residue relates to past motion at a distant one. Previously generated frames are cached as keys and values, just as in language-model-style generation.

To train this efficiently, the authors concatenate clean history frames with noisy target frames and apply a block-wise causal attention mask, so all frames can be denoised in one forward pass while preserving causality. To stop errors compounding during long rollouts, they add small random noise (τ between 0 and 0.1) to history frames during training and apply the same perturbation to each generated frame before it becomes history at inference. Because they condition on the physical time stride Δt (sampled log-uniformly from 10⁻² to 10¹ ns) through adaptive layernorm, the same model can generate at different temporal resolutions without retraining.

Evaluation uses the ATLAS dataset (100 ns trajectories for 1390 proteins) with standard splits, plus new 250 ns and 1 µs MD simulations produced with the original ATLAS protocols for long-horizon testing. Metrics span structural validity (Cα clash/chain-break checks plus MolProbity backbone and side-chain outlier checks), conformational coverage (Jensen-Shannon divergence and recall on principal components of the reference trajectory, restricted to valid conformations), and dynamic fidelity (tICA lag-time correlation, RMSD, autocorrelation, and VAMP-2 score relative to MD). All trajectories are aligned to the first reference frame by Cα superposition.

Why This Matters

Research impact. The paper argues that long-horizon generative MD has been blocked less by modeling capacity than by architecture: fixed-window models lose memory, pairwise-feature models hit cubic costs and quadratic KV-cache growth, and factorized spatial/temporal modules cannot express coupled dynamics. STAR-MD demonstrates that a single joint attention mechanism over residue-frame tokens, combined with noise-perturbed history, can remain stable an order of magnitude beyond the training horizon. The Mori-Zwanzig framing also gives a theoretical account of why temporal history and joint attention are necessary in coarse-grained models, which the authors present as a justification for a broader class of recent temporal architectures.

Real-world applications (each tied to the biological processes the paper cites):

• Studying catalysis and enzyme function, where conformational changes gate reactivity.
• Modeling binding events between proteins and partners.
• Characterizing allosteric regulation, where motion at one site controls activity at another.
• Accelerating the exploration of protein function generally, by replacing some microsecond-scale MD simulations with generated trajectories.

Industry relevance. MD simulation is a computational bottleneck in structure-based drug discovery and protein engineering, where microsecond behavior is often what matters. A model that generates physically plausible trajectories at a fraction of the cost — and in the authors' complexity analysis, with O(NL) cache memory instead of O(N²L) — could make large-protein, long-timescale dynamics screening more practical. The work comes from ByteDance Seed with academic collaborators at Georgia Tech and UCLA, and the authors note the first author's contribution was made during an internship at ByteDance Seed.

Future Directions

1. Closing the gap to the MD oracle. STAR-MD narrows but does not eliminate the distance to reference MD on JSD and recall (for example, JSD 0.43 versus 0.31 on the 100 ns task), leaving room for better coverage of the conformational landscape.

2. Opening the black box of context and noise. The paper reports that historical-context noise is critical for stable rollouts but does not characterize the best noise schedule or how much history is actually needed; systematic study of both is an open question.

3. Scaling to larger proteins and longer horizons. The authors' complexity analysis is asymptotic, and the paper motivates the O(NL) KV cache as a scalability advantage, but the largest test proteins already broke AlphaFolding's memory; whether STAR-MD extends to substantially larger systems and millisecond horizons remains to be shown.

4. Broadening the benchmark. Because ATLAS only provides 100 ns trajectories, long-horizon results rest on the authors' own newly run 240 ns and 1 µs simulations for 32 and 8 proteins respectively — a small sample that invites independent replication and extension.

Target Audience

This paper is best suited to machine learning researchers working on generative models for scientific data, computational chemists and biophysicists interested in replacing or accelerating MD simulation, and practitioners in protein engineering or structure-based drug discovery who need long-timescale dynamics at scale. Readers will get the most out of it with prior background in diffusion models, SE(3)/SO(3) geometry, and standard MD analysis metrics; the theoretical Mori-Zwanzig section in particular assumes comfort with coarse-graining arguments from statistical mechanics.

Authors’ 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.