Generative Modeling of Molecular Dynamics Trajectories
Authors: Bowen Jing
Hannes Stark
Tommi Jaakkola
Bonnie Berger
Venue: NeurIPS 2024 poster (2024)
Keywords: molecular dynamics
molecular simulation
transition paths
Boltzmann distribution
proteins
Primary area: machine_learning_for_physical_sciences

Abstract: Molecular dynamics (MD) is a powerful technique for studying microscopic phenomena, but its computational cost has driven significant interest in the development of deep learning-based surrogate models. We introduce generative modeling of molecular trajectories as a paradigm for learning flexible multi-task surrogate models of MD from data. By conditioning on appropriately chosen frames of the trajectory, we show such generative models can be adapted to diverse tasks such as forward simulation, transition path sampling, and trajectory upsampling. By alternatively conditioning on part of the molecular system and inpainting the rest, we also demonstrate the first steps towards dynamics-conditioned molecular design. We validate the full set of these capabilities on tetrapeptide simulations and show preliminary results on scaling to protein monomers. Altogether, our work illustrates how generative modeling can unlock value from MD data towards diverse downstream tasks that are not straightforward to address with existing methods or even MD itself. Code is available at https://github.com/bjing2016/mdgen.

Generative Modeling of Molecular Dynamics Trajectories
Bowen Jing"
Hannes Stärk" Tommi Jaakkola'
Bonnie Berger'
'CSAIL, Massachusetts Institute of Technology
Dept. of Mathematics, Massachusetts Institute of Technology {bjing, hstark) Cairttmmmmt...mmm edu, Commi@esail .dit.edu, babêmit.
Abstract
Molecular dynamics (MD) a powerful technique for studying microscopic phe- nomena, but its computational cost has driven significant interest in the development deep learning-basc based surrogate models. We introduce generative modeling of molecular trajectories
paradigm for learning flexible multi-task surrogate models of MD from data. By conditioning on appropriately chosen frames of the trajectory, we show such generative models can be adapted to diverse tasks such as forward simulation, transition path sampling, and trajectory upsampling. By alternatively conditioning we also demonstrate the first steps towards dynamiis-conditioned molecular design. We validate the full set of these capabilities preliminary results on scaling to protein monomers. Altogether, our work illus- trates how generative modeling can unlock value from MD data towards diverse downstream tasks that are not straightforward to address with existing methods or even MD itself. Code is available at https //gi /gggttu...... com/ /njjjji ing2016 mndgen.
part of the molecular system and inpainting the rest,
tetrapeptide simulations and show
Introduction
Numerical integration of Newton's equations of motion atomic scales, known molecular dynamics (MD), is a widely-used technique for studying diverse molecular phenomena in chemistry, biology, and other molecular sciences (Alder and Wainwright, 1959; Rahman, 1964; Verlet, 1967; McCammon et al., 1977). While general and versatile, MD is computationally demanding due to the large separation in timescales between integration steps and relevant molecular phenomena. Thus, vast body of literature spanning several decades aims to accelerate or enhance the sampling efficiency of MD simulation algorithms (Ryckaert et al., 1977: Darden et al., 1993; Sugita and Okamoto, 1999; Laio and Parrinello, 2002; Anderson et al., 2008; Shaw et al., 2009). More recently, learning surrogate models of MD has become an active area of research in deep generative modeling (Noé al., 2019; Zheng et al., 2023; Klein et al., 2024; Schreiner et al., 2024; Jing et al., 2024). However, existing training paradigms fail to fully leverage the rich dynamical information in MD training data, restricting their applicability to a limited set of downstream problems.
In this work, we propose MDGEN, novel paradigm for fast, purppse surrogate modeling of MD based on direct generative modeling of simulated trajectories. Different from previous works, which learn the autoregressive transition density or equilibrium distribution of MD, we formulate end-to-end generative modeling of full trajectories viewed as time-series of 3D molecular structures. Akin to how image generative models were extended to videos (Ho et al., 2022), our framing of the problem augments single-structure generative models with additional time dimension, opening the door to a larger set of forward and inverse problems to which our model can be applied. When
Equal contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Smndetio
Drwand
?
?
?
Trajectory coordinates
SF(3)-invarian roto trann-rraiaioonn offeet and forskon aiglle lokenication
?
?
?
Upanpiru
?
?
Framoo
Figure 1: (Left) Tasks: generative modeling of MD trajectories addresses several tasks by conditioning on different parts of trajectory. (Right) Method: We tokenize trajectories of T frames and L residues into an (T x L) -aaray of SE(3)-invariant tokens encoding roto-translation offsets from key frames and torsion angles. Using stochastic interpolants, generate arrays of such tokens from Gaussian noise.
provided (and conditioned on) the initial "frame" of given system, such generative models familiar surrogate forward simulators of the reference dynamics. However, by providing other kinds of conditioning, these "molecular video" generative models also enable highly flexible applications to variety of inverse problems not possible with existing surrogate models. In sum, we formulate and showcase the following novel capabilities of MDGEN:
Forward simulation given the initial frame of trajectory, sample potential time evolution of the molecular system.
Interpolation-given the frames at the two endpoints of a trajectory, we sample a plausible path connecting the two. chemistry, this known transition path sampling and is important for studying reactions and conformational transitions.
Upsampling-gii given trajectory with timestep At between frames, we upsample the "fram- erate" by a factor of M to obtain a trajectory with timestep A//M. This infers fast motions from trajectories saved at less frequent intervals.
Inpainting given part of a molecule and its trajectory, generate the rest of the molecule (and its time evolution) to be consistent with the known part of the trajectory. This ability could be applied design molecules scaffold desired dynamics.
These tasks are conceptually illustrated in Figure 1. the typical modeling paradigm of approximating the data-generating process, the others represent novel capabilities
While the forward simulation task aligns with
scientifically important inverse problems not straightforward address even with MD itself. As such, our framework presents a new perspective on how to unlock value from MD simulation with machine learning towards diverse downstream objectives. We highlight further exciting possibilities opened up by our framework in Section 5.
We demonstrate our framework on MD simulations of tetrapeptides (i.e., length-4 peptides), with pre- liminary extensions full-sized protein monomers. do so, we parameterize molecular trajectories in terms of sidechain torsions and residue offsets with respect to conditioning key frames, obtaining generative modeling task over 2D array of (3))-inraiii tokens rather than residue frames or point clouds. In this parameterization, we can then employ a Scalable Interpolant Transformer (SiT) (Ma et al., 2024) as our flow-baser generative backbone, avoiding the more restrictive geometric architectures commonly used for molecular structure. Furthermore, by replacing the time-wise atten- tion SiT with the long-context architecture Hyena (Poli et al., 2023), we provide proof-of-concee of scaling up to trajectories of 100k frames, enabling wide range of timescales and dynamical processes to be captured with a single model generation.
We evaluate MDGEN on the forward simulation, interpolation, upsampling, and inpainting tasks on tetrapeptides in a transferable setting (i.e., unseen test peptides). Our model accurately reproduces free energy surfaces and dynamical content such as torsional relaxation and Markov state fluxes, provides realistic transition paths between arbitrary pairs of metastable states, and recovers fast dynamical phenomena below the sampling threshold of coarse-timestep trajectories. In preliminary

---

steps toward much higher sequence recovery than inverse folding methods based on one or two static frames. Finally, we evaluate MDGEN on simulation of proteins and find that it outperforms MSA subsampling with AlphaFold (Del Alamo al., 2022) in terms of recovering ensemble statistical properties.
design, we show that molecular inpainting with MDGEN obtains
2 Background
Molecular dynamics, At a high level, the aim of molecular dynamies is the integrate the equations of motion M, = xxx Ux1 ......,,, (N) ) for each particle i in molecular configuration (x1 xN) E R3N where M, the mass and is the potential energy function (or force field) U R + IR.
However, these equations of motion are often modified to include a thermostat in order to model contact with surroundings at a given temperature. For example, the widely- uueed Langevin thermostar transforms the equations of motion into a stochastic diffusion process:
dx; - Pa/M, dt, dp: = Vx; U - >p; + 22M dw
where P, are the momenta. By design. this process converges to the Boltzmann distribution of the system p(xi XN) x -U/// To incorporate interactions with solvent colecules-ubiquiin biochemistry-one includes a box of surrounding solvent molecules as part of the molecular system (explicit solvent) or modifies the force field U to model their effects (implicit solvent). In either case, only the positions x; of non-solvent atoms are of interest, and their time evolution constitutes (for our purposes) the MD trajectory.
in
Deep learning for MD. An emerging body of work seeks to approximate the distributions over configurations X (X1 ... XN) arising from Equation with deep generative models. Fu al. (2023), Timewarp (Klein et al., 2024), and ITO (Schreiner et al., 2024) learn the transition density p(X++t: X,) and emulate MD trajectories via simulation rollouts of the learned model. On the other hand, Boltzmann generators (Noé et al., 2019; Köhler et al., 2021; Garcia Satorras et al., 2021; Midgley al., 2022, 2024) directly approximate the stationary Boltzmann distribution, forgoing any explicit modeling of dynamics. In particular, frames from MD trajectories have demonstrated promising scalability and generalization to protein systems (Zheng et al., 2023; Jing et al., 2024). However, these works have focused exclusively on forward simulation and have not explored joint modeling of entire trajectories (X,... X1+NA() or the inverse problems accessible under such formulation.
diffusion models trained with
Stochastic interpolants. We build MD trajectory generative model under the stochastic in- terpolants framework: Given a continuous distribution p1 = Pdata over R", stochastic interpolants (Albergo and Vanden-Eijnden, 2022; Albergo et al., 2023; Lipman et al., 2022: Liu et al., 2022), provide a method for learning continuous flow- models dx = (x, L) di transporting prior distribution po (e.g., Po = (0, I)) to the data pl. To do so, one defines intermediate distributions X₁ Pt, t (0, 1) via X₃ = a,X1 + oxXo where X. and X1 P1 and the interpolation path satisfies ag 0i 0 and 01 do
A neural network Uo R" [0,1] - R" is trained to
approximate the time-evolving flow field
(x 1) R u(xe, = Ex..xI x. (iEX1 + oxx0]
(2)
which satisfies the transport equation 0p: 0t + V (prer) O. Hence, convergence, noisy samples Xo ~ Po can be evolved under Up to obtain data samples x, ~ Pt. When parameterized with transformers (Vaswani et al., 2017), stochastic interpolants are state-of-the-art in image generation (Esser et al., 2024). In particular, we adopt the notation, architecture, and training framework of Scalable Interpolant Transformer (SiT) (Ma et al., 2024), to which refer for further exposition.
3 Method
3.1
Tokenizing Peptide Trajectories
Given a chemical specification of a molecular system with atoms, our aim to learn a generative model over time-series x [X,,... ..) of corresponding molecular structures X, E R3N for some trajectory length T. In this work, we specialize to MD trajectories of short peptides (Sections 4. 4.4) or single-chain proteins (4.4). Thus, our chemical specifications are amino acid sequences
3

---

A (1 20 and we adopt an 2021; Yim et al., 2023). In this parameterization, the all-atom coordinates of each amino acid residue are implicitly described by roto-translation (i.e., element of (F)) corresponding to the rigid body motion of the residue, and seven torsion angles describing conformation:
parameterization of peptide structures (Jumper et al.,
x = 9g. .........T TP.. g E ( x ([SE(3) xT1")²
Throughout, subscripts indicate time and superscripts residue indices. The undefined torsion angles can be randomized and are unsupervised for residues with fewer than seven torsion angles.
Traditionally, equivariant architectures have been required for geometry-awars processing of polypep- tide structures. However, to learn a scalable generative model over this space of roto-translations and torsion angles, we seek to represent each X in terms of an (3ii-naariia feature vector-a a loken suitable for processing by vanilla transformers. To obtain such a vector, we leverage the fact that we are concerned with conditional trajectory coneration-meaninn that there always exists at least one frame in the trajectory with un-noised roto-translationn which we do not need to generate and can reference in the modeling process. Inspired by analogy video compression, we refer to such frames as key frames. We can then obtain (E(3)-invarinl tokens by parameterizing the roto-translations of remaining structures relative to the key frames.
In more detail, given K key frames at times t,.. we tokenize residue j in frame as:
xd =
[ [66
.......... ..... ((....((o))n))))) (I 96)
C 7k++1
(4)
where of E SE(3) represents the roto-translation and [T'']s the torsion angles of residue j at frame Here, Φ (3) + RR? parameterizes an element of SE(3) in terms of a unit quaternion and translation vector, and 3o T + I²2 converts torsion angles to points on the unit circle. We thus obtain (7K + 14)-dimensional array for each residue in every frame. Because the relative roto-translations and torsion angles are both (3))-iinvarial in this manner we can represent a polypeptide molecular trajectory as an (T L) -array of S -invariant tokens.
To untokenize a generated trajectory of tokens to all-atom coordinates X; E R3N we first convert each predicted quaternion and translation vector relative roto- -ro-strannlttion and apply it to the key frame(s), obtaining absolute roto-translations. We then read off the torsion angles from the unit circle and assemble the all-atom coordinates as implemented in Jumper et al. (2021), averaging the reconstructions from different key frames if needed.
3.2 Flow Model Architecture
Our base modeling task is to generate distribution over RTXLX(7K+14) conditioned on roto- translations of one or more key frames Sr, ... ..gtt:: and (in most settings) amino acid identities A. To do so, we learn flow-based model via the stochastic interpolant framework described in SiT (Ma
et al., 2024) and parameterize velocity network vo(- 9t, .t. Itx A) (TXxx(7K+11) x [0.
[TXXX(7K+14) 14) To condition on the key frames and amino acids, we first provide the sequence embedding to several IPA layers (Jumper al., 2021) that embed the key frame roto-raanslation;; these conditioning representations (which are S added to the input embeddings. The main trunk of the network consists of alternating attention blocks across the residue index and across time, with the construction of each block closely resembling DiT (Peebles and Xie, 2023). Sidechain torsions and roto- tro-trannslatinn offsets, when available, are directly provided to the model as conditioning tokens. Further details are provided in Appendix A.
are broadcast across the time axis and
In the molecular inpainting setting where we also generate the amino acid identities, we additionally require a generative framework over these discrete variables. While several formulations of discrete diffusion or now-matching are available (Hoogeboom et al., 2021; Austin et al., 2021; Campbell et al., 2022, 2024), we select Dirichlet flow matching (Stark et al., 2024) as it is most compatible with the continuous-space, continuous-time stochastic interpolant framework used for the positions. Specifically, we place the amino acid identities on the 20-dimensional probability simplex (one per amino acid), augment the token representations with these variables, and regress against a T x L x (7K 14 + -immensiiol vector field. Further details are provided in Appendix A.2.

---

3.3 Conditional Generation
We present the precise specifications of the various conditional generation settings in Table 1. Depending on the task, we choose the key frames to be the first frame 91 or the first and last frames 91. 9r. Each conditional generation task is further characterized by providing the ground -tutt tokens of known frames or residues as additional inputs to the velocity network. Meanwhile, mask tokens are provided for the unknown frames and residues that the model generates. For example, in the upsampling setting, we provide ground-truth tokens every M frames, while mask tokens are provided for all other frames. We note that in the inpainting setting, the model accesses the toottranslations of designed residues at the trajectory endpoints via the key frames, constituting slight departure from the full inpainting setting. However, these residues are not observed for intermediate frames, and their identities are never provided to the model.
Table 1: Conditional generation settings, g: roto-trans tro-trranllaiinns т: torsions. A: residue identities M: upsampling factor. Superscripts indicate residue index and subscripts indicate frame (time) index. For inpainting, we find that excluding identities and torsions reduces overfitting
Setting
Key frames
Generate
Conditioned on
Token dim.
Forward simulation
91
91...T TI..T
91. T1, 91,7; T1,T, A 91+(12, JM: T T (1,2.1.. YM, A 91..-
21
Interpolation
91.9T
TI...T
28
Upsampling Inpainting
91
91.,T; TI.-T
21
91, 9T
g1.,,; A
(+20)
Experiments
We evaluate MDGEN its ability to learn from MD simulations of training molecules and then sample trajectories for unseen molecules. We focus on tetrapeptides as our main molecule class for evaluation as they provide nontrivial chemical diversity while remaining small enough to tractably simulate to equilibrium (Klein et al., 2024). Sections 4.1-4.3 thoroughly evaluate our model on forward simulations, interpolation transition path sampling, and trajectory upsampling on test peptides. Section 4.4 provides proof-of-concep and preliminary exploration of additional tasks- namely, inpainting for cenamics-conditions design, long trajectories with Hyena (Poli et al.. 2023), and scaling to simulations of protein monomers. Separate models are trained for each setting.
To obtain tetrapeptide MD trajectories for training and evaluation, we implicit- and explicit- solvent, all-atom simulations of =x300 training, 100 validation, and 100 test tetrapeptides for 100 ns. For proteins, we use explicit-solvent, all-atom simulations from the ATLAS dataset (Vander Meersche et al., 2024), which provides three 100 ns trajectories for each of 1390 structurally diverse proteins. Unless otherwise specified, models are trained with trajectory timesteps of At 10 ps. Our default baselines consist of replicate MD simulations ranging from 10 ps to 100 ns, with additional comparisons in each section as appropriate.
Our experiments make extensive use of Markov State Models (MSMs), a widely used coarse-grained representation of molecular dynamics (Prinz et al., 2011: Noé et al., 2013). We obtain an MSM to represent a system by discretizing its MD trajectory (parameterized with torsion angles) into 10 metastable states and estimating the transition probabilities between them. Appendix B provides further details on constructing MSMs and other experimental settings. Additional results, including structural validations and further comparisons with related methods, can be found in Appendix C.
4.1 Forward Simulation
In the forward simulation setting, we train a model to sample 10 trajectories conditioned on the first frame. By chaining together successive model rollouts at inference time, we obtain 100 ns trajectories for each peptide to compare with ground-truth simulations. We evaluate if these sampled trajectories (1) match the structural distribution of trajectories from MD, (2) accurately capture the dynamical content of MD, and (3) traverse the state space in less wall-clock time than MD.
Distributional similarity. We report the Jemsen-Shannon divergence (JSD) between the ground-truth and emulated trajectories along various collective variables shown in Figure 2 and Table 2. The first
5

---

A
Backbone torsions
B
MDD FES
C
Sample
MSM state occupancy
TIC-O speedup
n=
r=0.89
ALDA
our
GLIL
10
100
1000
Torsional relaxation (ns)
MD
10
r=
Oues
Ours
FLRH
0.55
MD
Figure 2: Forward simulation evaluations on test peptides. (A) Torsion angle distributions for the six backbone torsion angles from MD trajectories (orange) and sampled trajectories (blue). (B, C) Free energy surfaces along the top two TICA components computed from backbone and sidechain torsion angles. (D) Markov State Model occupancies computed from MD trajectories versus sampled trajectories, pooled across all test peptides (n = 1000 states total). (E) Wall-clock decorrelation times of the first TICA component under MD versus our model rollouts. (F) Relaxation times of torsion angles computed from MD versus sampled trajectories, pooled across all test peptides-500 backbone (blue) and 722 sidechain (orange) torsions in total.
tetrapeptide AAAA colored by the decorrelation time computed from MD (top) and from rollout trajectories (bottom).
(G) Torsion angles
set of these are the individual torsion angles (backbone and sidechains) in each tetrapeptide. The second of variables are the top indepen- dent components obtained from time-lagged independent components analysis (TICA), rep- resenting the slowest dynamic modes of the peptide. By each of these collective variables, MDGEN demonstrates excellent distributional similarity to the ground truth, approaching the accuracy of replicate 100-ns simulations. To more stringently assess the ability to locate and populate modes in the joint distribution over state space, we build Markov State Models (MSMs) for each test peptide using the MD trajectory, extract the corresponding metastable states, and compare the ground truth and emulated distributions over metastable states. Our model captures the relative ranking of states reasonably well and rarely misses important states or places high mass on rare states (Figure 2D).
Table 2: JSD between sampled and ground-truth distributions, with replicate simulations baselines. 100 ns represents oracle performance.
C.V.
Ours
10 ns
ns
100 ps
100 ns
Torsions (bb)
.130
.145
.212
.311
.103
Torsions (sc)
.093
.111
.261
.403
.055
Torsions (all)
.109
.125
.240
.364
.076
TICA-0
.230
.323
.432
.477
.201
TICA-0.1 joint
.316
.424
.568
.643
.268
MSM states
.235
.363
.493
.527
.208
Runtime
60s
1067s
107s
11s
3h
Dynamical content. We compute the dynamical properties of each tetrapeptide in terms of the decorrelation time of each torsion angle from the MD simulation and from our sampled trajectory. Intuitively, this assesses if our model can discriminate between slow- and fast-relaxing torsional barriers. The correlation between true and predicted relaxation timescales is plotted in Figure 2F, showing excellent agreement for sidechain torsions and reasonable agreement for backbones. To assess coarser but digher-dimenssionll dynamical content, we compute the flux matrix between all pairs of distinct metastable states using ground-truth truth and sampled trajectories and find substantial Spearman correlation between their entries (mean p = 0.67 = 0.01; Figure 8). Thus, our simulation rollouts can accurately identify high-flux transitions in the peptide conformational landscape.
Sampling speed. Averaged across test peptides, our model samples 100 ns-equivalent trajectories in ~6 GPU-seconds, compared to s33 GPU-hours for MD. To quantify the speedup more rigorously,

---

State
State 2
State 3
State 4
State 5
State 6
MDGen path in FES
Method
1
100ns MD
5
50ns MD
0.6
30ns MD
20ns MD
6
0.4
0.004
5ns MD
MD
0.002
0.0
0.000
JSD
Valid path
Average path probability
rate
Figure 3: Transition path sampling results. (Top) Intermediate states one the 1-nanosecond interpolated trajectories between two metastable states for the test peptide IPGD. (Bottom Left) The corresponding trajectory on the 2D free energy surface of the top two TICA components (more examples in Figure 9). (Bottom Right) Statistics averaged over 100 test peptides and 1000 paths for each of them. Shown are JSD, fraction of drawn paths that are valid transition paths, and average path likelihood of our discretized transitions under the reference MSM compared to discrete transitions drawn from the reference MSM or alternative MSMs built from replica simulations of varying lengths.
we compute the decorrelation wall-clock times along the slowest independent component from TICA, capturing how quickly the simulation traverses the highest barriers in state space. These times are plotted in Figure 2E, showing that our model achieves a speedup of 10x-1000x over the MD simulation for 78 out of 100 peptides (the other 22 peptides did not fully decorrelate).
4.2 Interpolation
In the interpolation or transition path sampling setting, we train model to sample 1 ns trajectories conditioned on the first and last frames. For evaluation, we identify the most well-separated states (i.e., with the least flux between them) for each test peptide and sample an ensemble of 1000 transition paths between them. Figure shows an example of such sampled path, which passes through several intermediate states on the free energy surface to connect the two endpoints.
To evaluate the accuracy of these sampled transitions, we cannot directly compare with MD jectories since, in most cases, there are zero or very few 1-ns transitions between the two selected states (by design, the transition is metastable states and evaluate the path likelihood under the transition path distribution from the reference MSM (details in Appendix B.3). We also report the fraction of valid paths (i.c., non-zero probability) and the JSD between the distribution of visited states from our path distribution versus the transition path distribution of the reference MSM. For baselines, we sample transition paths from MSMs constructed from replicate MD simulations of varying lengths and compute the same metrics for these (discrete) path ensembles under the reference MSM.
MSM event). Thus, instead discretize the trajectory
As shown in Figure 3, our paths have higher likelihoods than those sampled from any replicate MD MSM shorter than 100ns, which is the length of the reference MD simulation itself. Moreover, MDGEN's ensembles have the best JSDs to the distribution visited states the reference MD MSM and the highest fraction on valid non-zero probability paths. Hence, our model enables zero-shot sampling of trajectories corresponding to arbitrary rare transitions for unseen peptides.
4.3 Upsampling
Molecular dynamics trajectories are often saved at relatively long time intervals (10s-100s of picoseconds) to reduce disk storage; however, some molecular motions occur at faster timescales and
7

---

Figure 4: Recovery of fast dynamics via trajectory upsampling for peptide GTLM. (Left) Autocor- relations of each torsion angle from (-) the original 100 fs-timestep trajectory, the subsampled 10 ns-timestep trajectory, and (...) the reconstructed 100 fs-timestep trajectory (all length 100 ns). (Right) Dynamical content as function of timescale from the upsampled vs. ground truth trajectories, Dynamical content
Autocorrelation
Backbone
Sidechain
Upsampled
Ground truth
1.0
0.4
10 0.8
10
10 100 1000
10 ps
100 1000
stacked for all torsion angles (same color scheme). The subsampled trajectory contains only the shaded region and our model recovers the unshaded region. Further examples in Figure 10.
MD
Ours
1.00
Table 3: Sequence recovery for the in- ner two peptides when conditioning on the partial trajectory (MDGEN), the two terminal frames (DynMPNN), or single frame (S-MPNN).
0.75
0.50
0.25
Method
High Flux
Random Path
0.00
0.1
10
100
1000
1
10
100
1000
ps
ps
MDGen DynMPNN
52.1 % 17.4%
62.0%
Figure 5: Autocorrelation functions of MDGEN sidechain tor- sion angles computed from 10-ns MD trajectory (left) sus a single 100k-frame model sample with Hyena (right), capturing dynamics spanning four orders of magnitude.
24.5%
S-MPNN
16.3%
13.5%
would be missed by downstream analysis of the saved trajectory. In the upsampling setting, we train MDGEN to upsample trajectories saved with timestep 10 ps to a finer timestep of 100 fs, representing 100x upsampling factor. To evaluate if the upsampled trajectories accurately capture the fastest dynamics, we compute the autocorrelation function (6: (tt each torsion angle in the test peptides function of lag time Al ranging from 100 100 ps.
Representative examples of ground truth, subsampled, and reconstructed autocorrelation functions for two test peptides are shown in Figure (further examples in Figure 10). We further compute the dynamical content as the negative derivative of the autocorrelation with respect to log-timescale, which captures the extent of dynamic relaxations occurring at that timescale (Shaw al., 2009). These visualizations highlight the significant dynamical information absent from the subsampled trajectory and which are accurately recovered by model. In particular, model distinctly recovers the oscillations of certain torsion angles as seen in the non-monone monotonottn the autocorrelation function at uub-picosecont timescales; these features completely missed the original sampling frequency.
4.4 Additional Tasks
Inpainting Design. We aim to sample trajectories conditioned on the dynamics of the two flanking residues of the tetrapeptide; in particular, the model determines the identities and dynamics of the two inner residues. We focus on dynamics seaffolding as one possible higher-level objective of inpainting: given the conformational transition of the observed residues, we hope to design peptides that support flux between the corresponding Markov states. Thus, for each test peptide, we select 100-ps transition between the two most well-connected Markov states, mask out the inner residue identities and dynamics, and inpaint them with model. To evaluate the designs, we compute the fraction of

---

high-flux path or (as sanity check) random path Table 4: Median results test protein ensembles
MD
Ours
(n = 82); evaluations from Jing et al. (2024). Run- times are reported per sample structure or frame.
MDGEN AlphaFlow MSA sub.
Pairwise RMSD
0.48
0.48
0.22
Global RMSF r
0.50
0.60
0.29
Per-target RMSF 1 0.71
0.85
0.55
Root mean W22 dist. +
2.69
2.61
3.62
MD PCA W22 dist. +
1.89
1.52
1.88
% PC-sim > 0.5 f
10
44
21
Weak contacts J
0.51
0.62
0.40
Exposed residue J t
0.29
0.41
0.27
Figure 6: MD vs generated ensembles for 6uof_A A. with Ca RMSFs plotted by residue index (Pearson 0.74).
Runtime
0.2
70
4
generated residue types that are identical to the tetrapeptide in which the target transition is known to occur. We compare MDGEN with a bespoke inverse folding baseline that is provided the two terminal states (i.e., two fully observed MD frames), and thus designs peptides that support the two modes (rather than additionally a partially-observed transition between them). We call this baseline DYNMPNN, and it otherwise has the same architecture and settings MDGEN (more details in Appendix B.3). We find (Table 3) that MDGEN recovers the ground-truth peptide substantially more often than DynMPNN when conditioned from the reference simulatiom.
Scaling to Long Trajectories. Although Section 4.1 showed that our model can emulate long trajectories, this was limited to rollouts 1000 frames at a time with coarse 10 ps timesteps. potentially missing faster dynamics or disrupting slower dynamics. Thus, we investigate generating extremely long consistent trajectories that capture timescales spanning several orders of magnitude within single model sample. To do so, we replace the time attention in our baseline SiT architecture with a non-causal Hyena operator (Poli et al., 2023), which has O log N) rather than O(N²) overhead. We overfit on 100k-frame, 10-ns trajectories of the pentapeptide MDGEN and compare the torsional autocorrelation functions computed from a single generated trajectory with a single ground truth trajectory (Figure 5). Although not yet comparable to the main set of forward simulation experiments due to data availability and architectural expressivity reasons, these results demonstrate proof-of-concept for longer context lengths in future work.
Protein Simulation. To demonstrate the applicability of our method for larger systems such as proteins, we train model to emulate all-atom simulations of proteins from the ATLAS dataset (Vander Meersche et al., 2024) conditioned on the first frame (i.e., forward simulation). We follow the same splits as Jing et al. (2024). Due to the much larger number of residues, we generate samples with 250 frames and 400 ps timestep, such that a single sample emulates the 100 ns ATLAS reference trajectory. The difficulty of running fully equilibrated trajectories for proteins prevents the construction of Markov state models used in our main evaluations. Instead, we compare statistical properties of forward simulation ensembles following Jing et al. (2024). Our ensembles successfully emulate the ground- -tuth ensembles at a level of accuracy between AlphaFlow and MSA subsampling while being orders of magnitude faster per generated structure than either (Table 4; Figure 6).
5 Discussion
Limitations. Our experiments have validated the model and architecture for peptide simulations; however, a few limitations provide opportunities for future improvement. Due the reliance on key frames, the model is not capable of unconditional generation or inpainting of residue roto- translations. The weaker performance on protein monomers relative to peptides suggests that scaling to larger systems will likely require additional data or methodological innovations. Fine-tuning of single structure models for co-generation of the key frames and trajectory tokens, similar to the content-frame decomposition of video diffusion models (Yu et al., 2024), may provide improvement. Since our tokenization scheme is specific to polypeptides, alternative strategies will be needed to model all-atom trajectories of more general systems, such as organic ligands, materials, explicit

---

solvent. More ambitious applications (see below) may require the ability to model trajectories not of a predefined set of atoms but over a region of space in which atoms may enter and exit. As such, we anticipate advancements in tokenization and architecture to be a fruitful direction of future work.
Opportunities. Similar to the foundational role of video generative models for understanding the macroscopic world (Yang et al., 2024), MD trajectory generation could serve as a multitask, unifying paradigm for deep learning over the microscopic world. Interpolation can be more broadly framed as hypothesis generation for mechanisms of arbitrary molecular phenomena, especially when only partial information about the end states is supplied. Molecular inpainting could be a general technique to design molecular machinery by scaffolding more fine-ge ame-grained and complex dynamics, for example, redesigning proteins to enhance rare transitions observed only once in simulation or (with ab initio trajectories) de novo design of enzymatic mechanisms and motifs. Other types of conditioning not explored this work may lead to further applications, such as conditioning over textual or experimental descriptors of the trajectory. Future availability of significantly more ground truth MD trajectory data for diverse chemical systems could a chief enabler of such work. Lastly, considerations unique to molecular trajectories, such equilibrium vs non-equilibrium processes, Markovianity, and the reversibility of the microscopic world contrasted with the macroscopic world (e.g., the missing arrow of time), could provide ripe areas for theoretical exploration.
Acknowledgments and Disclosure of Funding
We thank Felix Faltings, Jason Yim, Mateo Reveiz, Gabriele Corso, and anonymous NeurIPS reviewers for helpful feedback and discussions.
This work was supported by the National Institute of General Medical Sciences of the National Institutes of Health under award number R235GM141861: the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Department of Energy Computational Science Graduate Fellowship under Award Number DESC0022158; the National Science Foundation under Grant No. 1918839; the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium; the Abdul Latif Jameel Clinic for Machine Learning in Health; the DTRA Discovery of Medical Countermeasures Against New and Emerging (DOMANE) threats program; and the DARPA Accelerated Molecular Discovery program. This research used resources of the National Energy Research Scientific Computing Center (NERSC), a Department of Energy Office of Science User Facility using NERSC awards ASCR-ERCAP0027302 and ASCRERCAP00788 18
References
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijnden. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint trXii:2303.08797 2023.
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2022.
Berni J Alder and Thomas Everett Wainwright. Studies in molecular dynamics. i. general method.
The Journal of Chemical Physics, (1(2):459 466, 1959.
Joshua A Anderson, Chris D Lorenz, and Alex Travesset. General purpose molecular dynamics simulations fully implemented on graphics processing units. Journal of computational physics, 227(10):5342-5335 2008
Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 14:17981-17993, 2021.
Andrew Camphell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnad Doucet. A continuous time framework for discrete denoising models. Advances Neural Information Processing Systems, 35:28266-28279, 2022.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.049997 2024.
10

---

John D Chodera and Frank Noé. Markov state models of biomolecular conformational dynamics. Current opinion in structural biology, 25:135-144, 2014.
Tom Darden, Darrin York, and Lee Pedersen. Particle mesh ewald: An n log (n) method for ewald sums in large systems. The Journal of chemical physics, 98(12): 100090000922 1993.
Diego Del Alamo, Davide Sala, Hassane S Mchaourab, and Jens Meiler. Sampling alternative conformational states of transporters and receptors with alphafold2. Elife, 11:c75751, 2022.
Peter Eastman, Jason Swails, John D Chodera, Robert T McGibbon, Yutong Zhao, Kyle Beauchamp. Lee-Ping Wang, Andrew C Simmonett, Matthew P Harrigan, Chaya D Stern. et al. Openmm 7: Rapid development high performance algorithms for molecular dynamies. PLoS computational biology, 13(7)e 1005659, 2017.
Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. arXiv preprint 1rXiv:2403.03206 2024.
Xiang Fu, Tian Xie, Nathan J Rebello, Bradley Olsen, and Tommi S Jaakkola. Simulate time- integrated coarse-grained molecular dynamics with multi-scale graph networks. Transactions on Machine Learning Research, 2023.
Victor Garcia Satorras, Emiel Hoogeboom, Fabian Fuchs, Ingmar Posner, and Max Welling. E (n) equivariant normalizing flows. Advances in Neural Information Processing Systems, 34:4181-4192.
2021.
Jonathan Ho. Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633-8646, 2022.
Emiel Hoogeboom, Alexey A Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. arXiv preprint arXiv:2110.02037, 2021.
Brooke E Husic and Vijay S Pande. Markov state models: From an a science. Journal of the American Chemical Society, 140(7):236-2396, 2018.
Bowen Jing, Bonnie Berger, and Tommi Jaakkola. Alphafold meets flow matching for generating protein ensembles. arXiv preprint arXiv:2402.04845, 2024.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Židek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. Nature, 596(78733
2021.
Leon Klein, Andrew Foong, Tor Fjelde, Bruno Mlodozeniec, Mare 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.
Jonas Köhler, Andreas Krämer, and Frank Noé. Smooth normalizing flows. Advances in Neural Information Processing Systems,
2021.
Alessandro Laio and Michele Parrinello. Escaping free-energy minima. Proceedings of the national academy of sciences, 99(20): 12562-22256222 2002.
Yaron Lipman, Ricky TQ Chen, Heli Ben- Gen-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arX:v:2210.02747, 2022.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning generate and transfer data with rectified flow. arXiv preprint arXiv:2209.0s8 00000,, 2022.
Nanye Ma, Mark Goldstein, Michael S Albergo, Nicholas M Boffi, Eric Vanden Fijnde, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. arXiv preprint MXXi::24010087440 2024.
J Andrew McCammon, Bruce Gelin, and Martin Karplus. Dynamics of folded proteins. nature, 267(5612):585-590 1977.

---

Laurence Midgley, Vincent Stimper, Javier Antorán, Emile Mathieu, Bernhard Schölkopf, and José Miguel HeráandezLoobate Se (3) equivariant augmented coupling flows. Advances in Neural Information Processing Systems, 36, 2024.
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobato. Flow annealed importance sampling bootstrap. arXiv preprint arXiv:2208.01 2022.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 147, 2019.
Frank Noé, Hao Wu, Jan-Hendrik Prinz, and Nuria Plattner. Projected and hidden Markov models for calculating kinetics and metastable states of complex molecules. The Journal of Chemical Physics, 139(18):184114 11 2013.
Vijay S Pande, Kyle Beauchamp, and Gregory R Bowman. Everything you wanted to know about markov state models but were afraid ask. Methods, 52(1):99-105, 2010.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 4195-4205 2023.
Guillermo Pérez-Hermández, Fabian Paul, Toni Giorgino, Gianni De Fabritüs,, and Frank Noé. Identification of slow molecular order parameters for markov model construction. The Journal of chemical physics, 139(1), 2013.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pages 28043-2801 PMLR, 2023.
Jan-Hendrik Prinz, Hao Wu, Marco Sarich, Bettina Keller, Martin Senne, Martin Held, John D. Chodera, Christof Schütte, and Frank Noé. Markov models of molecular kinetics: Generation and validation. The Journal of Chemical Physics, 134(17):174105, 2011.
Aneesur Rahman. Correlations in the motion of atoms in liquid argon. Physical review, 136(2A): A405, 1964.
Susanna Röblitz and Marcus Weber. Fuzzy spectral clustering by pecat: application to markov state models and data classification. Advances in Data Analysis and Classification, Jun 2013.
Jean-Paul Ryckaert, Giovanni Ciccotti, and Herman JC Berendsen. Numerical integration of the cartesian equations of motion of a system with constraints: molecular dynamics of n-alkanes. Journal of compulational physics, 2333)322--34, 1977.
Martin K. Scherer, Benjamin Trendelkamp-Schhnn Fabian Paul, Guillermo Pérez-Hernández, Moritz Hoffmann, Nuria Plattner, Christoph Wehmeyer, Jan-Hendrik Prinz, and Frank Noé. PYEMMA 2: A Software Package for Estimation, Validation, and Analysis of Markov Models. Journal of Chemical Theory and Computation,
555525-5542, 2015.
Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit transfer operator learning: Multiple time-resolution models for molecular dynamics. Advances in Neural Information Processing Systems, 36, 2024.
David E Shaw, Ron O Dror, John K Salmon, JP Grossman, Kenneth M Mackenzie, Joseph A Bank, Cliff Young, Martin M Deneroff, Brannon Batson, Kevin Bowers, al. Millisecond-scnt molecular dynamics simulations on anton. In Proceedings of the conference on high performance computing networking, storage and analysis, pages 1-11, 2009.
Hannes Stark, Bowen Jing, Chenyu Wang, Gabriele Corso, Bonnie Berger, Regina Barzilay, and Tommi Jaakkola. Dirichlet flow matching with applications to dna sequence design. arXiv preprint arXiv:2402.0s5 202.44 2024.
Yuji Sugita and Yuko Okamoto. Replica-exchange molecular dynamies method for protein folding. Chemical physics letters. 11411-2) 1999.
12

---

Yann Vander Meersche, Gabriel Cretin, Aria Gheeraert, Shan-Christophe Gelly, and Tatiana Galochk- ina. Atlas: protein flexibility description from atomistic molecular dynamics simulations. Nucleic Acids Research, 52(D1))D384-DDS) D392, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Loup Verlet. Computer" experiments" classical fluids. i. thermodynamical properties of lennard- jones molecules. Physical review, 159(1):98, 1967.
Christoph Wehmeyer, Martin K. Scherer, Tim Hempel, Brooke E Husic, Simon Olsson, and Frank Noé. Introduction to markov state modeling with the pyemma software-v0..
Sherry Yang. Jacob Walker, Jack Parker-Holder, Yilun Du, Jake Bruce, Andre Barreto, Pieter Abbeel, and Dale Schuurmans. Video as the new language for real-world decision making. arXiv preprint arXiv:2402. 17139, 2024.
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. arXiv preprint arXiv:2302.02277, 2023.
Sihyun Yu, Weili Nie, De-An Huang, Boyi Li, Jinwoo Shin, and Anima Anandkumar. Efficient video diffusion models via content-frame motion-latent decomposition. arXiv preprinl arXiv:24031441/ 2024.
Shuxin Zheng, Jiyan He, Chang Liu, Yu Shi, Ziheng Lu, Weitao Feng. Fusong Ju, Jiaxi Wang, Jianwei Zhu, Yaosen Min, et al. Towards predicting equilibrium distributions for molecular systems with deep learning. arXiv preprint arXiv:2306..0 05445, 2023.
13

---

A Method Details
A.1 Flow Model Architecture
Notation. Here, as in the main text, we use the following notation:
T:: number of trajectory frames
number of amino acids
K: number of key frames, with indicies tt...tk
In Algorithms 1-3 below, modify the architectures of Transformeetteeeeeeee
and Diffusion TransformerFinalT..e from DiT (Peebles and Xie, 2023). Elements from these layers are also then incorporated into our custom Invariattoon.AAttentoonngg
Algorithm I: Velocity network
Input: noisy tokens X B7xxx((K+114) conditioning tokens Xcond RTxxx(7K+14) key frame roto-translations 9t, 9tk (SE(3)))* flow matching time t,
amino acid identities A (1....20) conditioning mask m E (0, 1) TxLx(7K+10)
Output: flow velocity ITTxLx(7K+14)
t + Embed(t);
for k +- to K do
Xk: Embed(A) (A) + Ex Linear(gr,, gt)) for / + 1 to num__po Laye__syrssao do X InvariantPoint Atetioottttii gk.
Σx X Linear(()) Linear(Xcond @ m) + Embed(m):
for / <- to mum_trans formarsooorrrppaa do
x DiffrsoorTeaafooommerrttttttteer t)
return Diffudlo//aassoome//aaaalaas////
Algorithm 2: DiffrusioTramsoommmettttttttoat
Input: x E XX time conditioning /.
(a, 3, 7),,f Linear(t);
2 X += ge ① Attention Poth @ LayerNorm(x) + 3e, dim = 1):
9r @ AttentionWithhRo//
@ LayerNorm(x) + 3, dim 0);
4 X = Im @ MLP(1m. @ LayerNorm + Bm));
return
Algorithm 3: InvaranfPPint/AAeettonllaayz
Input: x L time conditioning roto-translations g E SE(()4
(a, 3, y),, Lincar(t);
2 += InvarianttoinntAAttennioo (Layoeooox))a
+ @ Attention WitthioPlytt @ LayerNorm(x) + 3z);
X Im @ MLP(2m. @ LayerNorm Bm));
return x
A.2 Integrating Dirichlet Flow Matching
To additionally generate amino acid identities along with the trajectory dynamics, integrate
SiT flow matching framework with Dirichlet flow matching (Stark al., 2024). Specifically, we
now parameterize velocity network
(IR7 @ 20 TXL
10, 1]
- (R²
20 TXL No
14

---

architecture modifications are necessary other than augmenting the tokens with one-hot tokens of residue identity, broadcasted across time, At training time, we sample from the Dirichlet probability path (rather than the Gaussian path) for those token elements. However, the parameterization is subtle as Dirichlet FM trains with cross-entropy loss, contrary to the standard flow-matching MSE loss. Thus, during training time we minimize the loss
L C= E [ [...,, 20] (x(Xi x,)ll CrossEntropy (atthoopySSoomxxxx)) (..0::)))(........))))...... A)] (5)
That is, we interpret the last 20 outputs in the channel dimension as logits over the 20 residue types. At inference time, on the other hand, we convert these logits to the Dirichlet FM flow field:
r/2
Concat vo [....:-.22.. -20], Σ
Softmax( [.... .:0....,,,2:]]]]. -20:]); UDDE( x1 = =1)
(6)
where UDFM is the appropriate Dirichlet vector field from Stark et al. (2024).
A.3 Conditional Generation
We control the conditional generation settings by simply setting appropriate entries of the conditioning mask m in Algorithm to 1 or O. Specifically,
[1 t=1
For the forward simulation setting, m m = (0 t
LE (1,
For the inpainting setting, m|t, c,c]
0 t t {1,7}
1% M =1
For the upsampling setting,
where M is the upsampling factor.
1 5% M #
l Sknown where Sknown is the set of residues $ Sknown
For the inpainting setting, ml(, f,
in the known part of the trajectory.
We use indexing to be consistent with the main text. In practice, in the inpainting setting we also mask out all torsion angles and withhold the amino acid identities for all residues. Further, do not train the model to generate the torsions as all, such that the tokenization yields X E R These interventions were observed to be necessary prevent overfitting.
B Experimental Details
B.1 Markoy State Models
A Markov State Model (MSM) is representation of a system's dynamics discretized into states (1...r} and a discrete timesteps separated by time lag 7 such that the dynamics are approximately Markovian (Husic and Pande, 2018; Chodera and Noé, 2014; Pande al., 2010). An MSM is parameterized with a vector л that assigns each state a stationary probability and a matrix T containing the probabilities for transitioning from state s; to 8141 after timestep, i.e., T p(8r+1 = 8t = i).
To build a Markov state model, we use PyEMMA (Scherer et al., 2015; Wehmeyer et al.) and its accompanying tutorials. Briefly, we first featurize molecular trajectories with all torsion angles as points on the unit circle, obtaining 2m-ddimensional invariant trajectory where m is the number of torsion angles. We run TICA on these trajectories with kinetic scaling and then run k-means clustering with A 100 over the first few (5-10 chosen by PYEMMA) TICA coordinates. We then estimate an MSM over these 100 states and use PCCA+ spectral clustering (Röblitz and Weber, 2013) to further group these into 10 metastable states. Our final MSM is built from the discrete trajectory over these 10 metastable states. In all cases we use timelag т = 100 ps.
Unconditionally sampling an MSM. To unconditionally sample a trajectory of length N from an MSM, we first sample the start state from the stationary distribution, i.e., 81 π. We then iteratively sample each subsequent state s141 ~ Te..-
15

---

Sampling an MSM conditioned on a start state. To sample trajectory of length N conditioned on starting state $1, we iteratively sample each subsequent state as 8t+1 sT..:
Sampling an MSM conditioned on a start and end state. For our transition path sampling evaluations in Section 4.2, employ replica transition paths sampled from an MSM by conditioning
on a start state $ and end state sN. To do so, we iteratively sample each state between the conditioning states by utilizing the probability
p(s++1 = 1,88 = k)) p(NN =k 8141 8, 1)p(94+1 = 8 =)
p(SN = 1: =(
Firstly, the term p(5(-1 = St = i) is available in out transition matrix as T.p- Secondly, we obtain p(sN k matrix. Specifically p(sN = St = = T:N-0) where the superscript denotes a matrix exponential. Lastly, we obtain the term p(SN k 8t-1 = =j.8t = i) by realizing that under the Markov assumption p(sN k 8141 8, =) p(SN k 88+1 = j). Further,
entry of the (N - t)th matrix exponential of the transition
p(sx = s = j) = T -
Replacing the terms in Equation results in
T T IN-- T
p(8 = = i,SN k))
(8)
T -2)
Thus, sample states 82 ... 88-1 iteratively
T: (N-1--1 1
:.a
Rest
St-1
T
IN-0 SL.8N
(9)
B.2 Tetrapeptide Molecular Dynamics
We run all-atom molecular dynamics simulations in OpenMM (Eastman et al., 2017) using the amber14 14 force field parameters with gbn2 implicit solvent or tip3pfb water model. Initial structures generated with PyMOL, prepared with pdbfixer, and protonated at neutral pH. For explicit solvent, we prepare solvent box with 10 Á padding and neutralize the system with sodium or chloride ions. All simulations are integrated with Langevin thermostat at 350K with hydrogen bond constraints, timestep 2 fs. and friction coefficient 0.3 ps (explicit) or 0.1 ps (implicit). For explicit solvent, nonbonded interactions are cut off 10 Á with long-range particle-mesh Ewald. We first minimize the energy with L-BFGS and then equilibrate the system the NVT ensemble for 20 ps. We then run 100 ns of production simulation in the NVT ensemble (implicit) or NPT ensemble with Monte Carlo barostat at I bar (explicit). We write heavy atom positions every 100 fs.
For explicit-solvent settings (forward simulation, interpolation, inpainting). we run simulations for 3109 training, 100 validation, and 100 test peptides. For implicit-solvent settings (upsampling). we run simulations for 2646 training. 100 validation, and 100 test peptides. All peptides are randomly chosen and split. Additionally, 5195 training and 100 validation implicit solvent simulations are run for the pentapeptide MDGEN.
B.3 Evaluation Details
Trajectory Featurization We featurize trajectories by selecting the sine and cosine of all torsion angles the collective variables. Specifically, we featurize vi, ó backbone angles and all x sidechain torsion angles for each peptide. We then reduce dimensionality with Time-lagged Independent al., 2013) PYEMMA (Scherer et al., 2015).
Components Analysis (TICA) "Pérez-Hernááne
Jensen-Shannon Divergence
We compute the JSD implemented in scipy, i.e.,
D(p m) + D(g m)
(10)
2
16

---

where m = (p+ 9)/2. For the 1-dimenional JSD over torsion angles. we discretize the range -,, x| into 100 bins. For the I-dimensional JSD over TIC-O, we discretize the range spanning the maximum and minimum values into 100 bins. For the 2-dimensional JSD over TIC-0,1 we discretize the space into 50 50 bins.
The autocorrelation of torsion angle 0 at time lag At is defined as (cos(0) B(+20)), corresponding to the inner product of f. B++A1 on the unit circle. To compute the decor- relation time of torsion angle, we subtract the baseline inner product (cos + (sin 0)?, this is analogous to removing the mean of a real-valued time series before computing the autocorrelation. The decorrelation time is then defined as the time required for the autocorrelation to fall below 1/e of its initial value (which is always unity), with the subtracted baseline computed from the reference trajectory. In small number of cases (21 torsions), the MDGEN trajectory did not decorrelate within 1000 frames (10 ns), and we exclude the angle from Figure 2F.
Autocorrelation
To compute the decorrelation time for TIC-O, we now define the autocorrelation as
[ u)( =n]/o²
(11)
where μ4, a are computed from the reference trajectory. Hence, when computed for a sampled MDGEN trajectory, the autocorrelation may not start at unity and may not decay to zero. We report a decorrelation time if starts above and falls below 0.5 within 1000 frames (10 ns), which happens in 74 out of 100 cases as shown in Figure 2E.
Interpolation In our interpolation or transition path sampling experiments, we sample 1000 trajectories of length Ins for each of our 100 test tetrapeptides. We first select start state 81 and an end state AN that exhibits non-trivial transitions. To do so, we consider a reference MD simulation of 100 ns for the tetrapeptide and obtain an MSM as described in Appendix B. From the MSM's transition matrix I and stationary distribution π, we compute the flux matrix F' = T @ Pi where Pi the square matrix with π in each column. The chosen start and end state is the row and column index of the smallest non-zero entry in F.
With the start state $1 and end state s selected, sample 1000 start frames x; and end frames x from the states. The 1000 start frames are sampled from all frames in the reference MD simulation that belong to state 81. Analogously, the end frames sampled from all frames belonging state 8N- Using the 1000 pairs of start and end frames, we condition MDGEN on them and generate discretize these trajectories under the 10-state
trajectories of 100 frames (1 ns). For evaluation, clustering determined by the MSM of the reference MD simulation as described in Appendix B. Note that with the MSM lag time of 100 ps. these discrete trajectories are of length 10.
MD baselines. To sample transition paths of ns between our selected start and end states, we employ MSMs built from replica MD simulations of varying lengths. For instance, for replica MD simulation of 100ns, we first discretize its trajectory with the cluster assignments of the reference MD simulation (the same cluster assignments as we use to discretize the MDGEN ensemble and that use for evaluation). Next, we estimate an MSM from the discretized trajectory. We then proceed to sample 1000 transition paths from the MSM as described in B. where the path is conditioned on an end and start state. In the event that the replica MSM has zero transition probability for transitioning out of the start state or zero probability for transitioning into the end state (this occurs if the replica MD simulation never visited the start or end state), we treat all 1000 paths of the replica MD as having zero probability for our evaluation metries which are further detailed in the following.
Computing TPS metrics. As described above, we obtain ensembles of 1000 discretized Ins paths of 100 frames for both MDGEN and the replica MD simulations. For these, in Figure 3, we show JSD, the rate of valid paths, and the average path probability. These metrics are computed with respect to the MSM of the reference MD simulation of length 100 ns.
To compute the JSD, we draw 1000 discrete transition paths from the reference MD sim- ulation and compute the probability of visiting each state from the frequency with which each state is visited. We do the same for the transition path ensemble of MDGFN (or the baseline) and compute the JSD between the categorical distributions as described above.
The average path probability for an ensemble is the average of its paths' probabilities for transitioning from the start to the end state under the reference MSM. This probability can be computed as described in Appendix B.
17

---

* The valid path rate is the fraction of paths that have non-zero probability.
Inpainting In our inpainting experiments, we set out design tetrapeptides that transition between two states. Considering the residue indices 1, 2, 3 and 4, we call the residues 1, 4 the flanking residues which we condition on and 2.3 3 the inner residues which we aim to design. Specifically, we condition MDGEN on the trajectory of the flanking residues' backbone coordinates and generate the residue identities of the inner two residues. To carry out this design for a single tetrapeptide, we draw 1000 samples from MDGEN to estimate the mode of its joint distribution over the inner two residues.
The conditioning information (the trajectories of the outer two residues' backbone coordinates) is different for the two evaluation settings of designing transitions with high flux or for designing arbitrary transitions. However, for both of them, the start and end frames are provided as conditioning information via the key frames. In the high flux setting, the conditioning information is obtained by sampling 1000 paths from the reference MD simulation of length 10 ps with 100 frames that start and end in the desired states. These states are determined as those with the maximum flux between them (see the paths about interpolation above for a description of flux). When designing residues that give rise to arbitrary random paths, the trajectories are randomly sampled from the reference simulation.
After sampling 1000 pairs of residues for the inner occurring pair as the final design. For this design, we report the sequence recovery (the fraction of residues that match the original sequence of the MD simulations from which the conditioning information was sampled).
residues, we select the most frequently
Inpainting Baselines. We aim to assess the benefit that is obtained by the trajectory- inference of MDGEN over a baseline that only takes the start frame or the start and end frame as input for construct DYNMPNN and S-MPNN.
designing residues that transition between two states. Thus, These baselines use the same architecture as MDGEN in the inpainting setting, but DYNMPNN only obtains the start and end frames as key frames and via their roto-translation offsets for the first and last frames. S-MPNN is the analog with only the first frame.
Notably, in the inpainting setting. MDGen and the baselines do not treat torsion angles, and all torsion angle entries of the SE(3)-invariant tokens are set to 0. Furthermore, the model does not take the amino acids of the flanking residues as input. We make this choice of withholding all information about amino acid identities since otherwise, the models overfit on the arbitrary identities of the flanking residues and do not generalize to the test set.
Protein Simulations For training and evaluation on proteins, we use trajectories from the ATLAS dataset (Vander Meersche et al., 2024), which includes 3 replicates of 100 ns explicit-solvent, all-atom simulations for each of 1390 non-membrane protein monomers. The proteins are chosen from the PDB representatives of all available ECOD domains and are thus structurally non-redundant. We split the dataset into 1265 training, 39 validation, and 82 test proteins by PDB release date following Jing al. (2024). At training time, randomly select protein, select one of the three replicates, subsample every 40 frames, obtaining a training target with 250 frames. We train with random crops of up to 256 residues, but draw samples for the full protein at inference time. To compute statistical similarity of the MDGEN ensembles with the ground truth MD ensembles, we compare the 250 frames with 30k pooled frames from all three trajectories. Baseline metries and runtimes for AlphaFlow and MSA subsampling are taken directly from Jing et al. (2024). Analysis and visualization code for Table 4 and Figure are provided courtesy of Jing al. (2024).
Runtime MD runtimes in Table 2 are tabulated on a NVIDIA T4 GPU. All MDGEN experiments are carried out on NVIDIA A6000 GPUs. AlphaFlow and MSA subsampling runtimes in Table 4 are tabulated on NVIDIA A100 GPUs by Jing et al. (2024).
18

---

Additional Results
C.1
Forward Simulation
LRHH
WWA
LIFE
FLRH
MAFM
VDRN
EHEV
WSAQ
ESIC
CSYR
Figure 7: Additional backbone torsion angle distributions (orange from MD, blue from samples) and free energy surfaces along the top two TICA components for 10 randomly chosen test peptides.
LIRH p= 0.40
LIFE p 0.80
MAFM p 0.68
EHEV = 0.65
ESIC p= 0.54
IVMA p= 0.79
FLRH p= 0.64
VDRN p= 0.40
WSAQ p 0.47
CSYR p 0.60
Figure 8: Flux matrices between MSM metastable states computed from reference trajectories (upper right) and MDGEN trajectories (bottom left) for 10 random test peptides (the matrices symmetric). Cells are colored by the square root of the flux, with darker indicating high flux. The Spearman correlation between the entries is shown.
19

---

C.2 Interpolation
MDGce FES
MDGen FES
NDGer FES
NDGer FBS
MDDGen pon FES
NDGan poth FES
NDGer path 3 FTFES
NDGer path 4 FES
AKIR
ASRE
MDDeen poth IES
NDGen both 2 TES
NDGer 2oth 3 TES
NDGHr euth IES
CFEE
CIYL
MDGen path 1n FES
NDGen 9oth 2 FES
NDGer path In FES
NDGen bath in FES
DLCG
FDNT
Figure 9: Four of 1000 transition paths MDGEN for several tetrapeptides in the test set.
20

---

C.3 Upsampling
Autocorrelation
Dynamical content
Backbone
Sidechain
Upsampled
Ground truth
1.0
PGKM
0.8
0.7
0.5
10
10
10
100
1000
10
100
1000
1.0
ps
PS
0.9
IWHF
0.7
0.6
10
0.9
SSNN
0.8
0.6
0.5
1.0
ps
10
ps
10
10
ps
100
1000
10
ps
100
1000
KSIY
0.8
0.7
0.5
1.0
05
10
10
10
100
1000
10
100
1000
0.9
ALDA
0.7
0.6
10
1000
Figure 10: Recovery of fast dynamics via trajectory upsampling for random test peptides. (Left) Autocorrelations of each torsion angle from (-) the original 100 fs-timestep trajectory, (o) the subsampled 10 ns-timestep trajectory, and (...) the reconstructed 100 fs-timestep trajectory (all length 100 ns). (Right) Dynamical content as function of timescale from the upsampled vs. ground truth trajectories, stacked for all torsion angles (same color scheme). The subsampled trajectory contains only the shaded region and our model recovers the unshaded region.
21

---

C.4
Inpainting
Zero diagonal Start End
NSM
GCIC
Zero diagonal Start End
RES
MSM graph. Start End
TLRK
Zero diagonal Flux. Start End
FBS
NSM graph. Start End
MAFM
Zero diagonal Flux. Start 1. End 9.
rTS start
NSM as graph. Start End 9.
HELI
Zero diagonal Fux, Start 5. Lind it.
FES with start and and stare
MSM as graph. Start End B.
EHEV
Figure 11: For six tetrapeptides, we show the states that we chose in our design experiments when designing transitions between the highest flux states. Column I shows the flux matrix with zeros on the diagonal. Column 2, the free energy surface of a 100 ns simulation and the selected start and end states based on the highest flux in the flux matrix. Column 3, the MSM that was built from the MD simulation.
22

---

C.5
Structural Validation
In addition distributional similarity and dynamical content, we also assess the frequency of clashes or high-energy structures in MDGEN forward simulation rollouts. Specifically, we compute the distributions of:
The closest distance between any pair of nonbonded atoms
Nonbonded energy (Coulomb + Lennard-Jones)
Torsional energy
Heavy atom bond lengths
Radius of gyration
These distributions are shown and compared to the ground truth in Figure 12. We find that the vast majority of MDGEN structures are high quality (i.e., clashes are rare) and adhere closely the ground truth distributions.
Nonbonded energy (k)
Tersional energy
(A]
Reference MDGen
1.8
-2000
10000
2000
250
500
750
3000
1.00
1:25
1.50
1.75
2.00
Figure 12: Histograms of various structural validation metrics between MDGEN forward simulation and reference trajectories, pooled across all test tetrapeptides.
C.6 Additional Comparisons
In this section, we compare MDGEN to Timewarp (Klein et al., 2024) and ITO (Schreiner et al., 2024), generative models for autoregressively rolling out surrogate MD trajectories. Note that these comparisons are limited the forward simulation task as Timewarp and ITO are not capable of solving the other tasks.
For Timewarp (Klein et al., 2024), we use the 4AA model with weights from the authors and sample 100 trajectories by running 2000 inference steps with timestep 50 ps. We do not use MH acceptance steps as the authors found exploration of the energy landscape to be much more effective without them.
For ITO (Schreiner et al., 2024), transferable models across tetrapeptides are not available. We therefore re-train ITO on our tetrapeptide dataset with timesteps 500 ps. We then run 200 inference steps to sample 100 ns trajectories.
For both methods, we observe that trajectories are unstable without further intervention. To bolster the baselines, proceed with further analysis. We note that the Timewarp authors, in lieu of relaxation, rejected steps with energy increase of 300 kJ / mol (Klein et al., 2024): however, would reject the majority of proposed steps on generic tetrapeptides. In Figure 13, we visualize the free energy surfaces and torsion angle distributions for several peptides from Timewarp and ITO compared with MDGen. Table shows the Jensen-Shannon divergences from the forward simulation rollouts across all test peptides fc.F. Table 2). Qualitatively and quantitatively, our model obtains better consistency with the ground-truth free energy surfaces.
OpenMM (Eastman et al., 2017) relaxation steps between each timestep to
found that this strategy
23

---

Timewarp torsions
ALDA
GLIL
FLRH
Figure 13: Comparison of MDGEN, Timewarp, and ITO in terms of forward simulation evaluations on test peptides. (Left) Torsion angle distributions for the six backbone torsion angles from MD trajectories (orange) and sampled trajectories (blue). (Right) Free energy surfaces along the top two TICA components computed from backbone and sidechain torsion angles.
Table 5: Comparison of MDGEN, Timewarp, and ITO in terms of the JSD between sampled and
ground-truth distributions along various collective variables in the forward simulation setting.
C.V.
MDGEN
Timewarp
ITO
Torsions (bb)
.130
.325
564
Torsions (sc)
.033
.427
.462
Torsions (all)
.109
.383
.055
TICA-0
.230
.265
.538
TICA-0,1 joint
.366
.756
MSM states
.255
.414
Runtime
60s
599s
2083s
24

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We outline all contributions of the paper, emphasizing the ones with robust support and qualifying the ones for which the support is preliminary.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: See Section 5 for a discussion of limitations.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
25

---

Justification:
The paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide details sufficient to reproduce our model and experiments in Appendix A and Appendix B. In particular, the model architecture is detailed in Algorithm 1.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
26

---

Answer: [Yes]
Justification:
We have released our code and data under MIT license.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nin) cc/
http://ggiide//oeeuuieeeee httpibimsssiooooooiiooooooooeeeo lcy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
//ips.cc. wwt.cc//ppllcccc gudddee/ooooeuummmii /Cc////odeeesesseeeeeeeeeee onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Training and test details are provided throughout the main text, figure and table captions, and Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: Our results are reported across a large test set for the task (n = 100) and focus more on the new qualitative capabilities that we introduce rather than quantitative improvements on benchmarks. Thus, we did not feel it was necessary to report error bars but would be happy provide them upon request.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
27

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative етrоr rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Our models small by modern standards and did not feel it neces- sary to explicitly highlight the compute time required for experiments. We provide some information computational resources in Appendix B. Further, the runtimes reported in Tables 2 and 4 provide an idea of the efficiency of our models.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: htap::::////urpppe wwtppcpppp//eeeeeeeeeeeeeeeeeeee h ines?
Answer: [Yes]
Justification: Ethics.
All authors have read and adhered to, in every respect, the NeurlPS Code of
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification: This paper presents work whose goal is to advance the field of Al for Science. There are many potential socictal consequences of our work, none which we feel must be specifically highlighted here.
Guidelines:
28

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negatíve societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper does not release data or models that have high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We provide the original citations of major open-source software packages and datasets used in our work. We did not feel it was necessary to explicitly mention their open-source licenses as all are publicly available.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
29

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Model code and weights for reproducibility are provided via a public repository with instructions for replicating training and evaluation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question:
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
30

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
31


Reviews:
[{"review_id": "sOvr0i0Q2A", "reviewer": "Reviewer_QAiG", "summary": "The authors introduce MDGen as a novel approach for modeling MD trajectories. They demonstrate the capabilities of this method in tasks such as interpolation, upsampling, and inpainting of small peptides. The accuracy as well as speed of the new approach compared to the ground truth baseline is quantitatively evaluated. Initial experiments toward upscaling to small proteins are shown.", "questions": "1. How often do clashes and other high-energy structures occur in the generated trajectories?\n2. When comparing to other methods and approaches in the experimental section - do all of them use a similar reduced representation or do the other methods generate all-atom representations?", "limitations": "The approach is limited to peptides. Transfer to any other molecules is questionable due to a lack of suitable representation/tokenization.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The idea of MDGen is novel and very well presented in this manuscript. The results are convincing and interesting.", "weaknesses": "1. Parts of Sections 3.1 and 3.2 are very condensed and hard to follow. A more detailed description in the SI would be helpful, where the most important aspects of the cited work is also repeated.\n2. The suitability of the chosen representation for longer amino acid chains is questionable. This is also mentioned in the manuscript, but nonetheless, proteins are mentioned many times (more than 30) in the manuscript, while almost all experiments are actually performed on very small peptides. It should be stated in a more prominent place that upscaling to proteins is not trivial.\n3. The representation limits the model to learn MD trajectories of natural amino acids, as no all-atom representation is used directly. This should be made clearer in the manuscript.\n\nMinor points: A lot of figures have no proper axis labels (e.g. Fig 3, 4, 5, 6). This should be fixed. The best models in Table 4 should be indicated in bold.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review! To address your questions and concerns:\n\n---\n\n**Parts of Section 3 hard to follow**\n\nDue to space limitations, the description of our method in Section 3 was indeed a bit condensed. We will expand the exposition with the extra page allotted in the revision.\n\n**Suitability for proteins and non-AA molecules**\n\nWe sought to qualify our results on protein simulations and apologize if a different impression was conveyed. In the revision we will increase emphasis that the focus of the paper is on peptides and remove prominent mention of proteins from the abstract and introduction, as suggested.\n\nWith that said, we have focused on tetrapeptides as model systems in this work for two key reasons:\n* We can run simulations for thousands of systems in order to properly test the generalization abilities of our model.\n* We can build Markov State Models for each system, allowing the careful benchmarking for forward simulation and interpolation capabilities.\n  \nBoth of these aspects are important for the thorough, careful benchmarking of the new capabilities we demonstrate. We opted to prioritize these careful studies to support the core claims of our work, in lieu of expanding its scope to larger and more diverse systems, which we think are best left to future work.\n\n**Clashes in the generated structures**\n\nTo assess the frequency of clashes or high-energy structures in MDGen forward simulation rollouts, we compute the distributions of\n* The closest distance between nonbonded atoms\n* Nonbonded energy (Coulomb + Lennard-Jones)\n* Torsional energy\n* Heavy atom bond lengths\n* Radius of gyration\n\nThese distributions are shown and compared to the ground truth in **Figure 1** in the PDF attached to the global response. We find that the vast majority of MDGen structures are of high quality (i.e., clashes are rare) and adhere closely to the ground truth distributions.\n\n**Representations of other methods**\n\nThe peptide representations of the other methods are summarized as follows:\n\n* For most evaluations we compare to Markov state models, which emit discretized representations of the trajectory. These are much more coarse-grained than our reduced representation.\n* AlphaFlow emits residue frames and torsion angles, an equivalent representation to ours.\n* The inpainting baselines are inverse folding models and do not emit trajectories of any kind.\n* Ground truth molecular dynamics provides all-atom trajectories.\n\n**Minor points**\n\nThank you for noting these; we will incorporate them in the revision.\n\n---\n\nWe hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.", "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": "UNZ4f7c61A", "reviewer": "Reviewer_TTLp", "summary": "The paper presents a novel generative model for molecular dynamics (MD) trajectories called MDGEN. This model aims to serve as a flexible surrogate for MD simulations by generating entire trajectories conditioned on initial frames. It addresses tasks such as forward simulation, transition path sampling, trajectory upsampling, and dynamics-conditioned molecular design. The model is evaluated on tetrapeptide simulations and demonstrates its capability to generate reasonable ensembles of protein monomers.", "questions": "How can the model be adapted or improved to reduce its reliance on key frames?\n\nExploring techniques for unconditional generation or alternative ways to handle the roto-translations without predefined key frames could enhance the model's flexibility and applicability.\nWhat architectural changes or enhancements could improve the model's performance on larger molecular systems such as proteins?\n\nInvestigating more scalable architectures or hybrid approaches that combine the current method with other techniques tailored for large systems could address this limitation.\nHow does the computational cost of training the model compare to traditional MD simulations, and what are the implications for its practical use?\n\nProviding detailed information on computational requirements and potential optimizations could help in assessing the model's feasibility for widespread use.\nWhat alternative tokenization strategies could be explored to extend the model's applicability to a wider range of molecular systems?\n\nResearch into tokenization methods that can handle diverse molecular structures and dynamics could broaden the model's utility.\nHow can additional conditioning types (e.g., textual descriptions, experimental data) be incorporated into the model, and what benefits might they provide?\n\nExperimenting with and integrating various forms of conditioning could enhance the model's ability to generate more accurate and contextually relevant trajectories.\nWhat are the potential impacts of data quality and availability on the model's performance, and how can these challenges be mitigated?\n\nAddressing data-related challenges through techniques like data augmentation, transfer learning, or synthetic data generation could improve the model's robustness and applicability.\nCan additional evaluation metrics be developed to provide a more comprehensive assessment of the generated trajectories' quality?\n\nIdentifying and implementing new evaluation criteria could offer deeper insights into the strengths and limitations of the model's output.", "limitations": "Reliance on Key Frames:\n\nThe model relies on key frames for parameterizing the roto-translations, which means it cannot perform unconditional generation or inpainting of residue roto-translations. This dependency might limit its applicability to scenarios where key frames are not easily obtainable or where full trajectories need to be generated from scratch.\nScalability to Larger Systems:\n\nThe architecture shows weaker performance on larger systems such as proteins compared to smaller peptides. This suggests that the current model and architecture might not be well-suited for handling the complex motions and larger size of protein structures without further modifications or enhancements.\nComputational Resources:\n\nWhile the paper mentions significant speedups compared to traditional MD simulations, the computational resources required for training the model (e.g., GPU hours) are not explicitly discussed. This information is crucial for understanding the practicality and scalability of the approach.\n\nGeneralization to Diverse Systems:\n\nThe current tokenization and modeling strategies are tailored to peptides and proteins. For more diverse molecular systems such as organic ligands, materials, or explicit solvent systems, alternative tokenization strategies might be necessary. This limits the immediate applicability of the model to a broader range of molecular simulations.\nLimited Exploration of Additional Conditioning:\n\nThe paper primarily explores conditioning on initial frames and residue identities. Other types of conditioning, such as textual or experimental descriptors, are not explored but could open up further applications and improve the model's utility.\nData Availability and Quality:\n\nThe success of the model heavily depends on the availability of high-quality MD trajectory data. For many complex systems, obtaining such data can be challenging, which could limit the model's applicability and performance.\nEvaluation Metrics:\n\nWhile the paper uses several rigorous metrics for evaluation, the choice of metrics may not fully capture all aspects of the generated trajectories' quality. Additional metrics or more diverse evaluation criteria could provide a more comprehensive assessment.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Novelty and Scope: The approach introduces a novel paradigm for surrogate modeling of MD, extending the capabilities of existing models to handle a variety of tasks that are not straightforward with current methods.\nGenerative Framework: The use of generative modeling for entire MD trajectories is a significant advancement, as it allows for a broader range of applications including forward and inverse problems.\nComprehensive Evaluation: The paper evaluates MDGEN on several tasks, demonstrating its effectiveness in forward simulation, interpolation, upsampling, and inpainting. The results show promising performance in terms of distributional similarity, dynamical content, and computational efficiency.\nTechnical Implementation: The detailed description of the tokenization process and the flow model architecture provides a clear understanding of how the model operates. The use of SE(3)-invariant tokens and the scalable interpolant transformer (SiT) backbone are well-motivated choices.", "weaknesses": "Complexity and Accessibility: The model’s complexity might pose challenges for reproducibility and accessibility for researchers who are not deeply familiar with both molecular dynamics and advanced generative modeling techniques.\nEvaluation on Larger Systems: While the paper provides proof-of-concept evaluations on proteins, the primary focus remains on smaller tetrapeptides. The model's scalability and effectiveness on larger and more complex molecular systems need further exploration.\nDependence on Key Frames: The reliance on key frames for conditional generation limits the model’s ability to perform unconditional generation or inpainting of residue roto-translations, which could be a significant limitation in certain applications.\nComputational Resources: The paper lacks detailed information on the computational resources required for training and inference, which is crucial for understanding the practical implications of using MDGEN in various research settings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review! To address your questions and concerns:\n\n---\n\n**Complexity and accessibility**\n\nWe have aimed to provide a clear and reproducible method and exposition accessible to the average reader familiar with molecular machine learning. We aimed to make modeling choices that were as simple as possible, i.e.,\n* using a simple stochastic interpolants framework rather than diffusion\n* using a vanilla unmodified transformer from previous work rather than complex frame- or pair-based architectures\n* generating invariant tokens over Euclidean space rather than diffusing over the space of frames and tokens.\n\nTo promote reproducibility, we have\n* provided pseudocode for our architecture in Appendix A\n* provided experimental details sufficient for replications and conceptual justifications for design choices and metrics in Appendix B\n* will provide full source code in the revision\n  \nWhile our work bridges molecular dynamics and generative models, it should not require a deep background in either as we have distilled the essential aspects of the relevant literature in Section 2. With that said, if there are any specific aspects of the method that remain unclear despite our best efforts, please let us know.\n\n**Reliance on key frames**\n\nWhile it is true that the use of key frames impacts the ability to do unconditional generation, we nevertheless opted for this design choice as unconditional generation of trajectories is (to our knowledge) not a problem of scientific interest. On the other hand, the key frames allow the conditional modeling problems to be consideribly simplified and clarified, leading to the strong results shown in our experiments. Hence, the use of key frames could instead be considered a technical insight of our model that enables it to contribute to the development of real-world, impactful scientific problems, at the cost of problems of lesser interest.\n\n**Larger and more diverse systems**\n\nWe have focused on tetrapeptides as model systems in this work for two key reasons:\n* We can run simulations for thousands of systems in order to properly test the generalization abilities of our model.\n* We can build Markov State Models for each system, allowing the careful benchmarking for forward simulation and interpolation capabilities.\n  \nBoth of these aspects are important for the thorough, careful benchmarking of the new capabilities we demonstrate. We opted to prioritize these careful studies to support the core claims of our work, in lieu of expanding its scope to larger and more diverse systems, which we think are best left to future work.\n\n**Conditioning**\n\nWe agree that further types of conditioning would be exciting to explore. However, the conditional settings we have already provided represent significant conceptual shifts relative to existing work in learning surrogate models of molecular dynamics. Additional types of conditioning would further expand the technical scope of the work, and the ones we mention (i.e., text conditioning) are highly speculative with large uncertainty with regards to scientific utility, at least at present. With that in mind, it is not clear to us which additional settings are being referenced as weaknesses or limitations. If there are specific areas in mind, please let us know.\n\n**Data availability and quality**\n\nWe acknowledge that MD data is required to train the method, and obtaining long trajectories can be time-consuming. However, compared to molecular ML modalities requiring _experimental_ data, such as those that train on crystal structures, it is much easier to obtain high-quality data by running simulations than by running wet-lab experiments. Hence, works like ours that train on _simulated_ data actually help  mitigate the challenges of data availability and quality in molecular machine learning.\n\n**Evaluation metrics**\n\nWe have designed and implemented thorough and principled metrics that assess the ability of our model to produce trajectories:\n* With good distributional properties over structures, assessed by the Jensen-Shannon divergences and Markov state occupancies\n* With good dynamical properties over motions, assessed by the various transition path metrics, the decorrelation times, flux matrix correlations, autocorrelation functions, and dynamical content.\n  \nWhile other metrics are certainly possible, it is not clear to us which areas are concretely being referenced as weaknesses or limitations of the current evaluations. If there are specific areas in mind, please let us know.\n\n**Computational resources**\n\nWe report our training time for each model as follows, measured on A6000 GPUs:\n\n* Forward simulation: 412 GPU-hrs\n* Interpolation: 272 GPU-hrs\n* Upsampling: 292 GPU-hrs\n\nWe have reported inference time comparisons with MD simulations in Table 2 on the forward simulation tasks, all on one A6000 GPU. Our method is much faster than the baseline simulations, and much more accurate than an abridged simulation of the same wall-clock time.\n\nEmphatically, the training time is amortized across all systems since our method is _transferable_ by design. That is, for the fixed training cost we obtain a set of models that can be applied, with the reported accuracy, to _any_ tetrapeptide system at test time. Thus, our model training offers a clear and overwhelming advantage over running long MD simulations for individual tetrapeptide systems.\n\n---\n\nWe hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.", "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": "5iKIq4YDUd", "reviewer": "Reviewer_b3SL", "summary": "In this work, the authors proposed MDGen, a new framework that aims to model molecular dynamics trajectories via generative modeling techniques. By properly encoding the Protein MD trajectories according to the characteristics of key frames, MDGen adopts flow matching techniques (both continuous and discrete flow matching) to generatively model MD trajectories. As a unified framework, MDGen is able to perform diverse tasks including forward simulation, interpolation, upsampling and inpaiting. Extensive experiments are conducted to demonstrate the effectiveness of MDGen.", "questions": "N/A", "limitations": "The authors carefully discuss the limitations of this work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The problem this work aims to tackle is of great significance in scientific domains lie computational biology.\n2. The formulation of molecular (protein) trajectories by using key frame references is reasonable and compact for reducing the modeling difficulties.\n3. The experiments are comprehensive.\n4. The paper is well-written and easy to follow.", "weaknesses": "1. Lack of discussion on related works. This work does not discuss related works on the same topic. Some works are mentioned in the Introduction section, but I still recommend that there should be an independent Related Works section for comprehensive discussion.  Here are also several works that are worth discussing: (1) EGNO, which uses neural operator learning approach to also model the trajectory dynamics of molecules; (2) DiffMD, which uses diffusion models to simulate molecular dynamics. The quality of this work should be further improved if the authors could carefully discuss the differences between MDGen and these works and the strengths of MDGen compared to these works.\n\n2. Lack of ablation studies. MDGen is composed of several parts, including the design of the backbone model, the design choices of flow matching framework, and the adoption of Hyena architecture for efficiency consideration. In addition to the aimed tasks, it would further improve the quality of this work if the authors could conduct ablation studies on these aspects to help readers know what the influence of each part of MDGen is.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review! To address your questions and concerns:\n\n---\n\n**Further discussion on related works**\n\nThanks for mentioning these works. We are happy to discuss them in the revision alongside the existing related work in the Background which we will also expand as per your suggestion. Briefly:\n\n\n* EGNO is a fourier neural operator that, given a molecular structure and a time delta, *deterministically* predicts the structure after the time delta. The time delta is constrained to be between 0 and a maximum time which in practice is 5fs - this is 20,000,000 smaller than the duration of our trajectories in forward simulation. Notably, the possible structures after a time delta (in the absence of the initial velocity) are a distribution, which EGNO fails to model since it is not a generative model. \n* DiffMD adds noise to molecular structures and denoises them to generate a future structure. Repeating the process autoregressively yields a trajectory.\n\nWhen producing trajectories of timescales similar to MDGen, both methods produce their trajectories autoregressively which limits them to forward simulation as an application and prevents past frames from being informed by future frames. Meanwhile, MDGen generates all frames in a trajectory jointly, which can improve temporal coherence and enables addressing additional tasks such as interpolation, upsampling, and inpainting.\n\n\n**Lack of ablation studies**\n\nWe now provide new experimental results ablating components of the method. The ablations are as follows:\n\n* **No IPA embedding**: the model is not provided IPA embeddings of the key frames and amounts to a vanilla Scalable Interpolant Transformer\n* **No $SE(3)$ invariance**: there is no SE(3)-invariant tokenization and the model operates directly over the frame representations $g_j$ rather than frame offsets $g_i^{-1}g_j$\n* **No frames/torsions**: the model operates directly over all-atom coordinates rather than the reduced representation in $\\mathbb{R}^{21}$\n\nWe run these ablations on the forward simulation experiment and obtain the following Jensen-Shannon divergences (c.f. Table 2). For a fair comparison, we use the baseline model after training for the same number of epochs as the ablations. The ablations all perform worse than the baseline model.\n\n|C.V.| Baseline | No IPA embedding | No $SE(3)$ invariance | No frames/torsions |\n|:-| :-: |  :-: | :-: | :-: |\n|Torsions (bb)| **0.161** | 0.339 | 0.195 | 0.537 |\nTorsions (sc) | **0.106** | 0.249 | 0.262 | 0.502 |\nTorsions (all) | **0.130** | 0.287 | 0.233 | 0.517 |\nTICA-0 | **0.245** | 0.375 | 0.298 | 0.510 |\nTICA-0,1 joint | **0.345** | 0.500 | 0.416 | 0.657 |\nMSM states | **0.237** | 0.250 | 0.278 | 0.408 |\n\n---\n\nWe hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.", "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": "XSDFwbQNut", "reviewer": "Reviewer_zaG3", "summary": "The paper presents a new framework for generating trajectory of molecular geometries, ie, generative modeling for molecular dynamics. The paper proposes tokenization methods to tokenize the trajectory and learn flow models on the data. Experiments demonstrate the effectiveness of several tasks including forward sampling,  interpolation, and up sampling.", "questions": "I think existing methods can also tackle several tasks. For example, for the forward sampling task, previous generative MD models like Timewarp (Klein et al., 2024) and ITO (Schreiner et al., 2024) can also be used for the task. A numerical comparison with these baselines can help to justify the effectiveness of the proposed method.", "limitations": "Limitation is nicely discussed in the paper.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper tackles a new problem in molecular dynamics generation, which has not been explored in existing literature.\n\n2. The paper is in good structure and easy to follow.\n\n3. The paper provides a detailed analysis of several domain tasks on interested molecular structures, which demonstrate the critical usage in some scenarios.", "weaknesses": "1. Limited ML technical contribution, as all components exist in previous molecular generative models.\n\n2. The experiment is comprehensive from a domain perspective. However, I feel the experiments lack some benchmarking comparison with state-of-the-art molecular generative models for related tasks. See my question below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review! To address your questions and concerns:\n\n---\n\n**Limited ML technical contribution**\n\nIn writing the paper, we opted to placed more emphasis on the experimental results. Nonetheless, we respectfully disagree that our work has limited ML technical contribution.\n\nWe highlight that the following points are novel from a modeling perspective and, to our knowledge, have not been used in any prior molecular generative model:\n* The reduced representation of residue all-atom coordinates into $\\mathbb{R}^{21}$.\n* The use of key frames to obtain SE(3)-invariant tokens from molecular trajectories.\n* The use of scalable vanilla transformers for molecular generation while respecting the problem symmetries (i.e., without breaking equivariance like molecular conformer fields or AF3).\n\nIndeed, as we write in the Introduction, our approach provides a novel strategy for how to circumvent the expensive architectures normally used for equivariant molecular modeling, a technical contribution that could be useful for many future works.\n\nAdditionally, we believe that the insight of using trajectory generation to solve inverse problems like interpolation or upsampling in itself constitutes methodological novelty, akin to the novelty appreciated in the many surprising applications of diffusion models in RL.\n\n\n\n**Comparison with Timewarp and ITO**\n\nWe now provide new results comparing our work to Timewarp and ITO. Emphatically, these comparisons are limited to the forward simulation task as Timewarp and ITO are not capable of solving the other tasks. \n\n* For **Timewarp**, we use the 4AA model with weights from the authors and sample 100 ns trajectories by running 2000 inference steps with timestep 50 ps. We do not use MH acceptance steps as the authors found exploration of the energy landscape to be much more effective without them.\n* For **ITO**, transferable models across tetrapeptides are not available. We, therefore, train ITO ourselves on our tetrapeptide dataset with timesteps of 500 ps. We then run 200 inference steps to sample 100 ns trajectories.\n\nFor both methods, we observe that trajectories are unstable without further intervention (see **Figure 2** of the PDF attached to the global response). To bolster the baselines, we run OpenMM relaxation steps between each timestep to proceed with further analysis. We note that the Timewarp authors, in lieu of relaxation, rejected steps with an energy increase of 300 kJ / mol; however, we found that this strategy would reject the majority of proposed steps on generic tetrapeptides.\n\nIn **Figure 3** of the PDF attached to the global response, we visualize the free energy surfaces and torsion angle distributions for several peptides from Timewarp and ITO compared with MDGen. Our model obtains much better results across the board.\n\nQuantitatively, we obtain the following Jensen-Shannon divergences from the forward simulation rollouts (c.f. Table 2):\n\n|C.V.| MDGen | Timewarp | ITO |\n|:-| :-: |  :-: | :-: | \n|Torsions (bb)| **0.130** | 0.325 | 0.564 |\nTorsions (sc) | **0.093** | 0.427 | 0.462 |\nTorsions (all) | **0.109** | 0.383 | 0.505 |\nTICA-0 | **0.230** | 0.265 | 0.538 |\nTICA-0,1 joint | **0.316** | 0.419 | 0.756 |\nMSM states | 0.235 | **0.222** | 0.414 |\nRuntime (s)| **60** | 599 | 2083 |\n\nTo evaluate the dynamics, we also compare the Pearson correlation of predicted torsional relaxation times (c.f. Figure 2F). We could not compute these quantities for ITO as the vast majority of torsions did not decorrelate within the ITO rollouts.\n\n|C.V.| MDGen | Timewarp | ITO |\n|:-| :-: |  :-: | :-: |\nTorsion relaxation (bb)| **0.55** | 0.04 | ---|\nTorsion relaxation (sc) | **0.97** |  0.77 | ---|\n\n\nAltogether, the results confirm that MDGen is more successful and efficient at modeling thermodynamics and kinetics than Timewarp or ITO.\n\n---\n\nWe hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.", "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": "DygF786bk1", "reviewer": "Reviewer_hVZ7", "summary": "The authors propose MDGen -- a generative model to sample molecular dynamics trajectory conditioned on key frames. This is a direct application of video generation techniques to solve domain challenges in protein modeling. Specifically, SiT and flow matching models are used to sample SE(3)-invariant representation of all-atom protein representations. This work demonstrates the effectiveness of MDGen primarily on tetrapeptide systems, where the authors showcase four downstream application tasks including forward simulation, interpolation, upsampling, and dynamics-conditioned inpainting. \n\nIn general, I find this manuscript well-written and easy to follow. The model performance looks reasonable on tetrapeptides, yet the results are proof-of-concept in nature and generalization to larger proteins remain challenging. However, it is one of the pioneering work in AI protein modeling to directly emulate MD simulation trajectories using data-driven approaches. To that end, I think it would be beneficial for this work to gain visibility across the research community to inspire future studies.", "questions": "- L#141, when $K > 1$, how to ensure roto-translation prediction consistency across $K$ key frames and obtain a final $\\hat{g}_j$?\n- Table 2, with 100 ns being the ground truth, the non-zero JSD in the last column originates from subsampling the simulation trajectory?\n- Figure 2F. My understanding is that sidechains exhibit faster dynamics while backbone motions are slower. The low correlation for backbone suggests that MDGen is not good at learning slower dynamics, which are typically more interesting to researchers?\n- Temporal coherence between generated protein conformations is mainly evaluated using auto-correlation in this work. Is it possible to show other metrics to capture detailed structural quality and variation during time evolution?\n- Why is MDGen more effective at sequence recovery than MPNN? More explanation and analysis would be helpful here.\n- Would it be possible to emulate MD simulation trajectory of the 12 fast folding proteins from [Shaw 2009](https://dl.acm.org/doi/abs/10.1145/1654059.1654126)? They are smaller than ATLAS proteins and longer than tetrapeptides, with much longer simulation time and rich dynamics.\n- It would be nice to see if MDGen could infer a trajectory given an [apo/holo pair](https://arxiv.org/abs/2304.02198).", "limitations": "See weaknesses and questions.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- It is one of the pioneering work to adopt video generation techniques for MD trajectory generation. Although conceptually straightforward, good practices to generate time-coherent MD trajectories across different protein systems remain underexplored.\n- The authors demonstrated a variety of downstream tasks using the same model architecture. The underlying modeling framework seems versatile and transferrable across different applications..\n- Performance benchmark and analysis on tetrapeptides are comprehensive and provides insights to modeling these peptide systems.\n- I think it is a good idea to model residue offsets relative to the key frames in order to bypass the need to learn sequence-to-structure mapping. MD simulations always start from a seed structure, so I do not think this is a key limitation as mentioned in L#310-312.", "weaknesses": "- Benchmark and evaluation results on tetrapeptides, although comprehensive, are proof-of-concept in nature. It may not be sufficient to demonstrate transferability to general protein systems.\n- Performance on ATLAS (i.e., larger proteins instead of short peptides) does not seem promising. MDGen performance is worse than AlphaFlow in Table 4. I wonder if the main bottleneck is training data quality/availability, or model architecture?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review! To address your questions and concerns:\n\n---\n**Limited to tetrapeptides**\n\nWe have focused on tetrapeptides as model systems in this work for two key reasons:\n* We can run simulations for thousands of systems in order to properly test the generalization abilities of our model.\n* We can build Markov State Models for each system, allowing the careful benchmarking for forward simulation and interpolation capabilities.\n  \nBoth of these aspects are important for the thorough, careful benchmarking of the new capabilities we demonstrate. We opted to prioritize these careful studies to support the core claims of our work, in lieu of expanding its scope to larger and more diverse systems, which we think are best left to future work.\n\n**Comparison with AlphaFlow**\n\nIt is true that the results on ATLAS simulations are worse than AlphaFlow. However,\n* MDGen learns the harder task of providing temporal consistency across structures, which is beyond the capabilities of AlphaFlow. However, our metrics are obtained from the AlphaFlow paper and only assess the ensemble similarity of unordered sets of conformations, favoring their method.\n    * To more concretely demonstrate learned temporal consistency, we now show an MDGen trajectory connecting apo and holo states of adenylate kinase in **Figure 4** in the PDF attached to the global response. \n* AlphaFlow is a much larger model with significant transfer learning from O(10^5) structures via AlphaFold, and has a much better pretrained understanding about likely protein conformations. On the other hand, MDGen is trained from scratch on only O(1000) proteins.\n\n\n\nWith that said, we do not anticipate that this will be the final, definitive architecture for trajectory modeling of protein systems. Rather, our experiments serves as a proof-of-concept that these capabilities _can_ be demonstrated on proteins, even with the extremely limited available data. We anticipate that a combination of expanded datasets and further architectural exploration can improve results on ATLAS in future work.\n\n**Additional structural quality and variation metrics**\n\nWe note that in addition to autocorrelation, our torsion angle distributional metrics and free energy surfaces also assess the quality of structures in the trajectory. With that said, we are happy to report addition metrics as requested. \n\nTo more stringently assess structural quality in MDGen forward simulation rollouts, we compute the distributions of\n* The closest distance between nonbonded atoms\n* Nonbonded energy (Coulomb + Lennard-Jones)\n* Torsional energy\n* Heavy atom bond lengths\n* Radius of gyration\n\nThese distributions are shown and compared to the ground truth in **Figure 1** in the PDF attached to the global response. We find that the vast majority of MDGen structures are of high quality (i.e., clashes are rare) and adhere closely to the ground truth distributions.\n\nTo assess structural variation across time, we report the aligned RMSD between frames spaced at regular intervals, and compare to the same metric computed from the reference simulation. This gives an idea of how fast the structures should be moving. We observe that the trajectories have dynamics that closely resemble the ground truth.\n\n| Interval | MDGen | Reference MD |\n|:-| :-: | :-: |\n| 10 ps | 1.02 A | 0.99 A |\n| 20 ps | 1.19 A | 1.17 A |\n| 50 ps | 1.45 A | 1.44 A |\n| 100 ps | 1.65 A | 1.67 A |\n| 200 ps | 1.86 A | 1.89 A |\n| 500 ps | 2.11 A | 2.18 A |\n| 1000 ps | 2.28 A | 2.37 A |\n\nPlease let us know if you had other metrics in mind; we are happy to analyze the structures further!\n\n**Other questions**\n>L#141, when K>1, how to ensure roto-translation prediction consistency across K key frames and obtain a final g^j?\n\nWe predict all-atom coordinates using each set of key frames (and corresponding predictions) and average these coordinates.\n\n>Table 2, with 100 ns being the ground truth, the non-zero JSD in the last column originates from subsampling the simulation trajectory?\n\nSimilar to the other columns, the MD baselines indicate replicate simulations, i.e., an independent 100 ns simulation.\n\n>Figure 2F. My understanding is that sidechains exhibit faster dynamics while backbone motions are slower. \n\nIt is true that the slower dynamics are typically more interesting; unsurprisingly, they are also harder to simulate and learn. Our method does not perfectly recover these dynamics, but it is the first of its kind to report good results on this task.\n\n>Why is MDGen more effective at sequence recovery than MPNN? More explanation and analysis would be helpful here.\n\nMDGen is trained to use more input information about the peptide, namely the intermediate dynamics of the unmasked residues. Our inverse folding baselines are not able to make use of these partially observed intermediate structures. \n\n>Would it be possible to emulate MD simulation trajectory of the 12 fast folding proteins from Shaw 2009?\n\nThe DESRES trajectories are unusual in that they were simulated at elevated temperatures to invoke unfolding events. In terms of absolute displacement, these are vastly larger motions away from the starting structure than those seen in ATLAS simulations. As such, we do not think our architecture, with its dependence on key frames, would be optimal for modeling such trajectories. Earlier in the project, we considered implementing a distogram-based architecture for the DESRES proteins, but opted to focus on the transferable setting with ATLAS instead.\n\n>It would be nice to see if MDGen could infer a trajectory given an apo/holo pair.\n\nThanks for the suggestion! We have newly trained an interpolation model on ATLAS data and visualize a trajectory between the apo and holo states of adenylate kinase (1AKE / 4AKE) in **Figure 4** of the PDF attached to the global response.\n\n---\n\nWe hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.", "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": "tYSN9REPTg", "reviewer": "Reviewer_b8Jp", "summary": "The paper suggests a flow-based generative framework on molecular trajectories, with various downstream tasks such as forward simulation and transition path sampling. Additionally, the model is trained in a transferable setting, across tetrapeptides.", "questions": "1. Difference between downstream tasks\n\nCould the upsampling task be seen as a superset of interpolation? Since upsampling with two given frames would the same as interpolation.\n\n2. Training on one tetrapeptide\n\nJust curious, though the authors has presented a transferable setting, are there any results when the model is trained for a specific tetrapeptide and tested on downstream tasks?", "limitations": "1. Unconditional generation\n\nAs the authors have mentioned, unconditional generation is impossible since the model relies on key frames.", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Extensive experiments over various downstream tasks\n2. Transferable settings for tetrapeptides", "weaknesses": "1. Experiment baselines\n\nThe baselines of experiments are mostly the Markov State Models. I think it would also be good if there were some comparison between other models, though I understand that many prior works targeted Alanine dipeptide not tetrapeptides.\n- Forward simulation: ITO$^{[1]}$, Timewarp$^{[2]}$\n\n- Interpolation (Transition path sampling): PIPS$^{[3]}$\n\n2. (Minor) Necessity of additional tasks\n\nThe necessity of additional tasks relatively seems weak compared to tasks such as forward simulation, TPS, specially the inpainting design. Rather than additional tasks, ablations for stability might be a better? One can obviously see that scaling to long trajectories shows the stability against the time scale, and protein simulation shows the stability against space complexity.\n\n**Minor typos, suggestions**\n\n- Definition of S-MPNN only exists in the Appendix. It would great to point out that more details are in the appendix, in the first paragraph of section 4.4\n- Figure 6 is not referenced in the main paper, only the appendix\n- Figure 2F, reference of that blue indicates the side chains and orange indicates the backbones seems to be missing\n\n[1] Implicit transfer operator learning: Multiple time-resolution surrogates for molecular dynamics, NIPS 2023\n\n[2] Timewarp: transferable acceleration of molecular dynamics by learning time-coarsened dynamics, NIPS 2023\n\n[3] Stochastic Optimal Control for Collective Variable Free Sampling of Molecular Transition Paths, NIPS 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review! To address your questions and concerns:\n\n---\n\n**Experimental baselines**\n\nWe now provide new results comparing our work to Timewarp and ITO. Emphatically, these comparisons are limited to the forward simulation task as Timewarp and ITO are not capable of solving the other tasks. \n\n* For **Timewarp**, we use the 4AA model with weights from the authors and sample 100 ns trajectories by running 2000 inference steps with timestep 50 ps. We do not use MH acceptance steps as the authors found exploration of the energy landscape to be much more effective without them.\n* For **ITO**, transferable models across tetrapeptides are not available. We, therefore, train ITO ourselves on our tetrapeptide dataset with timesteps of 500 ps. We then run 200 inference steps to sample 100 ns trajectories.\n\nFor both methods, we observe that trajectories are unstable without further intervention (see **Figure 2** of the PDF attached to the global response). To bolster the baselines, we run OpenMM relaxation steps between each timestep to proceed with further analysis. We note that the Timewarp authors, in lieu of relaxation, rejected steps with an energy increase of 300 kJ / mol; however, we found that this strategy would reject the majority of proposed steps on generic tetrapeptides.\n\nIn **Figure 3** of the PDF attached to the global response, we visualize the free energy surfaces and torsion angle distributions for several peptides from Timewarp and ITO compared with MDGen. Our model obtains much better results across the board.\n\nQuantitatively, we obtain the following Jensen-Shannon divergences from the forward simulation rollouts (c.f. Table 2):\n\n|C.V.| MDGen | Timewarp | ITO |\n|:-| :-: |  :-: | :-: | \n|Torsions (bb)| **0.130** | 0.325 | 0.564 |\nTorsions (sc) | **0.093** | 0.427 | 0.462 |\nTorsions (all) | **0.109** | 0.383 | 0.505 |\nTICA-0 | **0.230** | 0.265 | 0.538 |\nTICA-0,1 joint | **0.316** | 0.419 | 0.756 |\nMSM states | 0.235 | **0.222** | 0.414 |\nRuntime (s)| **60** | 599 | 2083 |\n\nTo evaluate the dynamics, we also compare the Pearson correlation of predicted torsional relaxation times (c.f. Figure 2F). We could not compute these quantities for ITO as the vast majority of torsions did not decorrelate within the ITO rollouts.\n\n|C.V.| MDGen | Timewarp | ITO |\n|:-| :-: |  :-: | :-: |\nTorsion relaxation (bb)| **0.55** | 0.04 | ---|\nTorsion relaxation (sc) | **0.97** |  0.77 | ---|\n\n\nAltogether, the results confirm that MDGen is more successful and efficient at modeling thermodynamics and kinetics than Timewarp or ITO.\n\n**PIPS** We note that PIPS is a method for steering simulations towards a desired target state, whereas MDGen aims to learn from existing unsteered simulations in a transferable manner to sample transitions of unseen systems. Due to the difference in tasks, a direct comparison with PIPS is difficult to formulate.\n\n**Additional tasks vs ablations for stability**\n\nA central point of our paper is to highlight the novel capabilities afforded by our trajectory modeling framework. We chose several tasks out of reach of previous methods, and hence less studied in ML papers, but all of which we believe are scientifically meaningful. With that said, we acknowledge that different tasks will appeal to different readers, and appreciate your input about their relative importance. \n\nWe are happy to consider suggestions for further experiments for stability (time permitting) beyond the existing experiments you referenced.\n\n\n**Difference between the downstream tasks**\n\nYes, upsampling could be seen as a superset of interpolation in terms of the conditioning input. However, their conceptual aims are somewhat different: In the interpolation setting, we choose endpoints in different (often distant) macrostates of the energy landscape to study how they are connected. On the other hand, in upsampling the conditioning frames are chosen based on timestep, irrespective of whether they exhibit interesting transitions.\n\n\n**Training on one tetrapeptide**\n\nApart from the Hyena experiment (Section 4.4), we opted to focus on the transferable setting as we believe this is the arena in which ML-based MD emulators will be useful in practice. Additionally, since our model generates 10 ns at a time, we do not anticipate that learning from a single 100 ns trajectory would be meaningful, as this amounts to only 10 statistically independent training examples. \n\n**Unconditional generation**\n\nWhile it is true that the use of key frames impacts the ability to do unconditional generation, we nevertheless opted for this design choice as unconditional generation of trajectories is (to our knowledge) not a problem of scientific interest. On the other hand, the key frames allow the conditional modeling problems to be consideribly simplified and clarified, leading to the strong results shown in our experiments. Hence, the use of key frames could instead be considered a technical insight of our model that enables it to contribute to the development of real-world, impactful scientific problems, at the cost of problems of lesser interest.\n\n**Minor typos, suggestions**\n\nThank you for noting these; we will incorporate them in the revision!\n\n---\n\nWe hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.", "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:
# Overall Response

We thank all reviewers for their time taken in providing constructive feedback!

In addition to the individual responses, we also provide new **figures and visualizations in the PDF file** attached to this global response.

* In Figure 1, we compare the distributions of additional **structural and energy metrics** (e.g., clashes) from MDGen rollouts and from reference MD simulations, as suggested by **Reviewers QAiG** and **hVZ7**. We find that the vast majority of MDGen structures are of high quality and closely resemble the ground truth distribution on these metrics.
* In Figures 2 and 3, we provide **comparisons with Timewarp and ITO** as requested by **Reviewers zaG3** and **b8Jp**. In Figure 2, we show that Timewarp and ITO suffer from unstable inference rollouts out-of-the-box on generic tetrapeptides. However, we correct for these instabilities ourselves to obtain a meaningful comparison. In Figure 3, we visualize free energy surfaces and torsion angle distributions obtained from Timewarp and ITO compared with MDGen. We find that MDGen obtains better results across the board.
* In Figure 4, we visualize an **interpolation path between apo / holo states** of adenylate kinase (1AKE / 4AKE) as suggested by **Reviewer hVZ7**, obtained from a newly trained ATLAS interpolation model.

Meta-review: The paper presents a new approach to a timely topic -- modeling molecular dynamics -- with two accepted papers in NeurIPS last year. There are several contributions, lending ideas from CV and NLP (such as masking and impainting) to model a joint distribution of the simulation trajectory. As such, these new training strategies allow for the use of more data than previous approaches. The tasks the authors explore are in line with standard benchmarks. New prediction tasks are introduced, mostly of limited interest and practical use (e.g. trajectory upsampling). The major limitation and strength of the approach are its main innovations: 1) the molecular representation (tokenization), which locks in the models to be proteins or peptides, with no clear strategy to generalize to molecules more broadly, and 2) the trajectory representation, requiring end users to make hard decisions at training time, which may limit inference time usefulness. The underlying learning objective is not well justified from a physical perspective: e.g. a trajectory is learned as a joint distribution suggesting a time-dependence of the generator. There are settings where this is the case, however, in all the systems explored the paper the generating process is homogeneous and Markovian, and therefore at odds with the model and training setup. Nevertheless, the performance appears good in certain cases (small peptides), but poor on larger scale application, e.g. ATLAS protein dataset.
Decision: Accept (poster)