Scalable Spatio-Temporal SE(3) Diffusion for Long-Horizon Protein Dynamics
Authors: Nima Shoghi
Yuxuan Liu
Yuning Shen
Rob Brekelmans
Pan Li
Quanquan Gu
Venue: ICLR 2026 Poster (2026)
Keywords: Proteins
Molecular dynamics
Generative modeling
Diffusion models
Autoregressive modeling
SE(3)-equivariant diffusion
Spatiotemporal modeling
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: A scalable spatiotemporal SE(3) diffusion model that generates long-horizon protein trajectories with state-of-the-art fidelity.

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 spatiotemporal 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 spatiotemporal 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 spatiotemporal modeling enables robust dynamics simulation at biologically relevant timescales, paving the way for accelerated exploration of protein function. Project page: https://bytedance-seed.github.io/ConfRover/starmd

Published as a conference paper at ICLR 2026
SCALABLE SPATIO-TEMPORAL SE(3) DIFFUSION FOR LONG-HORIZAS HORIZON PROTEIN DYNAMICS
Nima Shoghi' : Yuxuan Liu 3 Yuning Shen', Rob Brekelmans', Pan Li Quanquan Gu "ByteDance Seed Georgia Institute of Technology SUniversity of California, Los Angeles {yuning. shen, quanquan. gu, rob. breke imans J@bytedance com yxliuêmath. ucla edu
(nimash, (aanii)@gtecc.eeee
ABSTRACT
Molecular dynamics (MD) simulations remain the gold standard for studying protein dynamics, but their computational cost limits access to biologically rel- evant timescales. Recent generative models have shown promise in accelerat- ing simulations, yet they struggle with long-horizon generation due to architec- tural constraints, error accumulation, and inadequate modeling spatio-temporal dynamics. We present STAR-MD Molecular Dynamics). a scalable -qquivaaaa diffusion model that gener- ates physically plausible protein trajectories over microsecond timescales. Our key innovation is a causal diffusion transformer with joint spatio-temporal atten- tion 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-ar performance across all metrris--ubssanniaaly improving conformational coverage, structural validity, and dynamic fidelity com- pared to previous methods. STAR-MD successfully extrapolates to generate stable microsecond-scale trajectories where baseline methods fail catastrophi- cally, 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. Project page: sttsraaaal
Autoregressive Rollout for
bttpp::/bytedae seed. www.giit....ooooooooooooieoo
INTRODUCTION
Protein functions emerge from conformational dynamics the continuous structural changes that underlie important biological processes such as catalysis, binding, and allosteric regulations (Mc- Cammon 1984 Berendsen & Hayward 2000). Classical molecular dynamics (MD) simulation remains the gold standard and relies on physical models to integrate atomic motions over time us- ing Newtonian mechanics. However, the need for small integration steps (~ femtoseconds) severely limits its practicality for exploring the
miiiicooommmiisocood timescales often required to capture biologically relevant events. Recent advances in protein structure prediction (Jumper et al.) [ Abramson et al. 2024 Lin et al. 2023) and generative modeling for conformational dynamics (Jing jtt 20244b Cheng et al., 20255 Shen al., 2025 Costa al., (2024) offer promising data-driven approaches to accelerate simulations by learning the dynamics directly from data and generating trajectories at coarser temporal resolutions. However, existing methods are still constrained by short time horizons (typically up nanoseconds) and struggle to scale to larger proteins. These limi- tations highlight an urgent need for generative models capable of producing physically plausible protein trajectories over extended timescales while remaining computationally efficient.
Scaling protein dynamics modeling to long timescales requires both an efficient structural repre- sentation and architecture capable of capturing complex spatio-temporal dependencies. Current approaches often rely on pairwise residue representations or computationally expensive architectures (e.g., Alphafoold--tyle triangular attention), leading to quadratic memory growth and cubic com- putational cost with respect to protein size. These challenges become even more pronounced when accounting for spatio-temporal coupling across multiple conformations during trajectory modeling.
Work done during internship ByteDance Seed. CCorresponding author.

---

Published as a conference paper at ICLR 2026
As result, most existing models treat spatial and temporal components separately, employing inter- leaved "spatial" and "temporal" modules (Jing et al. 2024b; Cheng et al. 2025 Shen et al. 2025), which limits their expressiveness in capturing coupled dynamics. These limitations confine current models to small proteins and short simulation horizons, ultimately hindering their ability to learn long-range temporal dependencies and generate high-quality conformations over extended rollouts.
We introduce Spatio-Temports Autoregressive Rollout for Molecular Dynamics (STAR-MD), an (equi--quvv autoregressive diffusion model for generating physically plausible protein trajec- tories over microsecond timescales, even for large protein systems, At its core, STAR-MD employs a causal diffissonn transformer with joint spatio-temporal attention, enabling improved autoregres- sive generation by dynamically computing attention over historical context during diffusion and capturing long-range dependencies through expressive yet memory-efficient spatio-temporal mod- eling.
Our main contributions are summarized as follows:
We present STAR-MD, novel autoregressive diffusion transformer that leverages efficient spatio- temporal attention to model the complex dependencies underlying protein dynamics.
Through several key technical improvements, including historical context noise addition dur- ing training and inference, flock-diffusion-style causal training, and architectural optimizations, STAR MD achieves efficient, scalable training and stable trajectory generation, significantly im- proving over current state-of-the-an models.
We perform an exhaustive evaluation across multiple different simulation timescales ranging from 100 ns to μs, providing comprehensive assessment of conformation quality, coverage, and dy- namic fidelity. Our analyses yield new insights into the limitations of current state-of-the-art models in long-horizon generation and offer valuable guidance for future model design.
2 RELATED WORK
Protein Conformation Generation. Several models (Jing et al. 2024a Lewis et al. 2024: Wang
et al.
2024) build on advances in protein structure prediction (Jumper al.] 2021 Abramson al., Lin et al. 2023) and diffusion models (Ho et al. (020) to directly generate time-independent conformations. These methods provide an efficient alternative MD by enabling parallel sampling. However, they capture only the equilibrium distribution of conformations and cannot model the temporal evolution of protein dynamics.
2024
MD Trajectory Generation. To model the temporal evolution, operator- aperator-based methods Klein et al. (2024) Costa et al. (024) aim learn transport operators that predict conformations lagged times. These methods approximate evolution as a Markovian process, and thus fail to capture the non-Markovian properties often present in partially observed systems such as protein dynamics data. In contrast, methods generating trajectories, akin to video generations, consider dependencies across multiple time frames. AlphaFolding (Cheng et al., 2025) incorporates higher-order information through additional "motion frames" and generates multiple future frames simultaneously. Similarly, MDGen (Jing et al.) (224b)) models the joint distribution of frames across 100-ns trajectories. How- ever, both methods only capture the dependencies in a fixed context and prediction window. When generating longer trajectories through extension, they discard memory of earlier windows, breaking temporal consistency. To address this, ConfRover (Shen et al., (025) adopts language model-style autoregression that can generate trajectories of arbitrary lengths while maintaining full memory via KV-caching. STAR-MD follows this approach to avoid the "memory break" in block-based models.
Scalable Structural and Temporal Modeling, Current structural models, including AlphaFolding and ConíRover, represent protein structures using single and pairwise features and process them with specialized architectures such as Pairformer and Invariant Point Attention (IPA) (Umper et al.) 2021 While these modules are highly expressive and preserved the required roto-translational symmetries, they are computationally and memory intensive. This limits model efficiency and scal- ability, especially when extended to modeling trajectories across multiple time frames. MDGen sidesteps this limitation by anchoring trajectories to key frames, encoding only single embeddings with standard transformers. While this allows modeling of long trajectories (c.g., 250 frames for AT- LAS), the reliance on key frames limited flexibility and model performance is suboptimal. Lastly, all the above approaches model dynamics between frames using interleaved spatial and temporal modules. While this decoupled design reduces computational cost, it limits the model's expressive
2

---

Published as a conference paper at ICLR 2026
power to capture complex spatio-tempore relationships. In contrast, STAR-MD proposes to ploy joint spatio-tempone attentions on the single embeddings' direct space-time time processing while keeping the memory footprint manageable.
3 METHOD
To generate long and physically plausible protein trajectories, we propose Spatio-Tempone Autoregressive Rollout for Molecular Dynamics (STAR-MD), spatio-temporal (E(3)diiffusson model that operates within an autoregressive framework. In the following sections, we first describe the overall autoregressive diffusion framework, then detail the architecture of our diffusion model, and finally explain the training and inference procedures.
Frozen OpenFold
Seq: MGLSD... PVEK
Structure
2ns
4ns
996ns
998ns
1.0ps
Inputs
Output
D Autoregressive Loop
Diffusion Loop
Diffusion Block
Nx
Joint SxT Attention
Invariant Point Attention
History
Edge and Backbone Updates
New Frame
Currant
Z
Backbone rigids
Model
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 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.
3.1 AUTOREGRESSIVE SE(3)-DIFFUSION MODELS
PROTBIN TRAJECTORIES
Autoregressive Trajectory Generation. We aim to generate protein conformation trajectories æ1:L in an autoregressive fashion I p(æz æ<et; Alz), where A represents the time interval between frames. This formulation amortizes trajectory generation into frame-level process conditioned on the entire past history, enabling the model to generate future dynamics of arbitrary length while maintaining flexible memory representation learned from data. Causal transformers, widely used in language models for efficient autoregressive sequence modeling (Touvron al. 2023). provide natural architectural choice for this task.
SE(3) Diffusion Models. To model the generation process of each frame p(acy ce Alr),
(conditional) diffusion model on the Riemannian manifold SF((3,, where x = =TT R which captures translation positions
{I J==1 J3N and rotations
{R' N E SO(3)* N for each amino-
acid residue in protein sequence of length N (Yim et al. 2023 Wang et al. 2024 Shen et al.) 2025). The forward diffusion process independently corrupts translations and rotations across the diffusion time
T- = Vo+To o + VI- arc, c~N(0, Iax)
(1)
R2 ~ IGSO, (Ro, [
(2)
where IGSO3 is the isotropic Gaussian on S3(3). Given a noisy structure x", diffusion time T, and condition denoising score network
89 (x. c)
15
88 (a c)
(3)

---

Published as a conference paper at ICLR 2026
is trained via denoising score matching to predict the noise added to both components. During inference, the learned score function can be used to sample from the reverse diffusion process and generate protein structures.
To capture temporal dependencies needed for autoregressive generation, the condition с for gener- ating the current frame should incorporate all past history ><t- Specifically, we design efficient and expressive causal diffusion transformer network to compute c (x22 <e: a 7) from the preced- ing clean frames and the current noisy frame x.. This design differs from prior works such as MDGen Jing et al., 2024b) & Alphafolding (Cheng et al. 2025) which do not include autoregres- sive conditioning, and ConfRover (Shen al. (2025) which compresses all preceding frames a static condition c <x<<e). In the following section, we detail the design of с (a<e, 27.7) and how it efficiently integrates spatio-temporal information from preceding frames.
3.2 ARCHITECTURE
In this section, we describe the architecture choices underlying STAR-MD. The main features of this architecture lie in the diffusion decoder, which uses spatio-temporal attention for autoregressive conditioning on previous frames and block diffiusion allowing for efficient training and KV caching.
Input Module. We mirror the Fraoeer module Shen et al. (2025), starting with frozen OpenFold FoldingModin e for sequence-level single-residue and pairwise features 8, x = {s,, Rig Jij=1. We incorporate these time-independent features with amino-acid metadata and pairwise Cs distance within each frame to obtain dimendependent features st = used initial input the diffusion blocks.
L
init, which are
Diffusion Blocks. Our diffusion module contains submodules stacked into blocks:
Invariant Point Attention layers (Jumper et al. 202T) update single features sy using pair features zy and noisy frames z/- This step is independent for each frame.
2. Joint Spatio-Temporal Attention layers update the single-residue features sæ by attending to those from previous frames 85c, which allows for exchange of temporal informaiton.
3. Pair features 2, are updated from single features 8, using the EdgeTransite iion layer from Jumper et al. (202T) Noisy frames a- are updated from single features sy using an MLP-based backbone update.
After stacking the above blocks, the final update of the coordinates x/ is used as the score function prediction and fed to the denoising score matching loss.
Joint Tentoolll Attention. Our joint spatio-temporee (SxT) attention mechanism integrates information across the temporal dimensions, departing from previous models that employ factorized, sspace-then-time attention which imposes restrictive bias that spatial and temporal dependencies are separable. Instead, our SxT attention operates on tokens representing residue-frame pairs (i, E). allowing it to directly model non-separable relationships, such as how motion at one residue is coupled to past motion at distant site. 2D Rotary Position Embedding (2D-RoPE (Heo et al.) 2024) is used to embed residue and frame indices, which allows extrapolation past the training number of frames.
This architectural choice is key to our model's scalability. We analyze the per-layer computa- tional complexity of our SxT attention (O(N2L2)) against general S+T (i.e., "space-then-time"" architectural paradigms (see Appendix B for details). Most baseline architectures employ Pair- former (Jumper et al., 2 backbone, incurring a cubic spatial complexity O(N*L) due to ex- pensive triangular attention operations. We consider two representative baseline configurations: (1) Pairformer + Pair Temporal Attention (e.g., ConfRover (Shen et al., 2025)): This configuration scales as ON + ~22 Here, the pairwise temporal attention alone incurs O(N2 L²) cost, matching that of STAR-MD, but the additional cubic spatial overhead makes it significantly more expensive. Furthermore, this approach incurs quadratic memory cost O for caching pair- wise key-value states during autoregressive generation, which becomes a severe bottleneck for long rollouts. In contrast, STAR-MD only requires O(N L) memory for caching single-residue features. (2) Pairformer + Single Temporal Attention: This configuration scales as O(N While the temporal scaling is theoretically superior to STAR-MD's 's O(N L?) in the limit L > N. the bic spatial cost O 3 L) dominates in practical regimes. STAR-MD thus offers favorable trade-off, enabling efficient modeling of long trajectories for large proteins without the cubic spatial overhead.
4

---

Published as a conference paper at ICLR 2026
3.3 ADDITIONAL MODEL DETAILS
Block-causal Attention for Efficient Training. Traditional video diffusion models (Peebles & Xie 2023; Ma al. 2024 denoise all frames simultaneously, using fixed window size. For autoregressive generation with variable context length, the model denoises one frame at a time, attending clean history contexts. This straightforward during inference, where we cache clean frames for efficient autoregressive generation. However, given the sequential nature of this process, extra work is needed for parallel training. To simultaneously denoise all tokens with clean history context, we follow AAriola et al. (2025): Teng et al. (025); Deng et al. (2024) to concatenate clean and noisy frames as input sequence and employ a special block-wise attention pattern: all frames only attend to clean history frames, preserving the causal structure during training. At the cost of doubling the sequence length, the model can optimize the training loss for all frames in a single forward pass, aligning parallel seacher-foreing training with sequential autoregressive inference.
Contextual Noise Perturbation for Robust Rollouts. To mitigate compounding errors during long autoregressive rollouts, we apply contextual noise perturbation technique inspired by Diffusion Forcing (Chen et al., 2024 Song et al., 2025). During training, we perturb historical context frames Iet by applying the forward diffusion process (Eqs. [22 with small, randomly sampled noise level J ~ [0, 0.1J to obtain noisy contexts x-rr- The model then learns to predict frame condi- tioned on this perturbed history. Critically, at inference time, we apply the same noise perturbation: after generating frame f, we add noise at level J before using it as context for subsequent frames. This maining-inffreece consistency ensures that the model experiences similar input distributions during deployment as during training, making it robust to its own prediction errors and preventing compounding drift in long rollouts.
time Conditioning, To handle the vast range of timescales in protein motion, we draw the physical stride A Log-oonnio
100] ns independently for every training example. By conditioning the network on At through adaptive layernorm (AdaLN, Peebles & Xie (2023), ) the model learns to modulate its internal activations as a function of both diffusion progress and physical timescale. Crucially, this training strategy decouples the physical duration of trajectory from the number frames in the context window. Even when restricted to small context windows during training (as most methods are), randomly sampling large strides exposes the model to relative time deltas spanning orders of magnitude. This allows the network to learn long-range temporal dependencies without the computational cost of training on long sequences or the need for complex context-length extrapolation techniques.
3.4 THEORETICAL JUSTIFICATION FOR STAR-MD ARCHITECTURE
In this subsection, we use the Mori-Zwanzig formalism MMori 1965 Zwanzig. 1961 (See Ap-
pendix for more details) to provide theoretical justification for two key aspects of our architecture: the necessity of temporal history in coarse graineddr modeling, and the specific requirement for joint spatio-temporal attention arising from the removal of explicit pairwise features.
Atomistic molecular dynamics simulations evolve according to Hamiltonian mechanics and are Markovian in the full phase space of atomic positions and momenta. However, practical generative models must operate on coarse-grained representations, such as per-residue coordinates for protein structures. The Mori-Zwanzig formalism shows that projecting Markovian dynamics onto coarse- grained variables yields a Generalized Langevin Equation with three terms (Eq. (5)): b: a Markovian component, a memory kernel encoding history dependence, and a stochastic force representing elim- inated degrees of freedom. Consequently, accurate modeling of coarse -rainedrrindd protein dynamics requires Mon-Markovi models that incorporate temporal history capture memory effects aris- ing from eliminated degrees of freedom. Our work makes this connection explicit, providing the theoretical justification underlying recent temporal architectures trajectory modeling (Jing al.) 2024b Cheng al. 2025 Shen et al. (025)..
Existing models typically rely on explicit pairwise features to capture spatial structure, but this in- curs prohibitive OL + LN*) or (LN'3 costs (see Section 3.2 STAR-MD circumvents this by projecting out pairwise features, a design choice we analyze through our Memory Infla- tion proposition (Proposiiion I Appendix A). We show that removing explicit spatial correlations "inflates" the memory kernel for the remaining residues, necessitating a significantly richer tem- poral history compensate. Crucially, this inflated kernel exhibits non-separable spatio-temporal
5

---

Published as a conference paper at ICLR 2026
coupling (Remark 2). meaning spatial and temporal modes cannot be factorized. This theoreti- cal insight directly motivates our architecture: instead of the interleaved spatial and temporal blocks used in prior work. STAR-MD employs joint -empool attention approximate this complex, non-separable memory kernel, balancing physical fidelity with computational scalability.
EXPERIMENTS
We conduct a comprehensive set of experiments to evaluate STAR-MD's ability to generate long- horizon protein dynamics. First, we benchmark STAR-MD against state-of-the-art models on the standard ns trajectory generation task (Section 4.2) Next, we assess its extrapolation capa- bilities by generating trajectories over longer, microsecond-scall horizons not seen during training (Section 4.3). Finally, we perform a series of analyses and ablations to investigate model stability, scalability, and the contributions of our key design choices (Section 4.4
4.1 EXPERIMENTAL SETUP
Dataset. ATLAS dataset (Vander Meersche et al. 2024) contains 100 ns MD trajectories for 1390 structurally diverse proteins. We use the standard train/val/test splits from prior works (Jing al. 2024b, Shen et al. 2025) to assess model performance in a transferable setting. To evaluate long- horizon generation-a a key focus of our work-we we extend the benchmark by running new MD sim- ulations to produce 250 ns and 1 su trajectories for selected proteins, using the original ATLAS simulation protocols for consistency. Further details are provided in Appendix C
Implementations. We follow the training procedure described in Section 3.3 and generate trajec- tories using the configurations listed Table We compare STAR-MD with three state-of-the-art trajectory generative models trained ATLAS: AlphaFolding (Cheng et al. 2025). MDGen (Jing ct al. 2024b), and ConfRover (Shen et al. 2025). All models parametrize proteins using SE(3) backbone rigids (translations and rotations) with torsional angles for side-chain atoms. To stan- dardize evaluation, trajectories from variable-stride models (STAR-MID, ConfRover) are generated directly the required intervals. Trajectories from fixed-stride models (AlphaFolding, MDGen) are first generated at their native resolution and then subsampled. Finally, we include an oracle reference based on MD simulations run independently to represent the target performance. All tra- jectories are aligned to the starting frame of the reference ATLAS trajectory via Ca superposition prior to analysis (Appendix D). Specific parameters for each benchmark are detailed in the relevant sections below, with full implementation details in Appendix FF
We assess the physical plausibility of generated conformations using a hi- erarchy of geometric and stereochemical metrics. First, we perform -aarse-grained checks for Ca-Ca clashes (non-bonded atoms too close) and chain breaks (consecutive Ca atoms too far apart). Second, use the MolProbity suite (Chen et al. 2010) for fine-grained, all-atom analysis of backbone Ramachandran and side-chain rotamer outliers. We report three distinct validity met- rics based on different criteria: Ca-level Validity, passing only Ca checks; All-Atom Validity, passing only MolProbity checks; and Combined Validity, passing all checks simultaneously. See Appendix for details on the thresholds.
Structural Quality:
Conformational Coverage: To evaluate how well generated trajectories explore the conforma- tional space of the reference MD simulation, we follow the protocol of Shen et al. (2025). We project all conformations into a low-dimensional space defined by the principal components of the reference trajectory and compute the Jensen-Shannon Divergence (JSD) and conformation call between the distributions. To ensure that coverage reflects physically plausible exploration, we report these metrics computed exclusively on structurally valid conformations ("Cov Valid").
Dynamic Fidelity: We assess temporal coherence using four metrics. First, we use tICA lag-time correlation quantify the preservation of slow collective variables (Molgedey & Schuster 1994 Kěrez-Hernándezz et al., 2013) Shen et al. 2025): crucially, this is computed only on valid transi- tions avoid artifacts from broken structures. Additionally, magnitude of structural change over varying intervals, autocorrelation to assess temporal mem- ory, and VAMP-2 score evaluate how well slow dynamical modes of the system are captured. For last three metrics, we report the deviation from MD reference, with details in Appendix D
evaluate RMSD measure the
4.2 STANDARD BENCHMARK: 100 NS TRAJECTORY GENERATION
We first evaluate STAR-MD on the standard 100 ns ATLAS benchmark. Following the protocol of Shen et al. 2025), generate trajectories of 80 frames at 1. ns interval. Results are summarized

---

Published as a conference paper at ICLR 2026
Table 1: Quantitative results on 100 ns trajectory generation (ATLAS test set). STAR-MD achieves state-of-the-ar performance across all metrics, with particularly significant improvements in con- formational coverage (JSD, FI) and dynamic fidelity (ICA, average difference of RMSD, autocor- relation and VAMP-2 score to MD reference).
Cov Valid
Dynamic Fidelity
Validity
Model
JSDL
Rect
UCAt
RMSD!
AutoCori
VAMP-2,
CA%
A
CA+AA
MD (Oracle)
0.31
D.67
0.17
0.00
0.00
0.02
98.37
98.07
96.43
MDGen
0.562a01
0.28
0.38 = am
0.05 0:000
0.38 0.00
71.83
95.03
68.31
AlphaFolding"
0,59+am
0.20=a0r
NIA
3,31
0.12 +.0
1.36
10.58+<um
0,82+0.10
0.47=004
ConfRover
p.52
0.36
0.15:uen
0.20 = 200
0.08 0:0
0.47 u
26.94
92.47:2025
52.061:026
STAR-MD
0,43+sm
0.54=am
0.17=am
0.07
0.02
0.10
86,81+164
98,18-005
85.29
(a)
(b)
MD (R 6.75 Vi 16mmo
MDGan (2: 0.76 VI 90%
Alcharsiding (h: 0.00
0.8
0.7
0.6
0.5
0.4
th 6.75, 100%)
0.
02
CantRawer (R: 0.20, SHNN
STAR-HD (l: 8.8.0, 50%0
05-
0.3
02
00
Lagtime (frames)
PC1
PC1
PC1
Figure 2: Kinetic fidelity and conformational coverage on the ATLAS 100 ns benchmark. (a) Comparing Ca coordinate RMSD (top) and autocorrelation (bottom) at varying lagtime for differ- ent models. STAR-MD better captures the overall trend and characteristic magnitudes, similar to the MD reference runs (dashed lines). Shaded bands represent #1 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 small region of the conformational landscape. AlphaFolding's generated trajectories consist of all struc- turally implausible frames, with a validity of 0%. In contrast, STAR-MD demonstrates significantly broader exploration (with recall value of 0.65), visiting two of the major modes observed in the MD reference, matching the diversity seen in independent MD runs.
in Table I It is important note that due to its high computational cost, AlphaFolding failed to run on the four largest proteins in the test set, highlighting its scalability limitations even on this standard benchmark. On the full test set, STAR-MD, achieves a superior balance of conformational coverage. dynamic fidelity, and structural quality.
The results demonstrate clear performance deficiencies in all baseline models. Baselines like MD- Gen and ConfRover exhibit low conformational coverage (0.30 and 0.42 Recall, respectively) and produce substantial number of structurally invalid frames, with validity rates of only 64.9% and 49.7%, respectively. AlphaFolding achieves higher raw recall (0.51), but this metric is mislead- ing as a negligible 0. 1% of its generated frames are structurally valid; when controlled for valid structures, this recall drops to 0.01, rendering AlphaFolding as very unreliable model for this task. In contrast, STAR-MD achieves the highest conformational coverage (0.58 recall) and struc- tural validity (86.66%) of any generative model, significantly narrowing the performance gap to the ground-truth MD simulation which serves as the oracle for this task.
In Figure aa and [ we compare the level of conformation changes different lag times, as an indicator how well the model captures the ground truth motion changes. STAR-MD shows im- proved coherence with MD reference trajectories, while AlphaFolding significantly overestimates the motion and MDGen and ConfRover tend to underestimate the dynamics level.
AlphaFolding results evaluated on 78/82 proteins due out-of-memory error for 4 large proteins.

---

Published as a conference paper at ICLR 2026
Table Results on molecular dynamics trajectory generation at 240 and I jss timescales. STAR-
MD demonstrates competitive performance across different temporal scales, with particularly strong quality metrics both timescales.
Cov Valid
Dynamic Fidelity
Validity
Model
Time
JSD4
Rect
RMSD4
AutoCor
CA%T
A
CA+AA%
ns
0.26
0.75
0.01
0.00
99.53
96.83
96.36
MD (Oracle)
1 js
0.23
0.91
0.00
0.00
96.25
86.50
82.75
MDGee
ns
0.52
0.38=a01
0.48
= 2001
0.25 # a.
63.25
87.83
56.60+2.10
1 us
0,56++mm1
0,36++003
0.37 042
0.39 + 0.01
36,11+734
56.99
24.81
ns
0.57
0.20
1. 1.76iuus
0.14
8.96
0,94
0.63
Alphafolding
1 jss
0:65
0.20
0.78
0.04
9.64=a02
0.19
0,06:1000
ns
0.42.amm
0.35 not
0.39 1 0.01
14.7
16.51
ConfRover-W
1 us
0.55 +0.0
0.45+aa2
0.33
0.38 +
L4.74+177
62.32+3.43
+6.91+1.20
ns
0.44
0.59
0.20 = auaz
0.03 :: 0.0
85.16:1.1
97.57.24.13
83.15
STAR-MD
1 us
0.46:a01
0.61+tu0
0.13 0.82
0.10 + 0.02
88.47 1.00
89.81
19.93+1.44
CA + AA Validity (%)
MD (Oracle)
MDGen
AlphaFolding
ConfRover
STAR-MD
100
50
50
100
100
200
0
500 Time (ns)
1000
Time (ns)
Time (ns)
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 100 ns (left), 240 ns (middle), and μs (right) horizons. Shaded bands represent +1 standard deviation across different repeats. While most models exhibit clear error accumulation simulation time, STAR-MD maintains high structural validity, regardless of the simulation time scale.
Finally, Figure provides qualitative assessment of conformational coverage using the principal component projection for an example ATLAS 100 ns protein. STAR-MD exhibits better significantly better exploration and diversity than comparison methods and is able to populate two distinct modes.
4.3 EXTRAPOLATING FOR LONG-HORIZON TRAJECTORIES
The ability to generate stable, physically realistic trajectories over extended time horizons critical for modeling functionally relevant protein dynamics. Most conformational transitions of biological interest occur beyond the 100 ns timescale, where existing generative models begin to fail. We assess STAR-MD's long-horizon capability through two increasingly challenging settings: 240 simula- tion with 32 proteins and 1 jus simulation with 8 proteins, with all proteins unseen during training. To rigorously evaluate temporal extrapolation, all models
trained exclusively the 100 AT- LAS trajectories and received no fine-tuning for longer horizons. This setup directly tests the ability generalize dynamics far beyond the training data distribution. Since the ATLAS benchmark only provides trajectories up to 100 ns, we generated our own molecular dynamics simulations for longer timescales to create proper reference data for evaluation. These extended simulations follow the same simulation protocols as ATLAS but continue the dynamics to longer timescales.
Due its scaling issues, ConiRover could not be evaluated on either 240 ns or us generation tasks. This is because ConfRover performs temporal attention over pair features, requiring previ- ous frames' pair features to be stored memory as KV cache (see Appendix B for more detailed analysis). As such, even with CPU offloading of the KV cache, ConfRover's memory requirements exceeded our hardware limits (1869 GB CPU RAM, 8x H100 GPUs with 80GB VRAM each). As result and for fair comparison, we utilize variant of ConfRover with windowed attention with attention sink tokens (as described in Xiao et al. (2023b)) to reduce memory usage. We report results for this windowed variant, labeled ConfRover- W. in Table 2
Table summarizes the results for these two extended timescales. Remarkably, STAR-MD main- tains stable and competitive performance at both the 240 timescale and the challenging 1 ys
8

---

Published as a conference paper at ICLR 2026
CA (A+. AA Validity (%)
1us 2.5ns/frame
10a
cache size (GB)
100
1us, 1.2ns/frame
10²
25668
101
80
10"
protein szz- 100
protein size= 300
60
10
protein szc- 500
10:2
protein
0
200
400
600
800
1000
600
800
Time (ns)
Trajectory length (frames)
Figure 4: Stability of STAR-MD across differ- ent temporal strides for 1 js generation. We plot the structural validity over time for two ju trajectories generated with different strides: 2.5
Figure 5: Memory footprint of KV caches with varying protein sizes and trajectory lengths. ConfRover is shown in solid lines and STAR-MD is shown in dashed lines. Both mod- els contain 8 attention layers with hidden di- mension of 128 (ConfRover) MD). STAR-MD show much smaller memory cost compared ConfRover. Horizontal lines marks the 256 GB and 2 TB memory caps.
ns/frame (400 steps) and more challenging 1.2 ns/frame (~883 steps). Solid lines show mean validity, while shaded bands represent 1 stan- dard deviation across test proteins. Our mod- els remain stable and maintain high structural quality even when generating much longer se- quences of frames than seen in training.
256 (STAR-
timescale, demonstrating the effectiveness of our joint spatio-temporal attention and contextual noise techniques for long-horizon stability. Further, STAR-MI) exhibits controlled error accumulation, with low clash and break rates that remain consistent throughout the extended trajectory (Figure 3 (right), Section 4.4). Among the methods achieving reasonable structural quality, STAR-MD attains the best balance of coverage and fidelity, nearing the oracle MD performance. In both the 240 ns and us case, AlphaFolding produces physically implausible, high-error structures despite good cover- age metrics, while ConfRover-V W and MDGen degrade significantly over the microsecond horizon. We further report kinetic metries for the long trajectory cases in Fig. 8 where STAR-MD shows the best alignment with the RMSD and autocorrelation of the MD) reference from among baseline methods.
4.4 ADDITIONAL ANALYSIS
Error Accumulation. key challenge in long-horizon generation error accumulation, where mi- nor inaccuracies compound over time, leading to trajectory degradation. Figure plots the average validity percentage over simulation time for the 100 ns, 240 ns, and jus simulations. As the sim- ulation horizon extends, baseline models exhibit rapidly deteriorating performance, with structural validity declining sharply. In contrast, STAR-MD remains stable, maintaining high structural valid- ity close the MD oracle across all timescales. This demonstrates model's robustness against error accumulation, quality largely attributable to the historical-contex noise mechanism, which enables stable long-horizon autoregressive generation.
Varying temporal resolution for long-horizon generation. The uu generation task above uses a stride of 2.5 ns. resulting in 400 frames. Thanks to our continuous-time conditioning, we can gener- ate at different temporal resolutions without retraining. We test this by generating I us trajectory with 1.2 stride, nearly doubling the number of frames to ~50. Figure ssows shows that STAR-MD remains stable even with this much longer sequence of frames, maintaining high structural qual- ity. This further underscores our model's robustness to different sampling intervals and trajectory lengths, showing strong capability for stable, long-horizon autoregressive generation.
KV Cache Analysis. A critical limitation of ConfRover is the need to maintain the KV-cache of single and pair embedding for temporal attentions. In contrast, STAR-MD only requires maintaining the KV cache for singles. By only maintaining attentions among single representations, STAR-MD has magnitudes lower memory footprints (see Fig. 5) for comparison). This advantage allows our model to maintain a full KV caches in GPU memory without resorting to CPU offloading (e.g,
additional overhead) or sliding window style KV caches (compromises temporal history).
4.5 ABLATION STUDIES ON KEY COMPONENTS
To validate key design choices, conduct systematic ablations of the main components of STAR-MD. Table summarizes the quantitative results on the 100 ns benchmark. Appendix H.2 contains full ablation tables for the 100 ns, 240 ns, and I μs benchmarks.

---

Published as a conference paper at ICLR 2026
SxT attention vs. separable attention. We compare our joint SxT attention with separable al- ternative that processes spatial and temporal dimensions sequentially. The separable model suffers significant drop in conformational coverage and modest drop in dynamic fidelity. This con- firms that joint SxT attention crucial for capturing the complex, non-separable spatio-temporal dependencies inherent in protein dynamics, which directly impacts the model's ability to explore the correct conformational space.
SxT layer placement. We evaluated placing the spatio-temporal attention module outside the dif- fusion decoder as a static conditioner which compresses the full trajectory history into a single conditioning vector, mirroring Shen et al. (2025). This variant underperforms our full model in both coverage and fidelity. confirming that integrating spatio-temporal attention directly within the diffusion module essential for effective context utilization.
Historical-context noise. Removing the contextual noise perturbation leads to a substantial drop in structural quality. This highlights the importance of this technique for maintaining generation stability, finding that is further supported by our long-horizon experiments in Section Figure 12 compares the CA+AA validity over time for models with and without contextual noise on the 240 ns and μs benchmarks, showing that the noise helps maintain structural quality over long horizons.
These ablation studies demonstrate that each component of STAR-MD addresses specific challenge in protein dynamics modeling. The SxT attention provides the expressivity needed for complex spatio-temporal dependencies. The historical-contex noise ensures stable long-horizon generation by preventing error accumulation.
Table 3: Ablation study 100 ns simulation. The full STAR-MD model is compared against variants with key components removed (no contextual noise, separate spatial and temporal attention, and placing SxT attention outside the diffusion block). These modifications affect model coverage and conformation quality. Metries degraded in each ablation setting are highlighted in red. See complete ablation results on longer trajectories in Appendix H.2
Cov Valid
Dynamic Fidelity
Validity
Model
JSD.
Reet
UCAt
RMSD.
AutoCori
VAMP-2J.
CA
AA%1
CA+AA
STAR-MD
0.42
0.57
0.17
0.09
0.02
0.12
86.62
98.28
85.18
w/o Noise
0.43
0.54
0.18
0.04
0.11
1.22
77.82
97.76
76.12
w/ Sep Attn
0.46
0.46
0.17
0.09
0.05
0.49
87.95
98.34
86.70
w/ Preproc Attn
0.47
0.48
0.17
0.09
0.05
0.54
84.55
97.81
82.56
5 CONCLUSION
We introduced STAR-MD, a novel causal diffusion transformer model for generating long-horizon protein dynamics trajectories. STAR-MD addresses the challenges of long-horizon generation through several key innovations: joint spatio-temporal attention that efficiently models complex space-time couplings, continuous-time conditioning for generation at arbitrary timescales. and noisy-context training scheme that mitigates error accumulation. Our comprehensive evaluations demonstrate that STAR-MD achieves new state-of-the-art performance on the standard 100 ns AT- LAS benchmark, outperforming previous methods in conformational coverage, structural quality, and kinetic fidelity. More importantly, our model successfully extrapolates to generate stable, high- fidelity trajectories up to the microsecond regime, where prior models often fail due to computational costs or compounding errors. By balancing expressiveness and scalability, STAR-MD establishes an efficient and empirically strong foundation for modeling protein dynamics at biologically relevant scales, paving the way for accelerated exploration of complex biological processes.
Limitations and future work. While STAR-MD represents significant step forward, there enues for future improvement. The temporal consistency of the generated trajectories, while strong, does not yet perfectly match that of oracle MD simulations. The model's performance could be further enhanced by training on larger and more diverse MD simulation datasets, such as MD- CATH (Mirarchi et al. 2024). Additionally, promising direction for future work to extend the model's 's capabilities to simulate the dynamics of protein complexes or their interactions with small molecules, which are crucial for understanding biological processes and drug design.
10

---

Published as a conference paper at ICLR 2026
AUTHOR CONTRIBUTIONS AND ACKNOWLEDGMENTS
Author Contributions. Y.S. and Q.G. conceived the study, designed the experimental protocols, and supervised methodology development. N.S., Y.L., and Y.S. developed the spatiotemporal causal diffusion architecture and analyzed model performance. P.L.. and Y.S. conceptualized the contextual noise perturbation, which N.S. implemented. N.S., P.L. and R.B. conducted the theoretical veri- fication of the Mori-Zwanzig formalism and validated the mathematical frameworks. All authors participated in the manuscript preparation and critical review for publication.
Acknowledgments. We would like to thank Chan Lu, Cheng- Yen (Wesley) Hsich, Yi Zhou, Zaix- iang Zheng, Yilai Li and other members of ByteDance for their valuable feedback and suggestions that helped improve this work.
REFERENCES
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, al. Accurate structure prediction of biomolecular interactions with alphafold 3. Nature, 630(8016):493-500, 2024.
Marianne Arriola, Aaron Gokaslan, Justin T Chiu, Zhihan Yang, Zhixuan Qi, Jiaqi Han, Sub- ham Sekhar Sahoo, and Volodymyr Kuleshov. Block diffusion: Interpolating between autore- gressive and diffusion language models. International Conference on Learning Representations, 2025.
Herman JC Berendsen and Steven Hayward. Collective protein dynamics in relation function. Current opinion in structural biology, 10(2):165-169, 2000.
Boyuan Chen, Diego Martí Monsó, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitz- mann. Diffusion forcing: Next-token prediction meets full-sequence diffusion. Advances in Neural Information Processing Systems, 37:24081-24125, 2024.
Vincent B Chen, W Bryan Arendall, Jeffrey Headd, Daniel A Keedy, Robert M Immormino, Gary Kapral, Laura W Murray, Jane S Richardson, and David C Richardson. Molprobity: all- atom structure validation for macromolecular crystallography. Biological crystallography, 66(1): 12-21, 2010.
Kaihui Cheng, Ce Liu, Qingkun Su, Jun Wang, Liwei Zhang, Yining Tang, Yao Yao, Siyu Zhu, and Yuan Qi. 4d diffusion for dynamic protein structure prediction with reference and motion guidance. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 93- 101, 2025.
Allan dos Santos Costa, Ilan Mitnikov, Franco Pellegrini, Ameya Daigavane, Mario Geiger, Zhonglin Cao, Karsten Kreis, Tess Smidt, Emine Kucukbenli, and Joseph Jacobson. Equi- jump: Protein dynamics simulation via so (3)-cquivariant stochastic interpolants. arXiv preprint arXiv:2410.09661 09667, 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in neural information processing systems, 35:16344-16359, 2022.
Chaorui Deng. Deyao Zhu, Kunchang Li, Shi Guang, and Haoqi Fan. Causal diffusion transformers for generative modeling. arXiv preprint 24Xi 2024.
Byeongho Heo, Song Park, Dongyoon Han, and Sangdoo Yun. Rotary position embedding for vision transformer. In European Conference on Computer Vision, pp. 289-305. Springer, 2024.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Proceed- ings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 97817 11122113.
Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. In Proceedings of the 41st International Conference on Machine Learning, pp. 22277-22303, 2024a.
11

---

Published as a conference paper at ICLR 2026
Bowen Jing, Hannes Stärk, Tommi Jaakkola, and Bonnie Berger. Generative modeling of molecular dynamics trajectories. Advances in Neural Information Processing Systems, 37:40534 97:40534-40564, 2024b.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin 2' idek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie, Bernardino Romera- Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Se- bastian Bodenstein, David Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Push- meet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with AlphaFold. Nature, 596(7873) 589, 2021. doi: (10088/415566001-033888
Lcon Klein, Andrew Foong, Tor Fjelde, Bruno Mlodozeniec, Marc Brockschmidt, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36, 2024.
Sarah Lewis, Tim Hempel, José Jiménez Luna, Michael Gastegger, Yu Xie, Andrew YK Foong, Victor Gareía Satorras, Osama Abdin, Bastiaan S Veeling, Iryna Zaporozhets, et al. Scalable emulation of protein equilibrium ensembles with generative deep learning. bioRxiv, pp. 2024-12, 2024.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637):1123-11
2023.
Jiarui Lu, Bozitao Zhong, Zuobai Zhang, and Jian Tang. Str2str: score-based framework for zero-shot protein conformation sampling. In The Twelfth International Conference on Learning Representations, 2024.
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden- and Sain- ing Xie. Sit: Exploring flow and diffusion- based generative models with scalable interpolant transformers. In European Conference on Computer Vision, pp. 23-40. Springer, 2024.
JA McCammon. Protein dynamics. Reports on Progress in Physics, 47(1):1, 1984.
Antonio Mirarchi, Toni Giorgino, and Gianni De Fabritii.. mdcath: A large-scale md dataset for data-driven computational biophysics. Scientific Data, 11(1):1299, 2024.
Lutz Molgedey and Heinz Georg Schuster. Separation of a mixture of independent signals using time delayed correlations. Physical review letters, 72(23))3634, 1994.
Hazime Mori. Transport, collective motion, and brownian motion. Progress of theoretical physics, 33(3):423-455, 1965.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In 2023 IEEE/CVF International Conference on Computer Vision (ICCV), 4172-4182. IEEE, 2023.
Guillermo Pérez-Hernández, Fabian Paul, Toni Giorgino, Gianni De Fabritiis, and Frank Noé. Iden- tification of slow molecular order parameters for markov model construction. The Journal of chemical physics, 139(1), 2013.
Yuning Shen, Lihao Wang, Huizhuo Yuan, Yan Wang, Bangji Yang, and Quanquan Gu. Simulta- neous modeling of protein conformation and dynamics via autoregression. Advances in Neural Information Processing Systems, 2025.
Kiwhan Song, Boyuan Chen, Max Simchowitz, Yilun Du, Russ Tedrake, and Vincent Sitzmann. History-guided video diffusion. arXiv preprint 20Xiv:2502.06764 2025.
Hansi Teng, Hongyu Jia, Lei Sun, Lingzhi Li, Maolin Li, Minggiu Tang, Shuai Han, Tianning Zhang, WQ Zhang, Weifeng Luo, et al. Magi-1: Autoregressive video generation at scale. arXiv preprint ar arXiv:2505.1321 2025.
12

---

Published as a conference paper at ICLR 2026
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arX1v:230..1397,, 2023.
Yann Vander Meersche, Gabriel Cretin, Aria Gheeraert, Jean-Christophe Gelly, and Tatiana Ga- lochkina. Atlas: protein flexibility description from atomistic molecular dynamics simulations. Nucleic acids research, 52(D1):D384 D392, 2024.
Yan Wang. Lihao Wang, Yuning Shen, Yiqun Wang, Huizhuo Yuan, Yue Wu, and Quanquan Gu. Protein conformation generation via force-guided se (3) diffusion models. In Forty-first Interna- tional Conference on Machine Learning. 2024.
Hao Wu and Frank Noé. Variational approach for learning markov processes from time series data. Journal of Nonlinear Science, 30(1):23 66, 2020.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. ArXiv, abs/2309.17453, 2023a. URL https://api.
semanticscholl ott.:/a..oorrrooorrrr00222211111
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309.17433 17453, 2023b.
Tao Xie, Tamjeed Saleh, Paolo Rossi, and Charalampos G. Kalodimos. Conformational states dy- namically populated by a kinase determine its function. Science, 10.1 N0.1126/science.abc222.
2020. doi:
Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. In Proceedings of the 40th International Conference on Machine Learning, pp. 40001-400303 2023.
Shuxin Zheng, Jiyan He, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng, Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, et al. Predicting equilibrium distributions for molecular systems with deep learning. Nature Machine Intelligence, 6(5):558-56) 107555857...0 2024.
Robert Zwanzig. Memory effects in irreversible thermodynamics. Physical Review, 124(4):983, 1961.
13

---

Published as a conference paper at ICLR 2026
CONTENTS
Theoretical Analysis
16
Theoretical Foundation for Learning Coarse-Grained Protein Dynamics
16
A.J Overview and Motivation
16
Formal Derivation Memory Inflation
16
A.3
Implications for Architecture Design
18
Complexity Analysis of
Attention
19
B.
Problem Formulation and Notation
19
B.2
Computational Complexity of Different Attention Mechanisms
19
B.3
STAR-MD vs. Factorized Temporall Architectures
20
B.4
KV Cache Efficiency During Inference
20
B.5
Practical Scaling Behavior
21
Data & Implementation Details
22
ATLAS Dataset
22
D Details on Evaluation Metrics
22
Architecture Details
25
Architecture Details: Pair Features
25
Baseline Implementation
26
Model Hyperparameters
28
Extended Results
30
Extended ATLAS Analysis
30
Additional Results on Dynamic Fidelity
30
Full Ablation Results
3
n00 Filtered Results
35
Additional Visualizations Conformational Coverage
37
Stability and Scalability Analysis
42
Extended Stability Analysis
42
Extreme-Horizon Generation (10
43
14

---

Published as a conference paper at ICLR 2026
L3
Inference Cost
44
J
Applications to Other Systems
44
Ensemble Generation
44
J.2 Case Studies on Conformational Dynamics
47
Miscellaneous
49
K
LLM Usage Disclosure
49
15

---

Published as a conference paper at ICLR 2026
Theoretical Analysis
A THEORETICAL FOUNDATION FOR LEARNING COARSE-GRAINED PROTEIN DYNAMICS
A.1 OVERVIEW AND MOTIVATION
Modeling protein dynamics in solution requires a complete description of the system (e.g., coor- dinates and momenta of all protein and solvent atoms), making it computationally prohibitive for current generative models. Any practical implementation must adopt coarse-grained representations. such as per-residue coordinates solely protein structures. The Mori-Zwanzig formalísm (Mori) 1963 Zwanzig 1961) shows that projecting full-atom dynamics into a reduced subspace (as done by coarse- grraningg introduces memory effects, rendering the dynamics non-Markovian and requir- ing an explicit memory kernel to preserve the original dynamics in this coarse-grand maase--rained space. For protein structural learning, the reduction from singles-and and-pairs representations singles-only rep- resentations introduces a complex spatio-tempone paatio-tmporal-couplle coupling memory kernel. Therefore, to remain scalable without the overhead from expensive pair representations, model must employ an archi- tecture capable of learning complex, non-separable spatio-temporal dependencies. This analysis motivates our departure from prior "space-then senneethen-time factorized models joint spatio-tempon attention mechanism.
Following sections formalize the effects coarse-graine granseggaiinng on protein dynamic learning: We intro- duce Mori-Zwanzig formalism in [Appendix 1 its specialized form in protein structural repre- sentations /Apppeddxx A.2.2 and the implications for neural architecture design in/ppendixx A.3)
A.2 FORMAL DERIVATION OF MEMORY INFLATION
A.2.1 BACKGROUND: MORI-ZWANZIG FORMALISM
The Mori-Zwanzig formalism Mori 1965 Zwanzig 1961 provides a mathematical framework
for deriving reduced dynamical models from bigh-dimensional systems. We begin with a high- dimensional system described by variables [tt) evolving according to:
dr dt
CF())
(4)
where L is the Liouville operator. When we project onto lower-dimensional subspace of variables A(t) using a linear projection operator P the Mori-ZWanzig theorem provides the exact general- ized Langevin equation (GLE) for their evolution:
dA(1)) = PLA(1) + ( K(I ((L=+)A(7)dr++ F(1)
(5)
dt
where:
A(L) represents our chosen variables (the projection of T(()).
PLA(() captures Markovian (instantaneous) dynamics.
Sh K(t - A)r(r)dr is the memory term encoding history dependence and K(t) PCllL-P)C((1 L. the memory kernel, which encodes the information the de- grees of freedom that are not explicitly modeled.
F(t) is the random force, rapidly-ftuctuaiing noise term arising from degrees of freedom in the (1 - P) subspace.
"In the context of the Mori-Zwanzig formalism, P is linear operator that projects onto chosen subspace of the full phase space.
16

---

Published as a conference paper at ICLR 2026
Accounting for the memory term induced by coarse-graine -ranseig is crucial to correctly capture the dy- namics of the original system.
A.2.2 PROTEIN DYNAMICS PROJECTIONS
In deep learning models for proteins, per-residue features (singles) and their pairwise representations (pairs) are commonly adopted to model their structures (JJumper et al. [202] Abramson et al. 2024]
Ding et al. 2024a Shen al. 2025 Chen et 2024) Here, we consider two specific projections:
Singles-and- Pairs Representation We project onto both per-residue features s(t) E Nxd,, and pairwise features z(t) E RW yielding state A (t) = (s(t), The resulting dynamics can be written in block form:
s(t) 2(t)
Qge Zzg
Dax Dgz
s(t) 2(t) +
K!l (t 7)
KP (t 1 7) K!l (t - k (t - (2(T)
(s(7)
dr (t (6)
Fs(t)
where the memory kernel (11) has block structure corresponding singles-singles (K ) singles- pairs (K!!), ) pairs-singles (K!!), ) and pairs-pairs (K!!) interactions.
Singles-Only Representation We project directly onto per-residue features s(t) alone, yielding the more compact state (t) st).. The dynamics become:
+ S (K00))))(())+++++0
(7)
A.2.3 MEMORY INFLATION ON SINGLES-ONLY REPRESENTATION
Proposition [Memory Inflation] Under linearization, the memory kernel K( for the singles- only representation relates to the dingles-and-a and-pairs memory kernel (()) by:
(2) (p) + (A.e + (D)
pp Pes (K) (S2s (p),,))
(8)
Inflation Term
where kp) is the Laplace transform of the memory kernel, p is the Laplace variable, I is the identity matrix, Q represents instantaneous (Markovian) dynamics, and the subscripts and 2 refer singles and pairs features, respectively.
Proof. We begin by taking the Laplace transform of the singles-and dynamics:
ps(p) - s0)
(Raa Paz N
5(p)
(p)s()
(2) (((p)
pā(p) =00 =
i(p)
+
K (p)s(p)
())(()))))))))
+
(2))
(9)
From the second row, we can express {(p) in terms 3(p):
E(p) = pp - -2 (p)
(0) + (2., + R2 + F.(P)]
(10)
Substituting this into the first row equation:
pā(p) - s(0) = 1p.)(p) + (()) +)) F+(p)+
(11)
(Nez - K (p)) pl - -2r 10x (p) [1]
(A.s + KK (2(2)
(12)
17

---

Published as a conference paper at ICLR 2026
Collecting terms with s(p) and comparing with the Laplace transform of the singles-only represen-
tation:
pš(p) - s(0) = (2)))(())) KK F
(13)
We identify:
1222 =Ω4
(14)
k (p) (p) (Dez + (K (p)) pl -2= -K)
(p)
(azs + (p))
(15)
(p) = Fx(p) + (M,, + (x))
(v)) pl A2y k2)
=0) + F(())
(16)
This completes the proof, showing that the memory kernel in the singles-only representation incor- porates additional terms that account for the eliminated pair dynamics.
A.3
IMPLICATIONS FOR ARCHITECTURE DESIGN
From the Memory Inflation result in Proposition we can draw two important remarks that guide architectural design:
A.3.1
MEMORY ENRICHMENT
Remark (Memory Enrichment). When spatial detail is removed from dynamical system, the memory kernel must become more expressive to preserve the system's 's physical accuracy.
This follows directly from the inflation term in Proposition 1. By eliminating the explicit representa- tion of pairwise features z(t), we force the memory kernel K to incorporate additional dynamics that were previously captured through the direct modeling pairs. This represents a fundamental trade-off between:
1. Spatial complexity: Using O(N2) variables to explicitly represent all pairwise relation- ships
2. Temporal complexity: Using a richer memory structure that implicitly encodes these re- lationships through time
The theorem quantifies exactly how much additional memory structure is required: it must include all dynamical information that would have flowed through the pairs variables in the more complex representation.
A.3.2 SPATIO-TEMPORAL COUPLING
Remark 2 -eneiial Coupling). The memory kernel in the singles-only representation can- not be factorized as independent spatial and temporal components, requiring models that capture non-separable coupling between sputial indices and temporal frequencies..
The inflation term:
(O + R:! (1) PZo kgy (w) (Azs + R-l2 (p))
(0)
(17)
has critical property: it cannot be factorized product of purely spatial and purely temporal operators. To see this, consider its structure:
(Oxz + R:(1) (p)) couples spatial indices (i.j) with temporal frequency p
2. The matrix inverse pI Ozo -K (p)
mixes these couplings in non-separable way
3. The final term (Pts R:(1) (p)) further couples the result with additional spatial indices
18

---

Published as a conference paper at ICLR 2026
In general, in the time domain, this corresponds to memory kernel with structure:
kip (t (t- u:ju(t -)
(18)
where j are residue indices. This non-separability means that the memory kernel cannot decom- posed into independent spatial and temporal components spatial relationships evolve with time, and temporal patterns differ across spatial relationships.
A.3.3 CONNECTION TO MACHINE LEARNING ARCHITECTURES
These theoretical results directly inform architectural design for protein dynamics models:
Factorized attention (spatial then temporal, or vice versa) cannot capture the non- separable coupling revealed by Remark 2
2. Joint spatio-temporal attention with tokens indexing both residue and time provides ex- actly the structure needed to learn the inflated memory kernel demanded by Remark
This analysis demonstrates that architectures that model protein dynamics without explicit pairwise features require sophisticated temporal modeling capabilities. The mathematical structure of memory kernel dictates the minimum expressivity requirements for any machine learning model that aims to capture the underlying physics accurately.
B COMPLEXITY ANALYSIS OF SPATIO- TEMPORAL ATTENTION
To understand the computational complexity and scaling behavior of different protein dynamics models, we compare the the spatial and temporal operations in ConfRover and our STAR-MD ap- proach.
B. PROBLEM FORMULATION AND NOTATION
For protein with N residues and trajectory with L frames, we define:
Sy E R Nxds; Per-residue (singles) features at frame
1z E NxNxss Pairwise features frame
The full representation across all frames would be:
[S S sy.l E Lxxxxd
(19)
Z = |21, Z2 2r.] E 1LxNxNxd;
(20)
B.2 COMPUTATIONAL COMPLEXITY OF DIFFERENT ATTENTION MECHANISMS
B.2.1
FULL PPATIO-TEMPORAL ATTENTION (TUEORETICAL)
The full spatio-temporal attention computes across all residues and all time frames for both singles and pairs, with computational complexity:
O(L x (N + (N+N2))
= = O(L
(21)
This is computationally prohibitive for any realistic protein system.
B.2.2 CONFROVER APPROACH: Two SOURCES OF COMPUTATIONAL COMPLEXITY
ConfRover faces computational challenges from two distinct sources:
Channie-Factorized Temporal Attention. The primary limitation in ConfRover its temporal attention mechanism. It applies temporal attention independently to each single feature and each pair feature:
For singles: N separate I I. temporal attention operations, complexity O(N 1.2)
19

---

Published as a conference paper at ICLR 2026
For pairs: 2 separate L x L temporal attention operations, complexity O(N2 x L2)
Total temporal attention complexity: O(N x L + N2 = O((N + N x L2)
While this factorization makes computation manageable, it assumes that temporal dynamics can be modeled independently per channel (whether single or pair), which fundamentally limits the model's ability to capture complex spatio-temporal couplings that govern protein dynamics.
Pairformer Operations for Spatial Modeling, In addition to the temporal attention, Con- [Rover employs Pairformer layer for spatial interactions at each time step, which has complexity:
((N3) per frame due the interaction between all residues and all pairs . O(N x L) for the entire trajectory
The combined time complexity is therefore O((N + N x L + N x L). For large proteins, the N term dominates, making scaling to large systems prohibitive both in terms computation time and memory usage.
B.2.3 STAR-MD APPROACII:
SPATIO-TEMPORA, ATTENTION
Our approach restricts spatio-temporal attention to only single-residue features:
We apply one (N L) x (N x L) attention operation on the flattened singles tensor Computational complexity: O((N x L)²) = O(
Critically, eliminate the Pairformer component and its O(N3) complexity. Instead, rely on the spatio-temporal attention mechanism to implicitly learn the necessary pairwise relationships through the Memory Inflation mechanism described in Section A.
B.3 STAR-MD Vs. FACTORIZED SPATIAL-TEMPORET ARCHITECTURES
STAR-MD employs a joint spatio-temporal attention mechanism with complexity O(N2 In contrast, some architectures utilize spatial Pairformer layers (O( and temporal attention lay- ers on singles (O(NL))) resulting in a total complexity of O(N NL²).
To identify the regime where each model is more efficient, equate their complexities:
2 = 2 + NL²
(22)
Assuming N, L. 0, can simplify this find the crossover point:
N
L
(specifically L = N N-1
(23)
This reveals two distinct regimes:
Regime 1 (L 2 N): When the context length exceeds the protein size (e.g., small peptides with very long history), the O(N temporal scaling of basclines is advantageous.
Regime 2 (N >> L): For realistic protein modeling, where system size N is large (hun- dreds to thousands of residues) and context history L fixed (e.g., 50-100 frames), STAR- MD is significantly more efficient. In this regime, the baselines" cubic spatial scaling (O(N3)) becomes prohibitive bottleneck.
As illustrated in Figure 6 for standard context length of L = 80, STAR-MD's computational cost scales much more favorably with protein size than baseline methods, enabling the simulation of large protein complexes that are computationally intractable for O(N3) methods.
B.4 KV CACHE EFFICIENCY DURING INFERENCE
An important aspect to consider in autoregressive protein trajectory generation the memory cost during inference, particularly related to the key-value (KV) cache used for efficient generation. This becomes especially important for long-horizon generation where hundreds or thousands frames must be maintained in context.
20

---

Published as a conference paper at ICLR 2026
Scaling with Protein Size C (b) Scaling with Contect Length
STAR-MD S+I (Stagles)
5-T thams 60RC kW:
1010> 11.00
proo1
2211
Si1T (Sindaa) -83)
HD o0V/27
CODe CONSWOO2 WOOSS cowdtergg STOTT dropand
10.00
230
Compatabinal
130
Se1 (Fazzl HMM + M/1
5-1 (Slegian WMC2 MP12
19
2013 30m C00
260 Lengrth
100
366 2:00 912
Figure 6: Computational complexity scaling, Comparison of computational cost (FLOPs proxy) for STAR-MD (O() Pairformer + Temporal Attention on Pairs (O(N2 + and Pairformer + Temporal Attention on Singles (O(N + L)). ). Left: Scaling with protein size N for a fixed context length
80. STAR-MD scales significantly better for large proteins due to the absence of cubic spatial terms. Middle: Scaling with context length L for fixed protein size N 256. While baselines have lower temporal complexity, their high spatial cost dominates in the realistic regime where N L. Right: Heatmap showing the speedup factor of STAR-MD over Pairformer + Temporal Attention on Singles as function of protein size N and context length L. STAR-MD achieves significant speedups in the realistic regime where N L.
ConfRover KV Cache Requirements. During autoregressive generation, ConfRover must store in its KV cache:
Singles features: O(N L d) memory
Pairs features: ONN L X d) memory
Total KV cache memory: O( + N2) L d)
For a medium-sized protein with = 200 residues, a context window of L = 32 frames, and embedding dimension d 256, this results in approximately:
Memory for each layer (200 + 200?) 32 256 bytes
(24)
= (200 + 40000) x 32 x 256 x 4 bytes
(25)
1.3 GB
(26)
STAR-MD KV Cache Requirements. for the singles features:
In contrast, STAR-MD only needs store keys and values
Memory for each layer
N x dl x 4 bytes
(27)
200 32 256 4 bytes
(28)
2 6.6 MB
(29)
This represents 196x reduction KV cache memory requirements compared to ConfRover, which becomes increasingly significant as the protein size increases when using longer context windows.
Scaling to Larger Proteins and Longer Trajectories. Figure 5 shows the estimated KV cache memory requirements for different protein sizes and trajectory lengths.
B.5 PRACTICAL SCALING BEHAVIOR
While the asymptotic complexity of STAR-MD with respect protein size N is still quadratic (12),, our approach offers significant practical advantages:
1. Elimination of cubic scaling terms: The most expensive O(N3) operations from the Pairformer eliminated.
21

---

Published as a conference paper at ICLR 2026
Simulation time
Intervals (snapshots)
Number of frames
100 ns
120
80
240 ns
120
200
250
400
Table 4: Inference task configurations for ATLAS dataset.
2. Efficient attention implementations: Our attention pattern amenable to highly opti- mized implementations like FlashAttention Dao et (2022). which provides significant practical speedups.
3. KV cache efficiency: As shown in Section B.4, our approach dramatically reduces memory requirements during inference by avoiding the need to store pair features in the KV cache.
Data & Implementation Details
C ATLAS DATASET
The ATLAS (Vander Meersche et al. 2024 dataset contains triplicated 100 simulation for 1390 proteins with diverse structures and dynamics, representative for different ECOD X-class domains. Following prior work (Jing et al. 2024a b) Wang al., 2024, Shen et al. 2025), we adopt a time- based split for ATLAS. Specifically, the train/validation/ee
sets are divided based on the release date of each protein, using cutoff dates of May 2018 and May 1, 2019. Proteins released before May 1, 2018 are used for training, and proteins released after May 1, 2019 are used for testing. The diverse nature of ATLAS makes a standard benchmark for evaluating protein conformation and trajectory generation under transferrable settings.
For each 100 ns trajectories, snapshots of atom coordinates were saved at 10ps intervals. We use the coordinates for all heavy atoms in protein for model training and evaluation. Similar to Shen et al. (205),, we exclude the training proteins longer than 384 amino acids, leading to 1080 training proteins. We include the inference task configurations for different simulation time in Table
D DETAILS ON EVALUATION METRICS
Trajectory Alignment. Prior to any quantitative analysis (including conformational coverage and kinetic fidelity), all model-generated trajectories aligned to the corresponding reference MD simulation. This is performed via Ca superposition to a common reference frame, ensuring that all comparisons independent of global rotational and translational differences.
Conformational Coverage. Conformational state recovery is evaluated by comparing the distri- bution of model-generated and reference conformations in PCA space. Each conformation is pa- rameterized by the 3D coordinates of Ca atoms. The PCA space is constructed using conformations from reference MD simulations ATLAS. To compare distributions, each principal component is discretized into 10 evenly spaced bins. After projecting the conformations into this space, we count their occurrences in each bin and compute the distribution similarity using Chssen-Shannon Distance (JSD). We also binarize the occupancy counts to compute precision, recall, and F1-score evalu- ating whether sampled conformations fall within known states, following prior works 2024
(Lu et al.
Wang et al. 2024 Zheng et al. 2024 Shen et al. 2025).
Kinetic Fidelity. Previous metrics operate on individual proteins and ensembles of proteins, inde- pendent of their temporal ordering. To characterize the temporal evolution of protein conformations, we evaluate four complementary metrics: coordinate root mean squared distance (RMSD), autocor- relation, tIA correlation (Shen al. 2025), and VAMP-2 score (Wu & Noë 2020)
22

---

Published as a conference paper at ICLR 2026
Let #1.L (#i,..... .L) denote a trajectory of L frames, where each frame 2y E R RNNm33 is the conformation at frame index L. Given a model-generated trajectory &1.L and MD reference L1.L, project all trajectories onto the top 32 principal components computed using MD reference Ca coordinates. For a particular lagtime r, the coordinate RMSD is computed via
Eg [R[SS( (e))
(30)
which is an approximate average velocity highlighting the dynamical timescales of conformational changes, Autocorrelation is computed via
(xt
(I(++ - (1) r²
Ex
(31)
which is a measure of covariance within each trajectory. The aim of each metric is to approximate how the corresponding measure behaves in the MD reference data (dashed lines), capturing not only their overall trend but also characteristic magnitudes.
For tICA correlation, follow (Shen al. (2025) and compare the principal dynamic modes of the generated and reference trajectories. Input trajectories are aligned and represented by Ca co- ordinates flattened to d 3N features. We first compute validity mask M E (0, 1)2 for the generated trajectory, retaining only frames that pass the "CA + AA Validity" checks defined in Ap- pendix Only time-lagged pairs (æz, dg-r) where both frames are valid (M₃ 1, Me++ are used for analysis. We fit separate tICA models to the reference and generated trajectories by solving C, X,Cov, where Cu and C, are the instantameous and time-lagged covariance matrices, respectively. We use lagtimes E [1,5, 10, 20), apply kinetic map scaling, and set the regularization cutoff to e 10-6 For the top independent components (PC1 and PC2), extract the left singular vector (k)) 3 and compute per-residue contribution score s(*) max ([vaaa, |0% (0,,=)) for each residue i. The final metric is the absolute Pearson correla- ref Sgen averaged across PC1 and PC2. Our evaluation procedure mirrors Shen et al. (2025), with the sole exception
tion coefficient between the reference and generated scores: |r PPersson (S(&)
that we only consider valid pairs via M. For any given lagtime, we report the tICA correlation only
if at least 30 time-lagged pairs are valid; otherwise, we denote the result as "N/A".
Finally, the VAMP-2 score measures how well a set of features captures the slow dynamical modes of the system (Wu & Noë (020). The Koopman matrix is approximated via
K = C 1/2 Co+ CT!// where Ca = Ez((zr, - -14 (2 - -1)]
(32)
and the VAMP-2 score is the squared Frobenius norm of the Koopman matrix. We compute VAMP-2 score function of lagtime, projecting all trajectories onto the top 32 principal components com- puted using MD) reference Co coordinates, consistent with the RMSD) and autocorrelation metries.
Structural Quality and Validity To assess the physical plausibility of generated protein confor- mations, we use a suite of metrics that measure different aspects of structural integrity. A generated frame is considered structurally valid if it simultaneously satisfies all of the following conditions, which are based on thresholds derived from oracle MD simulations:
No Ca Chain Breaks: The distance between consecutive C, atoms must be below a certain threshold.
No C. Clashes: threshold.
The distance between any two non-adjacent C. atoms must be above a certain
Plausible Backbone Geometry: The percentage of residues that are outliers in the Ramachandran plot must be below threshold.
Plausible Side-Chain Geometry: The percentage of residues with outlier rotamers must be below a threshold.
The specific thresholds for these metrics were determined by analyzing the distribution of these quality metrics in the soound-truthh 100 ns MD simulation trajectories (our "oracle"). For each met- ric, we computed its value across all frames of the oracle trajectories and set the threshold at the 99th percentile (i.e., accepting 99% of the oracle frames). An exception made for the C., chain break
23

---

Published as a conference paper at ICLR 2026
rate. The oracle MD simulations contained chain breaks, resulting in 99th percentile threshold of O. This too strict for current generative models. To allow for a more meaningful comparison, set small, non-zero tolerance for chain breaks.
The exact thresholds used for determining a valid frame are listed in Table 5 We define three sets of criteria for validity: Ca All-Atom, which requires passing all four checks; All-Atom Only, which checks only Ramachandran and rotamer criteria; and Ca-Only, which checks only for Co clashes and breaks.
Table 5: Structural quality thresholds for validity. Thresholds were derived from the 99th per- centile of the oracle 100 MD simulations, except for the C. break rate.
Metric
Threshold
Co + All-Atom Evaluation
Ramachandran Outliers (%)
< 4.12
Rotamer Outliers (%)
< 7.05
C, Clash Rate (%)
S 1.29
C. Break Rate (%)
S 0.2
All-Atom Only Evaluation
Ramachandran Outliers (%)
< 4.12
Rotamer Outliers (%)
S 7.05
Ca-Only Evaluation
C Clash Rate (%)
5 1.29
C, Break Rate (%)
$ 0.2
24

---

Published as a conference paper at ICLR 2026
Architecture Details
E
ARCHITECTURE DETAILS: PAIR FEATURES
STAR-MD modifies Shen et al. 2025 's standard AlphaFold-derivee architecture by removing the
Pairformer module to reduce computational cost, while retaining explicit pairwise feature updates.
FrameEncoder: STAR-MD takes as input the pre-trained pairwise features from frozen Open- Fold (Jmper et al. 2021 model. Then, employs a FrameEncoder module [Shen al. (2025) which incorporates pairwise geometric information into the pair features.
EdgeTrransiieg Transitoon. Pairwise features E R-VxNxd: are initialized by the FrameEncoder and updated via EdgeTransition layers. Unlike the global attention in Pairformer, EdgeTransition uses local MLP update after each spatio-temporal attention block. For residues and j:
Zij - Zij + MLP(s;, Zig)
(33)
This maintains spatial context without O(N3) complexity.
Attention. Joint spatio-temporal attention applied exclusively single-residue features 8. This avoids the ((2²T) memory cost of temporal attention over pairs. Pair features modulate the attention mechanism solely through bias terms, similar Invariant Point Attention (IPA).
Singles-Only
25

---

Published as a conference paper at ICLR 2026
F BASELINE IMPLEMENTATION
MD oracles. In 100 ns simulation, we use one of the triplicated trajectories in the ATLAS dataset as the oracle estimate the expected performance from an independent simulation run. To simulate longer trajectories, we follow the setup described in Vander Meersche et al. (2024) using provided equilibrated structures and Gromacs trr files to preproduce production runs at the desired lengths. All simulations are conducted using Gromacs (version 2023.2) 2) on NVIDIA V100 GPUs with 40 GB memories. Independent reference and oracle trajectories are generated by repeating simulations with different random seeds.
AlphaFolding (Cheng et al., 2025). We followed the official data preparation and in- ference procedures provided by the authors | using the recommended model weights trapolation script (rum.eval.extraolaatil n.motion=2, and n.frames=16. 6. This model employs block-wise autoregressive scheme: it generates fixed-size block of 16 frames conditioned on the final frames of the preceding block. This process is repeated to achieve the desired total simulation time by adjusting the extrapol at contime parameter.
pth). Trajectories were generated using the recommended ex-
sh) with parameters sample.step-1
AlphaFolding has fixed time step of p0, which differs from our model's variable stride capability. Generated trajectories were sub-sampled to match the frame rates required for evaluation.
MDGen [Jing et al. 2024b). We followed the official data preparation and inference procedures provided by the authors using the provided model weights (atlas catlas.ckpt) ckpt). The model was trained on ATLAS trajectories preprocessed with 400 ps intervals (stride 40 with base interval 10 ps). For inference, we used the forward simulation script (s (sim-inference, py) with pa- rameters num frames-22 250 and suffix- R1.
This model generates trajectories autoregressively in blocks of 250 frames with internal timestep of 400 ps per frame. A single rollout produces 100 ns of simulation time (250 frames x 400 ps ns). For 100 trajectories, we used num rollouts- we extended generation by setting num rol.rolllle louts to the required number of sequential blocks. Generated trajectories were then sampled post-generation to match the frame rates required for eval- uation.
For 240 ns and 1 ls trajectories,
ConfRover (Shen et al. 2025). We use the code and model weights provided by the authors. ConiRover is a frame-level autoregressive model similar to STAR-MD, and we adopt the same setup as our model to generate trajectories for each experiments. For 100 ns results, we enable CPU- offloaded caching to store full key-value history in system memory, while for 240 ns and uu results, we employ a sliding-winde window cache (size=14) with attention sinks (first two frames) (Xiao et al, 2023a), labeled ConfRover-W. W. This reduces its temporal complexity from ( OlI W) and prevents out-of-memory errors. Despite slight quality drop, ConfRover-W maintains comparable performance to the full-attention model, see Figure 7 and Table 6 for the validation on the 100 ns benchmark.
All generative model baselines are evaluated NVIDIA H100 GPUs with 80 GB GPU memory and 2 TB system memory.
http://hh.ttitiiiteee//////////t diffusion hhtps://gthuuh.om///ing000 6/mdgen
26

---

Published as a conference paper at ICLR 2026
CA AA Validity (%)
100
ConfRover
ConfRover-W
0
20
40
60
80
Time (ns)
Figure 7: Validation of ConfRover-W at 100 ns. Comparison of Structural Validity (Ca + All- Atom) between the standard full-attention ConfRover and the windowed ConfRover-W. The win- dowed variant maintains comparable performance, validating its use as a proxy for long-horizon experiments.
Table Comparison of ConfRover and ConfRover-W V at 100 ns. Quantitative metrics show that the windowed approximation (ConfRover-W achieves performance comparable to the full- full-attention model (ConfRover), justifying its use for longer horizons.
Cov Valid
Dyn.
Validity
Model
JSD
Reet
tICAt
CA %
AA%1
CA+AA% %6
ConfRover
0.52
0.36
0.15
56.43
92.37
52.02
ConfRover-N
0.51
0.37
0.15
50.08
89.05
45.70
STAR-MD
0.42
0.57
0.17
86.83
98.23
85.35
27

---

Published as a conference paper at ICLR 2026
G MODEL HYPERPARAMETERS
We include the details of our training and inference configuration in Tables to TO The model was trained on the ATLAS dataset following the train/val/test split in the previous works(Shen et al.) 2025 Jing et al. 2024a by Wang et al. 2024 Cheng et al. (2025). STAR-MD contains 4 diffusion blocks, each containing I IPA layer and SxT layers with hidden dimension 256 and 8 attention heads (Table 9). We train on trajectory snippets with context length L = 8 frames and global batch size of (Table 7).
We use the Adam optimizer and learning rate 5 10-5 with the similar loss setup as used in Shen al. (0255) (Table TOP. Diffusion parameters are provided in Table g
For distributed training, we employed DeepSpeed Stage 2 optimization with gradient checkpointing efficiently manage memory usage during training of large protein systems. All models are trained 8 NVIDIA H100 GPUs with 80 GB memory.
Table 7: Dataset and Training Configuration
Parameter
Value
Dataset Configuration
Training data
ATLAS train split (length < 384)
Global batch size
Frames per trajectories
Frame intervals
1~ 2 of 10 ps snapshots
Feature Representation
Single-residue feature dim
384
Pairwise feature dim Number of recycles
128
Data Augmentation
Clean noise max magnitude Clean noise sampling prob Noise sampling
0.1
0.75
Frame-level
Table 8: Diffusion Model Configuration
Parameter
Value
SE(3) Diffusioo
Coordinate scaling
0.1
Translation Buax
0.1 / 20.0
Translation schedule
Linear
Rotation /a max
0.1/ 1.5 Logarithmic
Rotation schedule
Sampling
Method
Euler SDE
Diffusion steps
200
max mmin
1.0 0.01
28

---

Published as a conference paper at ICLR 2026
Table 9: Model Architecture
Component
Configuration
Encoder (PseuddBecaPairnnccdee)
Residue index embedding size
128
Output size
128
Invariant Point Allention
IPA blocks
4
Single channel size
256
Pair channel size
128
Hidden channel size
256
Skip channel size
64
Attention heads
4
Query/key points
8
Value points
12
Spatiotemporal Attention
Model dimension
256
Number of layers
2
Number of heads
4
RoPE 2D
Enabled
Table 10: Loss Functions and Optimization
Parameter
Value
Loss Weights
Rotation score loss
0.5
Translation score loss
1.0
Torsion loss
0.5
Backbone FAPE loss
0.5
Sidechain FAPE loss
0.5
Backbone coordinates loss
0.25 (diffusion t < 0..55)
Backbone distance map loss
0.25 (diffusion f < 0.25))
Optimization
Optimizer
Adam
Learning rate
0.0002
LR schedule
Linear warmup + cosine decay
Warmup epochs
Total epochs
250
Warmup start LR factor
0.01
Minimum LR factor
0.1
Gradient clipping
1.0 (norm)
Precision
BF16 mixed precision
29

---

Published as a conference paper at ICLR 2026
Extended Results
H EXTENDED ATLAS ANALYSIS
H.1 ADDITIONAL RESULTS ON DYNAMIC FIDELITY
We visualize the corresponding kinetic metries for 100 ns, 240 ns, and us in Figure For each metric, we include 2 separate MD simulations for reference (dashed lines). In all but one case (autocorrelation for (s),, STAR-MD better approximates the trend and values of MD references, showing higher kinetic fidelity than the rest of models.
We also include the corresponding metrics for our ablation study (Section 4.5) Figure g where STAR-MD outperforms other variants in all but two cases (100 ns Ca coordinate RMSD and js Ca autocorrelation at large lagtime).
We further provide detailed breakdown of the tICA correlation metric in Figure 10 and Table I1 Figure I0 shows the per-component and per-lagtime correlations for all valid runs, highlighting the variability across different simulation instances. Table [1] summarizes these results, demonstrating that STAR-MD achieves the highest mean correlation among generative models, closely matching the reference MD baseline.
Table 11: Aggregated tICA Correlation Statistics. Mean and standard deviation of tICA correla- tions averaged over PCI, PC2, and all lagtimes. We report results on both valid samples (Valid) and all samples (Unfiltered).
Model
tICA (Valid)
tICA (Unfiltered)
MD (Oracle)
0.170
0.17
MDGen
0.130 0.039
0.12
AlphaFolding
N/A
0.
ConfRover
0.162 0.035
0.18 1001
STAR-MD
0.176 + 0.033
0.18
We also report VAMP-2 score as function of lagtime in Figure II STAR-MD better approximates the VAMP-2 score of MD) references for all lagtimes compared to other models and ablation variants, demonstrating better capture of slow dynamical modes.
30

---

Published as a conference paper at ICLR 2026
MD Reference 1
MDGen
ConfRover STAR:
Reference
AlphaFolding
MD Reference
MDGen
Confeoeer STAR-MD
Reference 2
AlphaFolding
0.8
0.7
autocorrelato latto
0.5
RMSD
0.4
0.3
0.2
0.1
0.0
12
15
18
12
18
Lagtime (frames)
Lagtime (frames)
(a) Ca coordinate RMSD vs lagtime (100 ns simulation) (b) Ca aulocorrelation vs lagtime (100 ns simulation)
ConfRover
MD Reference 2
AlphaFolding
STAR-MD
Reference 1 MD Referencc
MDGen
ConfRover STAR-MD
AlphaFulding
RMSD (A)
10a
autocorrllaton
0.5
0.4
0.3
0.2
0.1
0.0
18
12
15
18
Lagtime At (frames)
Lagtime At (frames)
(c) Ca coordinate RMSD vs lagtime (240 ns simulation) (d) Ca aulocorrelation vs lagtime (240 ns simulation)
MD Reference
MDGen AlphaFolding
ConfRover
Reference 1
MDGen AlphaFolding
ConfRever STAR-MD
MD Referencce
0.8
10F
autocoreeltionn
0.6
RMSD
6x22-
0.2
1x10"
2x20
0.0
12
18
12
15
10
Lagtime At (frames)
Lagtime (frames)
(e) Ca coordinate RMSD lagtime (1 μs simulation).
(l) Ca autocorrelation vs lagtime (1 simulation).
Figure 8: Kinetic metrics at different trajectory lengths: 100 ns, 240 ns, and 1 js.. Dashed lines
MD references (in most cases MD reference metrics very close, demonstrating the metrics are robust across independent MD reference simulations). STAR-MD better approximates the MD references, demonstrating higher kinetic fidelity than the other models.
31

---

Published as a conference paper at ICLR 2026
Reference STAR.MD
Preproc
w/o Noise
Reference
w/ Sep Altn
w/ Preproc Altn
STAR-MD
Noise
0.7
0.6
RMSD (A)
autocorrelation
0.5
0.4
0.3
0.2
0.0
0.3
12
12
15
18
Lagtime At (frames)
Lagtime At (frames)
(a) Co coordinate RMSD vs lagtime ns simulation) (b) Co autocorrelation vs lagtime (100 ns simulation)
MD Reference
w/ Sep Altn
w/ Preproc Attn
STAR,MD
MD Reference
w/ Sep Altn
w/ Preproc Altn
STAR-MD
w/o Noise
0.7
autocorelatio raat
0.6
RMSD
0.5
0.4
0.2
0.3
0.1
0.0
12
15
Lagtime At (frames)
Lagtime At (frames)
(c) Co coordinate RMSD vs lagtime (240 simulation) (d) Ca autocorrolation vs lagtime (240 ns simulation)
Reference
Preproc
STAR,MD
w/o Noise
Reference
w/ Sep Altn
w/ Preproc Altn
STAR-MD
Noise
0.8
0.8
0.7
0.7
RMSD (A)
locoreeattoo latton
correaation
0.6
0.5
0.4
0.3
0.4
0.2
0.3
0.1
12
12
15
Lagtime At (frames)
Lagtime At (frames)
(c) Ca coordinate RMSD vs lagtime (1 jus simulation)
(f) Ca autocorrelation vs lagtime (1 us simulation)
Figure 9: Kinetic metrics for ablation study. Dashed line is MD reference. STAR-MD has simi- lar/better performances than other variants in all but two cases (100 ns Ca coordinate RMSD and uus Co autocorrelation at large lagtime).
32

---

Published as a conference paper at ICLR 2026
Correlations
Correlations
MD 100rs MDGen ConfRover
ND 100es
MDOmn
ContRover
0.25
0.25
Pearson Irl
Pearson Irl
0.20
0.20
0.15
0.35
6.10
0.30
0.05
0.45
10.0
15.0
20.0
5.0
10.0
15.0
17.5
20.0
(frames)
Lagtima (frames)
Figure 10: Detailed tICA Correlation Analysis. We report the per-component (PC1, PC2) and per-lagtime tICA correlation values for runs that produced least 30 valid pairs for computation. Shaded regions represent the standard deviation across valid runs.
MD Reference
MDGen
CanfRaver
AlphaFolde BohaaFodiing
ND Referance STAR-MD
W/ Sop Aztn
w/ Proproc Attn
VAMP-2 score
19
VAMP-2 score
19
16-
15
12
15
18
Lagtime (frames)
Lagtime At (frames)
(a) Comparison with other models
(b) Ablation study
Figure 11: VAMP-2 score function of lagtime (100 simulation). Dashed line is MD reference. STAR-MD better approximates VAMP-2 score of MD reference than other models and variants, demonstrating better capture of slow dynamical modes.
33

---

Published as a conference paper at ICLR 2026
H.2 FULL ABLATION RESULTS
We include the full results Tables for ablation study in this section (Table 12 We observe consistent trend that contextual noise improves structure quality, and spatio-temporal attention improves conformation coverage.
Figure compares the CA+AA validity over time for models with and without contextual noise on the 240 ns and benchmarks, showing that the noise helps maintain structural quality over long horizons.
100
CA + AA Validity (%)
80
60
40
STAR-MD
240 ns
20
Noise
1 1s
200
400
600
800
1000
Time (ns)
Figure 12: Effect of Contextual Noise Long-Horizon Stability. Comparison of Structural Validity (Ca + All-Atom) over time for STAR-MD with and without contextual noise on 240 ns and 1 js benchmarks.
Complete ablation study results (240 ns). We compare our full model, STAR-MD, against variants with key components removed. All metrics reported for the 2s generation lask. Metrics degraded in each ablation setting are highlighted red.
Table 12:
Cov Valid
Dynamic Fidelity
Validity
Model
JSDJ
Reet
RMSD]
AutoCor
CA%
AA%t
CA+AAST
STAR-MD
0.45
0.59
0.20
0.02
86.44
97.61
84.62
w/o Noise
0.42
0.61
0.22
0.16
77.11
96.19
74.25
w/ Sep Attn
0.48
0.51
0.26
0.04
87.09
97.58
84.83
Preproc Attn
0.47
0.52
0.25
0.03
87.25
96.72
84.28
Table 13: Complete ablation study results (1 ps). We compare our full model, STAR-MD, against variants with key components removed. All metrics are reported for the μs generation task. Metrics degraded in each ablation setting highlighted in red.
Cov
Valid
Dynamic Fidelity
Validity
Model
JSDJ
Rec
RMSD
AutoCor
CA
AA%t
CA+AAZ1
STAR-MD
0.45
0.61
0.11
0.11
91.00
91.44
83,59
w/o Noise
0.48
0.63
0.17
0.10
75.84
91.22
68.66
w/ Sep Attn
0.51
0.49
0.30
0.10
93.12
92.03
85.31
w/ Preproc Attn
0.52
0.49
0.32
0.11
94.12
92.19
86.44
34

---

Published as a conference paper at ICLR 2026
H.3 100 NS FILTERED RESULTS
H.3.1
PERFORMANCE BY SEQUENCE LENGTH
We break down the performance of STAR-MD and baselines across different protein sequence length buckets analyze scalability and robustness. Tables 14 to summarize these findings.
Table Sequence Length Bucket: 0-100 (17 chains)
Cov Valid
Dyn.
Validity
Model
JSDJ
Reet
tICA
CA %
AA% %
CA+AA 1
MD (Oracle)
0.28
0.71
0.10
100.00
93.24
93.24
MDGen
0.48
0.38
0.11:001
$5.662111
87.34-11.6
75.91
AlphaFolding
0.56
0.28
N/A
30.07
3.60
2.06
ConfRover
0.48
2.39
0.11:1001
79.59=2.0
18
65.26:11.00
STAR-MD
0.37 00001
0.56::0.02
0.11+000
91.37
93.75++047
86.10
Table 15: Sequence Length Bucket: 100-150 (16 chains)
Cov Valid
Dyn.
Validity
Model
JSDJ
Reet
IICA
CA %1
AA%
CA+AA %-
MD (Oracle)
0.30
0.69
0.12
98.98
98.44
97.42
MDGen
0.5
0.35
0. 0.10
74.50=7.40
94.17+1.99
70.624762
AlphaFolding
0.58
0.10
N/A
14.38
0.23
0.16
ConfRover
0.53+4o.1
0.35 +0.01
0.12
55.89
86.38++078
STAR-MD
0.49
0.53
0.12
79.05
98.25
77.64
Table 16: Sequence Length Bucket: 150-225 (11 chains)
Coy Valid
Dyn.
Validity
Model
JSDJ/
Reet
ICAA
CA% %1f
AA%
CA+AA %6
MD (Oracle)
0.27
0.66
0.16
90.57
99.66
90.23
MDGen
0.52
0.32
0. 0.122am
66.8625.84
97.18 97.181991
65.23
AlphaFolding
0.72
0.10
N/A
4.55
0.34
0.1
ConfRover
0.47+40.2
0.39
0.15
73.39
98.77 4:034
72.43++244
STAR-MD
0.37
0.53
0.16
89.5
99. 99.82-+e22
89,41+1.81
35

---

Published as a conference paper at ICLR 2026
Table 17: Sequence Length Bucket: 225-400 (15 chains)
Cov Valid
Dyn.
Validity
Model
JSD
Rect
tICAt
CA %6
AA %1
CA+AA %6
MD (Oracle)
0.33
0.65
0.22
100.00
99.90
99.90
MDGen
0.62=a001
0.21 +0.01
0.14+a00
61.69=4.22
97.71+106
60.40:4475
AlphaFolding
N/A
N/A
N/A
1.15
0.00
0.00
ConfRover
0.53
0.36
0.22
54.00
99 .4..44-0008
53.77:2295
STAR-MI
0.44
0.50
0.22
82.38
99,9810M
82.37
Table 18: Sequence Length Bucket: 400+ (23 chains)
Cov Valid
Dyn.
Validity
Model
JSDJ
Reef
UCAA
1
AA%t
CA+AA %
MD (Oracle)
0.33
0.66
0.24
99.86
100.00
99.86
MDGen
0.65=0001
0.16=am
0.14+001
66.75-3.69
99.81+0.1
66.67
AiphaFolding
N/A
N/A
N/A
0.97
0,00
0.00
ConfRover
0.57.10.1
0.32
0.261001
1.14
99.61:001
28.54
STAR-M. RD
0.45
0.54
0.24
90.92
100.00
90.92
H.3.2 GENERALIZATION TO DISSIMILAR PROTEINS
To strictly evaluate generalization to unseen protein folds, we report results excluding two test pro- teins (7buy A and 7e2s A) that have high sequence similarity (> 80%) to the training set. Table 19 summarizes these findings.
Table 19: Non-similar Proteins, Excluding 7buy A, 7e2s A (80 chains)
Cov Valid
Dyn.
Validity
Model
JSDJ
Rech
tICA
CA 1
AA %
CA+AA %
MD (Oracle)
0.31
0.67
0.17
98.37
98.07
96,43
MDGen
0.56
0.28
0.1200000
71,83
95.03
68.311220
AlphaFolding'
0.58
0.22
N/A
10.98
0,92
0.52
ConfRover
0.52.1001
0.36
0.15
56, 06.94
92.47
52.06
STAR-MD
0.43
0.54=
0.17::n.00
86.81
98. 18:108
85.29
36

---

Published as a conference paper at ICLR 2026
H.4 ADDITIONAL VISUALIZATIONS OF CONFORMATIONAL COVERAGE
Figures 13 t2 show conformational coverage of 10 proteins from the test set. These visualizations follow the same format as Figure b in the main text (i.e., we filter the generated conformations to show only those that are structurally valid (passing all validity checks)). This provides a more rigorous view of the useful conformational space explored by each model.
(R: 0.65, 1.00)%
MDGen (R. 0.50, V. 29%)
AlphaFaiding IR: 0.00, V. 0%5
PC2
MDD 2 (R: 0.55, V: 100%)
PC2
ConfRover (R: O 55, V. 94%)
STAR MD (R: 0.60, V. 99%)
100%)
PC2
PC1
PC1
PC1
Figure 13: Conformational coverage for 6kty A (Valid conformations only).
1 0e: 0.55, V: 100%)
MDXaen (R: 0.15, 50%)
AlphaF-iding (R: 0.00, V: 0%%
PC2
2 (R: 0.70, 1000)
PC2
ConfRover (R: O 30, 38%)
STAR MD (R: 0.45, V: 99%)
(P: 0.60, 100%)
PC2
PC1
PC1
PC1
Figure 14: Conformational coverage for 6jwh_A (Valid conformations only).
37

---

Published as a conference paper at ICLR 2026
0.85, 100%)
MDGen (R 0.45, 50%)
AlphaFalding (R: 0.00, V. 0%6
PC2
MD 2 (R: 0.90, V: 100%)
PC2
Confiover (R: a 55, 3195
STAR MD (i: O 0.80, V: 95%)
PC1
PC1
Figure 15: Conformational coverage for jjpt-AA (Valid conformations only).
0.75, 100%)
MDGen (R. 0.35, 78%))
AlphaFolding (R: 0.25, V. 4%5
PC2
MD 2 (R: 0.60, V: 100%)
RC
ConfRover (R: O 40, 98%)
STAR MD (R: 0.50, V: 86%)
PC1
PC1
Figure 16: Conformational coverage for Sokd_C (Valid conformations only).
38

---

Published as a conference paper at ICLR 2026
(R: 0.60, 100%)
MDGen (R. 0.45, 99%)
AlphaFaiding (R: 0.00, V. 0%)
PC2
MD 2 (R: 0.85, V: 100%)
PC2
Confitvver (R: O 45, V: 80%)
S STAR-D (R: 0.50, V. 100%)
PC1
PC1
Figure 17: Conformational coverage for 6odd.B (Valid conformations only).
0.70, 100%)
MDGen (R. 0.20, V. 54%))
AlphaFoiding IR. 0.00, V. 0%5
PC2
MD 2 (R: 0.75, V: 100%)
PC2
ContRover (R: O 20, V. 21%)
STAR-D (R: 0.55, V. 100%)
100%)
PCZ
PC1
PC1
Figure 18: Conformational coverage for 7 (Valid conformations only).
39

---

Published as a conference paper at ICLR 2026
(R: 0.80, 100%)
MDGen (R. 0.70, V. 46%)
AlphaFolding IR: 0.00, V. 0%6
PC2
MD 2 (R: 0.65, 100%)
PC2
Confitvver (R: a 55, 62%)
STAR MD dR: 0.85, V: 78%)
PC1
PC1
Figure 19: Conformational coverage for 7 71p1 AA (Valid conformations only).
0.90, 100%)
MDGen (R. 0.35, 99%)
AlphaFoiding (R: 0.00, V. 0%5
PC2
MDD 2 (R: 0.40, V:: 100%)
ConfRover (R: 0.45, 82%)
STAR MD (R: 0.60, V: 9950
PC1
PC1
Figure 20: Conformational coverage for 6134_A (Valid conformations only).
40

---

Published as a conference paper at ICLR 2026
(R: 0.75, 100%)
MDGen (R. 0.15, 94%)
AlphaFalding IR: 0.00, V. 0%6
PC2
MD 2 (R: 0.75, V: 100%)
PC2
Confiover (R: O 45, V. 12%)
STAR-D (R: 0.80, 100%)
PC1
PC1
Figure 21: Conformational coverage for 5znj (Valid conformations only).
0.85, 100%)
MDGen (R. 0.25, 575)
AlphaFoiding (R: 0.00, V. 0%5
PC2
MDD 2 (R: 0.85, V: 100%)
CorfRover (R: 0.45, 74%)
STAR- (R: 0.75, V: 7150
PCZ
PC1
PC1
Figure 22: Conformational coverage for 6tly_A (Valid conformations only).
41

---

Published as a conference paper at ICLR 2026
1 STABILITY AND SCALABILITY ANALYSIS
I.1 EXTENDED STABILITY ANALYSIS
To further characterize the stability of STAR-MD, we provide a detailed breakdown of the vari- ability in structural validity (Ca All-Atom) for the 100 ns, 240 ns, and I js settings. Figure 23 decomposes the total variability into inter-seed and inter-protein components.
CA + AA Validity (%)
MD (Oracle)
MDGen
AlphaFolding
ConfRover
STAR-MD
100
50
50
100
0
100
200
500 Time (ns)
1000
Time (ns)
Time (ns)
(a) Per-seed variability
CA + AA Validity (%)
MD (Oracle)
MDGen
AlphaFolding
ConfRover
STAR-MD
100
100
100 Time (ns)
500 Time (ns)
1000
Time (ns)
(b) Per-protein variability
CA + AA Validity (%)
MD (Oracle)
MDGen
AlphaFolding
ConfRover
STAR-MD
100
50
50
100
0
100
200
500
1000
Time (ns)
Time (ns)
Time (ns)
(c) Combined variability
Figure 23: Breakdown of stability variability. We analyze the standard deviation of the Structural Validity (Co + All-Atom) metric over time. (a) Variability across random seeds for single protein (as reported in the main text). (b) Variability across different test proteins for single seed. (c) Combined variability incorporating both seed and protein variance.
42

---

Published as a conference paper at ICLR 2026
1.2 EXTREME-
-XTREME-HORIZON HORIZON GENERATION (10
To probe the limits of STAR-MD's s stability, we generated trajectories using the maximum training stride (s)0 for 1000 steps, resulting in total simulation time of approximately us. Table 20 reports the structural validity metrics averaged over the full 10 trajectory for the ATLAS test set. Figure 24 visualizes the stability of structural validity over the course of the generation.
Table 20: Structural validity for 10 us trajectories. Metrics are averaged over 1000 steps (~10 ns stride) for the ATLAS test set.
Metric
Validity (%)
CA Validity
85,21
AA Validity
90,93
CA AA Validity
77.28
CA + AA Validity (%)
100
75
50
STAR-MD
25
0
2000
4000
6000
8000
10000
Time (ns)
Figure 24: Structural validity over 10 j... The average CA AA Validity remains stable throughout the 1000-step generation process.
43

---

Published as a conference paper at ICLR 2026
1.3 INFERENCE COST
We evaluated the wall-clock time for generating 100 trajectories under our 100 ns benchmark setting for each methods, across proteins of varying lengths (50~750 amino acids). To estimate the runtime of molecular dynamic (MD) simulation, we used OpenMM with the Amber force field and implicit solvent, running a ns simulation and extrapolating the cost 100 ns. MDGGn, trained with a time step of 400 ps and window of 250 frames, can directly generate 100 ns trajectories in single joint diffusion process and therefore provides the greatest acceleration. AlphaFolding, trained on a fixed time step of 10 ps and a window of 16 frames, requires substantially longer inference time to complete a 100 ns simulation. Both ConfRover and STAR-MD support varying- -triing--tridee simulation and can generate under our evaluation configuration (stride=1.3 slightly slower than ConfRover for small proteins, it is more efficient when scaling to larger proteins and is faster 6SMS-A (724 AA). All methods evaluated on single NVIDIA H100 GPU.
80 frames). While STAR-MD is
Table 21: Inference cost for different methods. The wall-clock time (in seconds) to generate 100 ns trajectories across proteins with varying lengths (shown in parentheses). AlphaFolding encounters out-of-memory eтror for larger proteins and is reported as 'N/A', All methods are evaluated on single NVIDIA H100 GPU.
60KD-C (51)
6Q9C-A (155)
6XB3-H
7AQX-A (364)
7P41-D (448)
7MF4-A (554)
6L8S-A (650)
6SMS-A
Method
(241)
(724)
MD
6322.8
9610.8
14420.3
24272.4
34850.8
49294.4
66296.7
78451.4
AlphaFolding
3135.7
5846.5
10024.8
19036.3
27340.3
40472.5
N/A
N/A
MDGen
4.1
8.1
13.0
20.5
28.2
35.3
41.2
54.4
ConfRover
404.5
423.6
463.5
573.7
710.4
934.4
1170.2
1617.7
STAR-MD
766.3
807.4
818.4
795.8
902.5
1055.5
1232.4
1381.2
J APPLICATIONS TO OTHER SYSTEMS
J.1 ENSEMBLE GENERATION
We examine the performance of STAR-MD to approximate target ensemble distributions from MD simulation, and compare it against the state-of-the-arr ensemble emulator BioEmu (Lewis et al.]
2024). Due to several differences in training dataset and paradigms, we evaluate them on two bench- marks: (1) the standard ATLAS ensemble benchmark used in prior works (Ting et al. 2024a b). and (2) the CATHI ensemble benchmark introduced in BioEmu.
We follow the standard protocol from from/Jigg et al. 2024a), sampling 250 conformations for each the 82 test proteins and comparing the generated ensembles with MD reference ensembles. For STAR- MD, we repeat the 100 ns simulation procedure five times and randomly sample 250 conformations from the collected set. As shown in Table 22 STAR-MD better captures the expected diversity, more accurately reflects both global and local flexibility, and matches the reference ensemble at distributional and unsemble-observable levels. In contrast, BioEmu tends to produce overly diverse samples and shows lower similarity to the reference ensemble across most metrics.
We then evaluate ensemble generation performance the CATHI benchmark from BioEmu. This benchmark contains the longer trajectories at microsecond level for 17 short proteins (< 75 amino acids). We follow BioEmu's setup by sampling 10,000 conformations per system and projecting them onto the reaction coordinates to estimate free-energy and compute metrics using author's 's eval- uation code. For STAR-MD we sample from five different starting frames of the initial MD runs in the datasets, generating 80 frames at a stride of 1.28 ns. We repeat this process 25 times to col- lect 10,000 samples per protein.
shown in Figure 25 BioEmu achieves higher accuracy with lower error and higher coverage. When compared the ground-truth distributions, STAR-MD centrates around the low energy basins, while BioEmu explores a broader range of conformations, occasionally extending beyond the region supported by the ground truth.
To verify this performance difference is not caused by the forward simulation setup used for STAR- MD, we also test unconditional generation -STAR-M--HR by inferring single frames from masked token, similar to (Shen et al., 2025). The similar results between STAR-MD and STAR-MD-in sug- gest that the performance gap between STAR-MD and BioEmu CATH1 likely arises from factors
44

---

Published as a conference paper at ICLR 2026
other than the sejeettoy--gneraatee archiecture, such as differences in training datasets, training objectives, or fine-tuning procedures.
STAR-MD
BioEmu
RMSF (=1.63)
1.38
2.87
Pairwise RMSD (=2.76)
2.34
4.51
Pairwise RMSD r (1)
0.52
0.33
Global RMSF (1)
0.56
0.52
Per target RMSFr
0.86
0.82
RMWD (L)
2.85
4.33
MD PCA W2 (2)
1.50
1.64
Joint PCA W2 (1)
2.33
3.43
sim > (f)
35.4 %
26.8 %
Weak contacts (f)
0.52
0.49
Transient contacts J (1)
0.37
0.37
Exposed residue J (1)
0.57
0.54
Exposed MI matrix tho (1)
0.24
0.26
Table 22: ATLAS ensemble results. Reported metrics include diversity, flexibility accuracy, dis- tributional similarity, and ensemble observables. Parentheses indicate whether higher/lower values are the better, or the expected value estimated from the ground truth trajectory. The better value shown in bold.
BioÊmu
STAR-MD
STAR-MD-iid
MAE mkcal/mo))
1.00
0.00
RMSE (Icalmmol)
1.00
0.00
1.00
coverage
0.50
coth1 ceth1_bblng 1.b00022
cath1, 1c75A00
ceth1 L1:jA06
catn), Iw01A00
cath1, 121vA00
coth1, 2dkzA01
coth1, 2hbpA00
cath1, cath1_2heoss 2he0A00
coth1 2wg5502
coth1, 3a52D02
cath1 3371100
cath1, 310yA02
coth1 30ssC00
cath1, 3u6cA02
cath1, 4a53A01
cath1, 4jr1000
cath1, 4096A01
meen
Figure 25; Free energy accuracy and coverage from the BioEmu CATHI ensemble benchmark.
45

---

Published as a conference paper at ICLR 2026
reference
8735 samples
9644 samples
9280 samples
4.0
6
cath1 1bl0A02
5
0.0
3
-2.0
2
-5.0
0.0
5.0
10.0
5.0
0.0
5.0
10.0
-5.0
0.0
5.0
10.0
-5.0
0.0
5.0
10.0
7.5
7915 samples
9817 samples
9928 samples
5.0
cet11 1c73A00
0.0
bay
2
-5.0
0.0
0.0
0.0
samples
samples
samples
5.0
2.5
0.0
2..5
55.0
0
9787 samples
9914 samples
9896 samples
1zivA00
4.0
2.0
0.0
2..0
-10.0
-5.0
0.0
5.0
110.0
-5.0
0.0
5.0
10.0
5.00
0.0
5.0
-10.0
.5.0
0.0
5.0
9719 samples
9958 samples
9952 samples
5.0
3371A00
2.5-
2.55
-5.0
-5.0
0.0
5.0
-3.0
0.0
5.0
-5.0
0.0
5.0
-3.0
0.0
5.0
Figure 26: Example free emergy surfaces from the BioEmu CATHI ensemble benchmark. Five systems randomly selected for visualization. The left columns shows the reference MDD distribu- tions provided in BioEmu. For each model and system, the number of valid samples (out of 10,000) that passes BioEmu's quality filter is indicated in the upper-right corner of each subplot.
46

---

Published as a conference paper at ICLR 2026
J.2
CASE STUDIES ON FUNCTION-RELATE CONFORMATIONAL DYNAMICS
J.2.1
CRYPTIC-POCKET BINDING
Sampling conformational changes during ligand binding is challenging, especially for proteins with cryptic pockets, where cooperative and complex pocket dynamics are required to enable binding. To assess whether STAR- RMD can sample bound-state (holo) conformations from unbound-state (apo) structures in such settings, we selected three cases from BioEmu's cryptic-pocket dataset: Adeny- losuccinate Synthetase (AdS), an enzyme has several allosteric competitive binding substrates; TEM-1 Laattamase (TEM-1), forming complexes with allosteric inhibitors; and Adenylate Kinase (AdK), highly flexible protein that undergoes large structural arrangement each closed holo conformation with covered 'lid".
For each of the case, initialize from the provided apo structure and generate 200 frames stride of 2.56 ns, approximating 500 ns simulation. We compare sampled conformations with the reference holo structure and compute RMSD over the local binding residues involved in the con- formational c The best holo samples with minimum local RMSD are shown in Figure 27 Across all three cases, STAR-MI successfully samples holo-like conformations with RMSD < T.5 Á (Lewis al. 2024).. For AdS and AdK, the sampled conformations closely algin with the target holo structures. For TEM-1, I, we observe remaining differences in the left small helix where the sampled structure does not fully form the helix. Nevertheless, these results demonstrate promising capability of STAR-MD in predicting complex, cooperative, and long-range conformational dynam- ics.
Adenyloouccinate Synthetase (POA7D4)
B-Lactamase (P62588)
Adanylate Kinase (P69441)
Apu SINP
Apo 1ADE
1550 D.71
1o00 192223
Figure 27: STAR-MD samples Holo structures started from Apo. Structures superimposed and the local binding regions are highlighted. The Holo RMSD shown under each case is calcualted based on the local binding residues.
J.2.2 KINASE ACTIVATION
Abl kinase is signaling protein related to leukemia and other cancers. It undergoes conformational transitions between active and inactive states, which regulate its activity. These transitions involve several functional domains (P-loop, aC-helix, and the Asp-Phe-Gly DFG motifs), imposing complex millisecond-caale dynamics (Xie al. (020). We investigate whether STAR-MD can sample the active-inactive transition when initialized from one of the two states. Specifically, we generate trajectories of 500 frames with stride of 5.12 ns, corresponding to roughly 2.5 us of simulation.
To assess the states of generated conformations, we extract pairwise Ca distances and backbone dihedral angles (v. Φ) for residues around function domains involved in the transition (THR262- GLU227, VAI.287-LEU311, LYS397-GLYY417 from 20 active-state structures (PDB:6XR6) and 20 inactive-state structures (PDB:6XR7). We then apply principle component analysis to obtain
'Residue information:
assets/multiconf benchmark 0.//cryppicpoccee/tecall residinfo
47

---

Published as a conference paper at ICLR 2026
2D low-dimensional projection that separates the active and inactive conformations. We also define a core-RMi core-RMSD metric measured as the Ca-RMSD between sample and the reference structure over the aforementioned residues of interest.
As shown in Figure 28 when starting from the inactive state, STAR-MD is able to sample towards the active state with the best core-RMSD the active state 2.00 Å. When sampling from the active state, the majority of sampled conformations remain near the active-state region, while a subset of samples explores the inactive region. This behavior is in line with the underlying energetic profile, where approximately 88% of conformations correspond to the active state and only ~88 correspond to the inactive state (Xie al., 2020). The best RMSD the inactive reference (2.60 Å) is still smaller than the average pairwise RMSD between active and inactive conformations in PDB (3.30 + 0.34 Å).
Overall, this case demonstrates the potential of STAR-MD to sample complex conformational dy- namics such as active-inactive transition in kinases.
Sample from inactive state (6XR7) to active state (6XR6)
Active RMSD:
6XRS
15
inactive
Samples
Sample
-10
-10
10
RMSD: A
PC
Sample from active state (6XR6) to inactive state (6XR7)
Active RKSD: 1.89 A
Active
BXR7
Inactive RMSD: 2.60 A
inactive Samples
Sample
-10
-30 PC
10
RMSD: 2.60
Figure 28: Results on ABI Kinase sampled from the inactive state (top) and active state (bot- tom). Left: PCA projection of reference conformations from the active and inactive PDB ensembles, with STAR-MID samples shown in green. Right: Superposed structures of the best sample and the corresponding target state, with core dynamic regions highlighted.
48

---

Published as a conference paper at ICLR 2026
Miscellaneous
K LLM USAGE DISCLOSURE
We utilized Large Language Models (LLMs) as writing assistant during the preparation of this manuscript. For few sections, our process involved authoring detailed outlines and initial drafts to establish the core scientific arguments and structure. Subsequently, we used LLMs to refine the verbiage, improve grammar, and enhance the overall clarity of the text in those specific sections. scientific claims, experimental results, and theoretical assertions originated from the human authors, who take full responsibility for the final content of the paper, in accordance with ICLR 2026 policy.
49


Reviews:
[{"review_id": "mziK6lfbKh", "reviewer": "Reviewer_Qv5B", "summary": "This paper introduces STAR-MD, a spatio-temporal SE(3)-equivariant diffusion model designed for learning protein molecular dynamics directly from Cartesian coordinate trajectories. The model represents each residue as a rigid frame (translation + rotation) and predicts its time evolution conditioned on physical stride Δt, trained via diffusion-style denoising. It is evaluated on the ATLAS dataset and compared against baselines including MDGen, ConfRover, and AlphaFolding-4D. The authors report that STAR-MD reproduces MD-like short-horizon dynamics (RMSD vs lagtime, tICA correlations), achieves high backbone validity, and maintains stability over extended rollouts (up to 1 μs). Ablation studies test variants without diffusion noise and without joint spatio-temporal attention.\n\nThis is a technically solid and promising paper whose model is likely to influence future ML-based coarse-grained MD approaches. However, the evaluation would be much stronger with explicit clarification of (i) representational scale differences, (ii) alignment and Δt normalization procedures, and (iii) statistical variability across rollouts.\nThe current comparisons convincingly show stability, but not necessarily accuracy, and the role of coarse-graining deserves explicit discussion.", "questions": "1. Role of representation. The results suggest that representational scale (Cα vs. backbone vs. all-atom) may contribute more to stability than architectural details. Do the authors have evidence or experiments that could separate representational effects from architectural ones?\n2. Alignment. Were trajectories aligned (e.g., via Cα superposition) before computing RMSD and PCA/tICA metrics, and if so how?\n3. Δt comparability. How were lag times normalized when comparing models with variable vs. fixed strides?\n4. Shading in Fig. 3: What statistic do the shaded bands represent? Standard deviation, percentile range, or something else?\n4. Long-horizon ablations. Do the “w/o noise” and “w/ sep attn” variants remain stable in 240 ns – 1 µs rollouts?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. Clear technical contribution: Combines SE(3)-equivariant frame modeling and continuous-time conditioning in a diffusion framework which conceptually elegant and well-motivated.\n2. Strong empirical results: STAR-MD substantially outperforms prior coordinate-based models in conformational coverage, dynamic fidelity, and stability over long horizons.\n3. Method clarity (core model): The paper’s architectural description and use of per-residue SE(3) frames are well explained.\n4. Long-horizon evaluation: Extending rollouts to hundreds of nanoseconds or microseconds is ambitious and valuable for assessing numerical stability.\n5. Well-integrated metrics: Use of RMSD, tICA, and validity metrics ties the work to established MD analysis practices.", "weaknesses": "1. Comparative clarity. It’s difficult to realize from the text that the baselines use very different geometric representations: AlphaFolding-4D is all-atom, MDGen and ConfRover are backbone-level, and STAR-MD is Cα-only.\nBecause this difference likely affects both stability and metric scale, the paper should state it explicitly and discuss how Cα-level evaluation influences comparisons.\n\n2. Alignment and Δt normalization. The paper does not describe how trajectories are aligned prior to RMSD/PCA/tICA computation, nor how variable Δt conditioning was reconciled with baselines’ fixed strides.\n\n3. Single-trajectory visuals. Figures 2–3 show single rollouts; ensemble variability or error bars would make coverage and stability claims more convincing.\n\n4. Unexplained shading in Fig. 3. The caption omits what the shaded bands represent (presumably inter-protein variation).\n\n5. Minimal ablation scope. Only “w/o noise” and “w/ separate attention” are tested. These yield small metric changes, so it remains unclear which components truly drive the reported improvements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZuAcNZOxvc", "reviewer": "Reviewer_jt5g", "summary": "This paper introduces STAR-MD, a new approach to modeling temporal protein dynamics, focusing on long-term consistency, using generative diffusion models. To this end, the authors use a scalable residue frame representation, following previous work. However, this corresponds to a coarse-grained representation compared to the fully atomistic molecular dynamics simulations that are used to train the model. Based on this observation, it is possible to show with Mori-Zwanzig theory that an accurate modeling of the temporal dynamics in the coarse representation requires the incorporation of memory (opposed to a fully Markovian approach of some prior works) as well as complex spatial-temporal correlations. Therefore, the authors introduce an architecture that not only leverages a single previous frame but a context window of several frames, thereby capturing the system's history as required by Mori-Zwanzig theory. Moreover, they use explicit joint temporal-spatial attention layers to also capture the necessary complex correlations. That aside, the authors leverage established architecture modules, like invariant point attention and pair representations. Conceptually, the approach can be considered related to methods in the video generation literature. The authors then validate their method and outperform prior baselines. The simulated protein dynamics trajectories are scored in terms of structural quality, conformational coverage and dynamic fidelity. In particular for long roll-outs the method performs better than baselines, avoiding major drift and error accumulation.", "questions": "In section 4.4, the paper generates trajectories with different strides. How far can this be pushed? The model is trained with a stride length up to 10ns (section 3.3). What happens if we generated data with this maximum stride length and use ~850 frames/steps, as in the other experiments? What simulation time could we reach without error accumulation? Shouldn't this improve coverage and phase space exploration?\n\nThat aside, I don't have further questions, but some typos:\n- line 036: integrate -> integrates\n- line 107: embeddings -> embedding's?\n- line 124: blocks -> block\n\nAnd a suggestion: In Figure 2, the two MD reference curves are almost exactly on top of each other - I initially thought a curve is missing. I would strongly suggest the authors to address this and fix the visualization so it's clear there are two curves on top of each other.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- **Clarity**: Overall the paper is easy to follow (however, I believe the motivation for the approach could be explained better, see weaknesses below).\n- **Originality:** The paper makes a novel and original contribution in the area of protein dynamics by explicitly modeling long context windows and using joint spatio-temporal attention with conditioning noise to achieve long-term temporal generative modeling of molecular dynamics trajectories. However, the different components can be found elsewhere in the literature, e.g. in video generation, and the approach is therefore not overly creative from an architecture perspective. This is not a major concern, though, as the results seem mostly strong. On the other hand, the coarse-graining and Mori-Zwanzig formalism perspective is novel in the context of generative method design and interesting (but this could be presented better, see comments below).\n- **Experimental results:** Overall, the paper shows satisfactory experimental results, outperforming baselines. In particular the long-horizon results are strong.\n- **Significance**: Accelerating and emulating molecular dynamics simulation via generative methods is an important research topic with impact in biology, physics and chemistry. The paper makes a non-negligible step in this direction, and therefore the contribution can be considered significant.", "weaknesses": "- **Presentation of motivation:** In the main text, the authors argue without much explanation that long context and complex spatio-temporal modeling is necessary to accurately model molecular dynamics trajectories. Naively, it will be difficult to understand for most readers why that is, given that the original atomistic molecular dynamics data is essentially Markovian, corresponding to a simple integration of equations of motion (possibly with a thermostat/barostat). The reasoning for this is only explained in more depth in the appendix, when looking at the problem through the lens of coarse graining and Mori Zwanzig theory, which many readers won't know. Looking at the problem of generative modeling of molecular dynamics from this perspective is very interesting and novel, yet this is only \"hidden\" in the appendix. I would suggest the authors to restructure the paper and explain more of this in the main paper. Meanwhile, the diffusion formalism and other basics are completely standard and could be moved to the appendix. This would improve the paper, I believe.\n- **Results in Figure 2:** In Figure 2, right, the MD trajectory of STAR-MD covers more phase space than the baselines, but it still never explores the mode on the right hand side (high PC1, low PC2 values). While the method is evidently better than the baselines, it is still unclear how well STAR-MD explores the configurational state space.\n- **Configurational sampling and baselines:** More generally, long-horizon molecular dynamics, the core problem tackled in this paper, is meant to explore the full phase space. I believe the authors should study this in more depth and compare to other methods that sample molecular conformations. For example, how does STAR-MD compare to BioEmu (https://www.science.org/doi/10.1126/science.adv9817) in that regard, one of the most well-known and prominent methods to sample similar molecular systems (proteins)? Can we have full phase space diagrams for different test proteins and compare (e.g. the fast folding test proteins used in BioEmu)? The calculated coverage metrics are nice, but not very interpretable. I would expect the method to show conformational coverage similar to MD to be useful in practice -- this is a critical aspect.\n- **Free Energy Calculations:** Also related to the above, does the method enable free energy calculations? For instance of conformational transitions? Free energy calculation is one of the most important tasks of molecular dynamics. Hence, it would be good to know if STAR-MD can be used for such tasks, and it would be great to add results on this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rFiZ6LAJNB", "reviewer": "Reviewer_QE3i", "summary": "This paper introduces STAR-MD, an autoregressive SE(3) diffusion model designed to simulate long-horizon protein dynamics. The authors identify two key bottlenecks in existing generative models: 1) poor computational/memory scaling (due to $\\mathcal{O}(N^3)$ or $\\mathcal{O}(N^2)$ pair-based operations) and 2) error accumulation that leads to instability in long rollouts.\n\nSTAR-MD's methodology tackles both problems directly. For Scalability, it uses a novel joint Spatio-Temporal (S$\\times$T) attention that operates only on single-residue features, avoiding the expensive pair-feature bottleneck of competitors like ConfRover. For Stability, it employs a \"contextual noise perturbation\" (a form of diffusion forcing) that makes the model robust to its own small errors, preventing them from compounding over time.\n\nThe results are very strong. On the 100 ns ATLAS benchmark, STAR-MD achieves the best combination of structural validity and conformational coverage. More importantly, it is the first model to demonstrate stable, physically plausible generation out to the 1 microsecond timescale, a task where all baseline models (MDGen, ConfRover-W, AlphaFolding) fail catastrophically due to either memory limits or error accumulation.", "questions": "1. Can the authors please comment on the \"singles-only\" trade-off? How confident are they that this architecture can capture complex, cooperative phenomena like allostery, which are fundamentally defined by long-range pairwise correlations? Have they considered testing STAR-MD on a known allosteric benchmark protein to validate this?\n\n2. The validation of conformational coverage (JSD/Recall) is done in the PCA space of the reference trajectory. This metric is a good measure of ensemble overlap but does not validate the kinetic pathways or transition probabilities. A model could, in theory, match this distribution without being a good simulator. Could the authors comment on this choice and why a more kinetically-focused metric (like implied timescales) was not also used?\n\n3. The tICA metric details are sparse. Given the \"AlphaFolding Paradox\" (where a 0.11% valid model gets a perfect tICA score), this metric needs more justification. What featurization, lag time, and number of components were used?\n\n4. Related to the AlphaFolding baseline: the 0.11% validity is shockingly low. Can the authors confirm their setup for this baseline? It's hard to believe this is a representative result for that model, and it makes the comparison less meaningful.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The primary strength is the joint S$\\times$T attention on singles-only features. This is a very clever architectural trade-off. It avoids the $\\mathcal{O}(N^3)$ or $\\mathcal{O}(N^2 L)$ complexity of competitors (AlphaFolding, ConfRover), which is the main barrier to long-horizon simulation. The KV cache analysis (Fig. 5) proves this is a massive practical win (e.g., 6.6MB vs 1.3GB per layer).\n\n2. The \"contextual noise perturbation\" is the second key contribution. Autoregressive models are notoriously unstable in long rollouts. This method demonstrably fixes the problem. The stability plots (Fig. 3) and long-horizon results (Table 2) are direct proof that STAR-MD is robust to compounding errors.\n\n3. The results are unambiguous. On the 100ns benchmark (Table 1), STAR-MD has the best balance of coverage and validity. On the 1µs benchmark (Table 2), it's the only model that works, maintaining 80% validity while competitors are unusable (e.g., MDGen 23% validity, AlphaFolding 0.06%).\n\n4. The paper provides strong evidence that this is a trajectory simulator, not just an \"ensemble generator.\" The long-horizon stability (Fig. 3) and kinetic fidelity metrics (tICA and RMSD-vs-lag) show that the model generates a temporally coherent and physically plausible \"movie,\" not just a shuffled bag of good-looking frames.", "weaknesses": "1. The paper's greatest strength is also its biggest unproven assumption. The model trades physical explicitness (i.e., operating on pair features) for computational speed. It bets that its S$\\times$T attention is powerful enough to implicitly learn all the complex, long-range pairwise interactions (like allostery) just from single-residue features. While this seems to work for the ATLAS proteins, it is a major leap of faith that this will hold for more complex systems defined by subtle, cooperative, long-range motions.\n\n2. The \"Dyn.\" metrics are a bit superficial. The tICA score is just a single number, and the paper's own Table 1 shows that the \"AlphaFolding\" baseline (which is 99.89% invalid) can get a \"perfect\" 0.17 tICA score. This \"AlphaFolding Paradox\" proves that the tICA metric alone is an insufficient measure of a good simulator. The RMSD/autocorrelation plots are good, but the paper would be much stronger if it included a more rigorous kinetic analysis (e.g., implied timescales from an MSM, VAMP-2 scores, or transition rate calculations).\n\n3. For the long-horizon tests (Table 2), the ConfRover baseline had to be modified to a windowed variant (ConfRover-W) to manage its high memory requirements . While this strongly supports STAR-MD's superior scalability, it is not a direct comparison of kinetic fidelity, as the baseline is architecturally constrained by a limited temporal window. The AlphaFolding baseline appears to be a non-competitive comparator. With a reported structural validity of only 0.11% (Table 1), this model is fundamentally unsuited for this task . Its \"perfect\" tICA score (0.17) serves primarily to demonstrate the insufficiency of the tICA metric in isolation, rather than providing a meaningful performance benchmark.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "I8NcprPdDB", "reviewer": "Reviewer_yvwX", "summary": "This paper proposes STAR-MD, an autoregressive model for MD trajectories of protein systems. The key innovation is a causal diffusion transformer to predict the conditioning at a given frame using spatio-temporal attention in an efficient manner.  The model samples per-residue SE(3)-frames at each timestep via a diffusion process on SO(3), conditional on this conditioning vector. Experiments are performed on the ATLAS dataset, showing improved coverage of the conformational space as compared to previous methods, with significantly reduced memory requirements.", "questions": "* Can you explain the structural noise aspect in the section “Contextual Noise Perturbation for Robust Rollouts”? Is it that you sample $τ ∼U[0.0, 0.1]$ and the corresponding noisy frame $x_\\tau$ instead of keeping the clean history?\n* Is the main reason for improved efficiency due to significantly lesser memory? The authors point out $O(N^2 L^2)$ for STAR-MD vs $O(N^3 L)$ for previous methods, but in the large $L$ limit, I would expect STAR-MD to become more inefficient?\n* Please mention how the ATLAS test set was created (even if referring to prior work) since it helps put generalization results into context. Do you find that the model is able to generalize to new proteins with low sequence similarity?\n* Please mention the inference costs of all models to enable a comparison to the ground truth MD.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The model removes the Pairformer block present in existing protein models to save on memory. Instead, the model utilizes attention over the single-residue features across all frames and all residues. By clever training tricks such as feeding the clean and noisy frames together with a special attention mask, the training is parallelized over all frames, improving efficiency.\nThe experiments on long-horizon trajectory generation are good to see, and the model performs well in this challenging setting.", "weaknesses": "The empirical results in this paper are indeed impressive, but I would have liked to see a discussion and evaluation in identifying functional motions in proteins (eg. transitioning between active/inactive states, for example as represented by DFG in/out motifs in kinases. See https://www.rcsb.org/structure/6XR6 vs https://www.rcsb.org/structure/6XR7 for an example with Abl kinase.) Indeed, that is the main point of performing long MD simulations. ATLAS itself has very few functional motions, and tICA itself is not an ideal metric to extract these motions. That being said, I think these comments are more about the problem (trajectory generation) rather than this specific approach.\n\nFor the model itself, I would have liked to see generalization experiments to unseen lengths or sequences with low similarity to the training set, since that is a very common setting in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)