DynamicsDiffusion: Generating and Rare Event Sampling of Molecular Dynamic Trajectories Using Diffusion Models

See weaknesses.

Please find my detailed list of comments, ordered by section, below.

Where needed, I indicated that something might be a breaking issue. These issue will have to be either clarified or rectified before I will consider increasing my vote to accept.

Introduction:

1. This also pertains to the remainder of the paper, please make sure the citations are correctly placed between brackets unless when used as part of the narrative.
2. The first paragraph of the introduction state "However, their potential in the study of the dynamics of molecular system is still unexplored" when discussing the use of Diffusion models for generating ensembles of molecular conformations. This sentiment is reflected a number of more times throughout the introduction and the remainder of the paper, as such, I would expect the experiments to focus on properties related to the dynamics of the system that can not be determined from the ensembles provided by standard Diffusion models. However, given that the experimental section mostly focusses on evaluating Free Energy estimates (which can be derived from the equilibrium distribution), this does not seem to be the case. I would like the authors view on this observation.
3. Similarly, this paragraph also states "an accurate description is essential to understand the biological functions of proteins and other molecules.". I urge the authors to make these "biological functions" concrete. Are they talking about binding affinity, transition states or some other properties? This distinction is important to be able to judge the contribution of the proposed method.
4. Related to my prior statement, the last sentence of the first paragraph states "otherwise rare but important events". I would appreciate a clear statement as to why these events are important and would consequently expect the experimental evaluation to focus on correctly identifying them.
5. The authors state that "standard classical atomistic MD must use an integration time step of 2fs". While this is an often used value, there is no specific requirement for MD to use a 2fs timestep. The word "must" should be reconsidered.
6. Paragraph 3, 4, and 5 discuss different approaches for sampling rare events, but somewhat lacks organisation. This is largely due to not mentioning general themes or creating a clear topology of the different methods. My understanding is that paragraph 3 discusses ML based methods for reducing the computational cost of determining the molecular forces, paragraph 4 discusses enhanced sampling methods for enforcing transitions, and paragraph 5 discusses direct sampling from the equilibrium distribution using ML frameworks. These overarching themes should be made more clear. In general, I think the paper would benefit from a dedicated related work section that goes more in depth for these 3 paragraphs. The introduction would in that case only need to focus on the most important methods.
7. Paragraph 3 currently has to narrow of a focus when discussing neural force fields and coarse graining. These are large research fields and can't be done justice with only specific citations. For example, neural force fields come in many other flavours then just the one based on GNNs. To resolve this issue, I would suggest proving citations to survey papers. For example, for machine learning force field [1] is well suited, and for coarse graining [2] provides a clear overview.
8. For paragraph 4, the discussion of large energy barriers should be part of the earlier discussion on timescales. It should be made clear that "rare events" does not refer to events that happen over a long period in time, but instead to events that happen with very low probability.
9. For paragraph 4, based on the authors statement to focus on system dynamics, TPS is in my opinion an extremely related topic and requires some additional discussion. For example, an additional discussion of ML based alternatives/extensions here would be appropriate [3, 4, 5, 6]
10. For paragraph 6, the authors use the term "enhanced sampling". From my understanding, this term is often used for methods that augment MD simulations (often by introducing a bias potential or other driving force). This does not seem to be the case for the suggested method.
11. The list of capabilities of the proposed method states "Generate trajectories that are conditioned to a global parameter like temperature". Based on my own understanding and the presented results, I am not confident that other global parameters, for example friction, will be possible.
12. Similarly the list of capabilities of the proposed method states "Generate an ensemble of (reactive) trajectories by partially noising and denoising them". I did not find any experimental validation of this. Additionally, a statement such as this would need to be validated to show that the sampling converges to the right distribution over trajectories.

Diffusion Generative Modelling for Trajectories

1. On first reading it is very hard to understand the description of the dataset. And after reading it a couple of times, it is still unclear to me if "trajectory snippet" refers to a single molecular configuration or a short sequence of configurations. Based on the first paragraph, I would assume that the dataset contains trajectories. $p(x)$ would thus be the distribution over trajectories of the system. Ie. $x$ here refers to a trajectory of arbitrary length? If so, it would be beneficial if the authors could include a more formal definition of this probability in terms of the marginal distirbution $p(R_0)$ and the transition distribution $p(R_i | R_{i-1})$. However, based on the second paragraph, i get the impression that $x_i$ is a single configuration, and instead $\boldsymbol{x}_i$ is a short trajectory. This confusion is a breaking issue for me and should be clarified by the authors.
2. The third paragraph, discussing the complete data array is hard to understand. Do the authors intend to clarify in this paragraph that various length of trajectories can be used?
3. In equation one, $\nabla_x$ and $p(x)$ should respectively be $\nabla_{x_i}$ and $p(x_i)$. To prevent further confusion with indexing, I would also suggest to use t instead of i for the updates.
4. Below the equation the authors state "$\epsilon\rightarrow 0$ we sample one $x_k$ from $p(x)$". I'm unsure what this statement refers to. If $\epsilon\rightarrow 0$ we have no update and thus the sample remains constant. This is independent on whether or not $K\rightarrow \infty$.
5. Pertaining the same sentence, and related to my earlier point, is $p(x)$ here the distribution over samples, or the distribution over trajectories?
6. Regarding the last sentence of this same paragraph, the authors refer to a section in the appendix by name here. This should replaced by a forward reference. This happens on a number of occasions throughout the paper.
7. For equation 2, the notation $R^{(i)}_{\nu, t=0}$ is unclear. Instead of adding the $t=0$ in the subscript would $R^{(i)}_{0}$ not suffice?
8. I'm unable to follow the authors discussion of the conservative forces and the role the modelling of the forces as the gradient of an energy function plays in this. Especially the claim that UNets do not satsify this property should be substantiated.
9. Regarding the last sentence of this paragraph, the result presented in the ablation study shows that this difference is not significant.
10. My understanding is that the authors suggest to use a model architecture that predicts an energy independently for each sample, then sums thism and then takes the gradient to obtain the final force, which is used as the predicted noise. I would like to note that this formalism suggests that the individual samples in the trajectories are independent from each other. ie. a change in one sample does not influence the probability of another sample in the same trajectory. Given that we are looking at trajectories, this is not true. In other words, the inductive bias introduced by the authors does not hold for the modelling of trajectories. This is a breaking issue and I would like to get the authors view on this observation.

Experiments

1. For the discussion of the datasets obtained, it is unclear how the long trajectory is divided into shorter trajectories for the training set.
2. Regarding the referencing to figures, please correct this to use the correct referencing format where it reads as either "fig. x" of "figure x" instead of simply the number.
3. Regarding the discussion of the Brownian Motion in a Benchmark potential correctly sampling the wells, the conditional sampling changing the spread of the samples, and correctly identifying the Boltzmann distribution, I'm unsure if these results are informative. As it stands, given that the training data also correctly covers the Boltzmann distribution, it is not surprising that the proposed diffusion method can correctly sample this space. The issue with rare event sampling is that the temperature of the simulation is too low to effectively sample transitions and as a result, the ratio in Boltzmann distribution between different states is incorrect. The temperature for generating the dataset however seem to be high enough to already get a correct estimate. This is a breaking issue and I would appreciate the authors comments on this.
4. Additionally, an important aspect of the papers contribution is that they extend diffusion models for molecular conformations to the study of dynamics. As of now, the presented evaluation however has focussed on properties of molecular systems that can also be studied without taking into account the dynamics. Standard diffusions models should, in theory, be able to sample the Boltzmann distribution and provide accurate estimates of the FES. This is an breaking issue and I would like to see a more in depth study of the dynamical properties that can not be studied using standard diffusion models.
5. In figure 4, the description of (b1) and (b2) mention Free Energy profiles but the figure shows equilibrium distributions instead.
6. For 3.2 the authors state, "a (relatively) rare conformational transition occurs around the dihedral angle $\psi$". However, the energy barrier between these two states is considered to be relative low and can thus be sampled in relative short periods at low temperatures. As such, I'm unsure if this transition is suitable for evaluating the method. I would suggest to instead consider transition along the $\phi$ axis.
7. Again, given that the training for Alanine Dipeptie already correctly determines the FES I'm unsure what the added benefit is from modelling the transitions using a diffusion model. The FES can also be reconstructed using a sample based diffusion model.
8. The last sentence of section 3.2 reads "pointing to the generalisation capabilities of the method" regarding the observation of a rare transition along the $\psi$ angle. I respectfully disagree with the authors here. This transition is highly physically unrealistic as it flows over a region of high energy.
9. Regarding the ablation study, based on the presented error bars these results are all not significant. This is a breaking issue and I would like the authors comments on how this affects their results

Conclusion

1. The authors state "which can serve as a useful surrogate model of the physics-based simulator, due to its enhanced sampling capabilities". Aside from possible speedup, the presented results do not currently show any enhanced sampling capabilities. As such, I think it would be appropriate for the authors to rephrase this claim.
2. The authors discuss the extension of their work to generalise accross single molecules. I believe this to be a very good suggestion and hope that the authors continue their research in this direction.

References

[1] https://pubs.acs.org/doi/10.1021/acs.chemrev.0c01111

[2] https://pubs.acs.org/doi/10.1021/acscentsci.8b00913

[3] Differentiable Simulations for Enhanced Sampling of Rare Events

[4] Learning Free Energy Pathways through Reinforcement Learning of Adaptive Steered Molecular Dynamics

[5] https://arxiv.org/abs/2207.02149

[6] https://iopscience.iop.org/article/10.1088/2632-2153/acf55c (Only recently published but possibly relevant for the authors)

1. As shown by the authors, taking Jacobian of the predicted energy term can be costly than predicting the noise (Table A.1). However, the authors only benchmarked the error of models (predicting energy or force) in the 2D double-well system. It is not convincing enough to me that energy prediction+Jacobian is better than force prediction in terms of cost/accuracy tradeoff for larger systems.

2. In Figure 4.c, the trajectories generated by inpainting methods seems to capture the molecular transition dynamics between the two conditioned points. However, the left condition (green dot) is not connected to other states. Is that a plotting mistake?

3. How is the self-attention implemented along with the UNet architecture? The detail is missing in the manuscript.

4. The model architecture and problem formulation enables rare-event sampling with a fixed number of timestep ($\tau$) when the initial and end conditions are given. Based on my understand, the method can garuantee the transition states between the two condition to be sampled. However, the probability of those rare-events cannot be obtained using the model.

1. What is the math expression of the energy formulation?
2. What is the performance difference between your proposed 2 sampling methods and the baseline method?
3. Which part is related to Figure 2? The reviewer can guess but wants it to be precise.
4. Wrong reference table for Page 8. The author cited Table 9 for the ablation study of self-attention. But Table 9 is about Jacobian. Missing references on page 18.
5. Which part of the result shows it can promote symmetry? Miscellaneous typos need to be further proofread.

N/A