Thank you for the review! To address your questions and concerns:

Experimental baselines

We 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.

• 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.
• 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.

For 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.

In 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.

Quantitatively, we obtain the following Jensen-Shannon divergences from the forward simulation rollouts (c.f. Table 2):

To 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.

Altogether, the results confirm that MDGen is more successful and efficient at modeling thermodynamics and kinetics than Timewarp or ITO.

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.

Additional tasks vs ablations for stability

A 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.

We are happy to consider suggestions for further experiments for stability (time permitting) beyond the existing experiments you referenced.

Difference between the downstream tasks

Yes, 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.

Training on one tetrapeptide

Apart 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.

Unconditional generation

While 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.

Minor typos, suggestions

Thank you for noting these; we will incorporate them in the revision!

We hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.

Thank you for the review! To address your questions and concerns:

Limited to tetrapeptides

We have focused on tetrapeptides as model systems in this work for two key reasons:

• We can run simulations for thousands of systems in order to properly test the generalization abilities of our model.
• We can build Markov State Models for each system, allowing the careful benchmarking for forward simulation and interpolation capabilities.

Both 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.

Comparison with AlphaFlow

It is true that the results on ATLAS simulations are worse than AlphaFlow. However,

• 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.
  • 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.
• 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.

With 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.

Additional structural quality and variation metrics

We 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.

To more stringently assess structural quality in MDGen forward simulation rollouts, we compute the distributions of

• The closest distance between 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 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.

To 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.

Please let us know if you had other metrics in mind; we are happy to analyze the structures further!

Other questions

L#141, when K>1, how to ensure roto-translation prediction consistency across K key frames and obtain a final g^j?

We predict all-atom coordinates using each set of key frames (and corresponding predictions) and average these coordinates.

Table 2, with 100 ns being the ground truth, the non-zero JSD in the last column originates from subsampling the simulation trajectory?

Similar to the other columns, the MD baselines indicate replicate simulations, i.e., an independent 100 ns simulation.

Figure 2F. My understanding is that sidechains exhibit faster dynamics while backbone motions are slower.

It 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.

Why is MDGen more effective at sequence recovery than MPNN? More explanation and analysis would be helpful here.

MDGen 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.

Would it be possible to emulate MD simulation trajectory of the 12 fast folding proteins from Shaw 2009?

The 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.

It would be nice to see if MDGen could infer a trajectory given an apo/holo pair.

Thanks 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.

We hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.

Thank you for the review! To address your questions and concerns:

Limited ML technical contribution

In 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.

We 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:

• The reduced representation of residue all-atom coordinates into $\mathbb{R}^{21}$.
• The use of key frames to obtain SE(3)-invariant tokens from molecular trajectories.
• 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).

Indeed, 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.

Additionally, 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.

Comparison with Timewarp and ITO

We 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.

• 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.
• 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.

For 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.

In 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.

Quantitatively, we obtain the following Jensen-Shannon divergences from the forward simulation rollouts (c.f. Table 2):

To 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.

Altogether, the results confirm that MDGen is more successful and efficient at modeling thermodynamics and kinetics than Timewarp or ITO.

We hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.

Thank you for the review! To address your questions and concerns:

Further discussion on related works

Thanks 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:

• 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.
• DiffMD adds noise to molecular structures and denoises them to generate a future structure. Repeating the process autoregressively yields a trajectory.

When 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.

Lack of ablation studies

We now provide new experimental results ablating components of the method. The ablations are as follows:

• No IPA embedding: the model is not provided IPA embeddings of the key frames and amounts to a vanilla Scalable Interpolant Transformer
• 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$
• No frames/torsions: the model operates directly over all-atom coordinates rather than the reduced representation in $\mathbb{R}^{21}$

We 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.

We hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.

Thank you for your clarifications. Most of my concerns have been addressed. I choose to keep my positive rating.

Thank you for the review! To address your questions and concerns:

Complexity and accessibility

We 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.,

• using a simple stochastic interpolants framework rather than diffusion
• using a vanilla unmodified transformer from previous work rather than complex frame- or pair-based architectures
• generating invariant tokens over Euclidean space rather than diffusing over the space of frames and tokens.

To promote reproducibility, we have

• provided pseudocode for our architecture in Appendix A
• provided experimental details sufficient for replications and conceptual justifications for design choices and metrics in Appendix B
• will provide full source code in the revision

While 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.

Reliance on key frames

While 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.

Larger and more diverse systems

We have focused on tetrapeptides as model systems in this work for two key reasons:

• We can run simulations for thousands of systems in order to properly test the generalization abilities of our model.
• We can build Markov State Models for each system, allowing the careful benchmarking for forward simulation and interpolation capabilities.

Both 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.

Conditioning

We 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.

Data availability and quality

We 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.

Evaluation metrics

We have designed and implemented thorough and principled metrics that assess the ability of our model to produce trajectories:

• With good distributional properties over structures, assessed by the Jensen-Shannon divergences and Markov state occupancies
• With good dynamical properties over motions, assessed by the various transition path metrics, the decorrelation times, flux matrix correlations, autocorrelation functions, and dynamical content.

While 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.

Computational resources

We report our training time for each model as follows, measured on A6000 GPUs:

• Forward simulation: 412 GPU-hrs
• Interpolation: 272 GPU-hrs
• Upsampling: 292 GPU-hrs

We 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.

Emphatically, 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.

We hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.

Thank you for the review! To address your questions and concerns:

Parts of Section 3 hard to follow

Due 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.

Suitability for proteins and non-AA molecules

We 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.

With that said, we have focused on tetrapeptides as model systems in this work for two key reasons:

• We can run simulations for thousands of systems in order to properly test the generalization abilities of our model.
• We can build Markov State Models for each system, allowing the careful benchmarking for forward simulation and interpolation capabilities.

Both 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.

Clashes in the generated structures

To assess the frequency of clashes or high-energy structures in MDGen forward simulation rollouts, we compute the distributions of

• The closest distance between 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 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.

Representations of other methods

The peptide representations of the other methods are summarized as follows:

• 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.
• AlphaFlow emits residue frames and torsion angles, an equivalent representation to ours.
• The inpainting baselines are inverse folding models and do not emit trajectories of any kind.
• Ground truth molecular dynamics provides all-atom trajectories.

Minor points

Thank you for noting these; we will incorporate them in the revision.

We hope the new discussion and results address your concerns! Please let us know if there are further opportunities to improve the score.