Timewarp: Transferable Acceleration of Molecular Dynamics by Learning Time-Coarsened Dynamics

Summary

This paper addresses high-throughput MD simulations. The proposed sampling algorithm simulates the MD of 10^5-10^6 fs timestep by a pre-trained model of a normalizing flow. The normalizing flow can be used as a proposal distribution in an MCMC method. The experimental results show the pre-trained model trained on peptides with two or four amino acids generalized to unseen combinations of amino acids. The proposed method is faster than naive MD in many cases of two amino acids.

Strengths

* The paper is well-organized and easy to follow. * The proposed method can be integrated with a popular sampling method, MCMC. * The batched sampling method compensates for the low acceptance ratio of the proposed model. * The experimental results of the proposed exploration algorithm are also provided. * No speed-up for no slow transitions is reasonable. * The limitations are clearly shown in the experimental results.

Weaknesses

* No comparison with conventional MD-acceleration methods such as metadynamics or parallel tempering. I would like to know how the proposed method is competitive with manually-tuned those methods. * The proposed method needs to be evaluated on application-oriented metrics such as the reproducibility of some physical properties. * Although the proposed method largely improved the simulation speed, the MH acceptance rates for the model are low, so the diffeomorphism function seems to need to be learned better. * According to Figure 3, the proposed method cannot reproduce the transition states, i.e., the hill of MD's energy surfaces, so the method may not be applicable to chemical reactions.

Questions

* Please provide the computational environment for the baseline MD simulation, amber-14. Is the wall clock comparison fair enough in terms of their implementation? * Why is the proposed method slower on 4AA? Is this because of O(N^2) computational complexity of the attention mechanism or the size of the neural networks?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

The limitations are clearly stated in Section 7.

Summary

In this paper the authors introduce Timewarp, a method that relies on a normalizing flow as a proposal distribution for enhanced sampling of the Boltzmann distribution. The key results include an acceleration in the MD time step from ~1 fs to 10^5 fs and transferability between peptide systems.

Strengths

The core demonstration of generalization to unseen peptides in an all-atom representation is an important contribution. The method is clearly presented and the value of accelerated sampling is also clear.

Weaknesses

The method seems to break down already for 4-mer peptide systems. The limitations in expressivity of normalizing flows shown throughout the literature cast doubt on the scalability of the method to practical applications. Performance benchmarks with respect to classical MD approaches are difficult to evaluate, due to choices in force field, integrator, etc.

Questions

Can the authors expand on the choice to treat velocity as a non-physical auxiliary variable? It is clear that it makes a permutation equivariant coupling layer easier to construct, but is there no information loss that might contribute to the low MH acceptance rate? How many heads and how were length scales determined for each head in the kernel self-attention? Does ablating this tell you anything about the importance of including long-range interactions, or is the locality assumption implicit in the kernel self-attention sufficient for capturing the dynamics for the peptides considered? Can the authors provide additional references and discussion (perhaps in the appendix) of the evaluation metric presented in Eq 18? What ESS/s values are reported when comparing different classical MD algorithms? Since Timewarp is evaluated on small peptide systems, conformer generation with a semi-empirical method like xTB might also be an appropriate evaluation. Can the authors discuss their method in the context of such approaches? Are there any concrete proposals for future work that would address the current limitations in MH acceptance rate, system size, and the troubling trends in scaling to larger systems? Why do the authors use amber-14 in openMM rather than a more modern force field? The caption in Fig 4 does not match the figure labels. In Fig 3c the red dot indicating initial state is difficult to see.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes, although further discussion of possible avenues for addressing these limitations would be welcome.

Summary

This paper presents a Boltzmann generator for small molecular systems in implicit solvent. The generator works demonstrates transferability to test systems of the same class (dipeptides or tetrapeptides). It is the first Boltzmann generator that operates in Cartesian coordinates and demonstrates transferability. The authors estimate that for these very small molecular systems in their study, the sampling via a Boltzmann generator could provide wall-time acceleration compared to vanilla MD simulation.

Strengths

This paper is the first demonstration of a Boltzmann generator in Cartesian coordinates that generalizes to untrained systems. This is certainly commendable, despite the existence of previous generators for small molecules that relied on internal coordinates. The method is presented in a clear fashion and the authors plan to release their codes and data upon acceptance. The experiments are sensible and show that the slow modes of the MD simulations are covered with a similar distribution by the sampler. This paper is adding to a series of developments in the line of Boltzmann generators, an area that shows great promise, even though it is still at its infancy with respect to practical applications in drug discovery. As such, I think that it is a significant contribution to this conference.

Weaknesses

The correct operation of this generator requires a Metropolis Hastings step, which will limit the application of the current technology to only very small systems. This is not necessarily a problem with all Boltzmann generators, but this limitation is very clearly demonstrated in the work and in leads to a reduction in the acceptance rate as the number of particles doubled. The authors don't discuss this possible limitation of the current work in their paper, nor suggest ways to think about scaling to systems of practical interest. The work is only dealing with molecules in implicit solvent. These implicit solvent simulations are not of practical use in drug discovery, so it would be useful to have a discussion of possible generalizations towards explicit simulations (either generating the full environment, or at least learning from such simulations and generating whatever is feasible). It is hard to judge reproducibility without the availability of the code/data; I hope this will be alleviated once the paper is published.

Questions

The work is limited to molecules in implicit solvent. How difficult would it be scale these methods to also generate the full environment of water, ions, lipids and small molecules in typical molecular dynamics simulations in the context of drug discovery? Perhaps in a more practical fashion, are there any limitations in the current architecture Did the authors try to train the model on explicit solvent simulations and use only the non-water/non-ion coordinates in their systems? The generation of the training trajectories does not seem particularly time consuming compared to the rest of the compute effort that went into the project and the possible wall-time improvements might be even larger (though of course this depends on the software and hardware that is employed). The very limited number of total atoms in the current models would suggest that existing enhanced sampling methods for MD simulations (including those that include kinetic considerations) would also work for these small systems; however, such sampling techniques might be trickier to optimize in explicit water than in implicit water. (Even without any sophisticated techniques, I would guess that the implicit-solvent dipeptides and tetrapeptides would reduce their autocorrelation much more quickly by sampling them at ridiculously high temperatures (or fluctuating temperatures as in simulated tempering), and could then be projected back to the baseline with an MH criterion.) Why did the authors not mention the Torsional Diffusion paper in their introduction? Wasn't that work also a generalizable Boltzmann generator (using internal coordinates), and probably the first of this kind?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

4 excellent

Limitations

No potential negative societal impact

Summary

In this work, the authors present a method for speeding up MD simulations. The speeding up of MD simulations is needed to be able to simulate processes that occur over longer timespans (eg. protein folding). The method proposed by the authors, Timewarp, does so by learning the transition probability of MD simulations over timesteps beyond those normally used. The transition probability is learned using a normalising flow based on a Transformer architecture. Transition steps sampled using the normalising flow are pushed through a Metropolis-Hasting acceptance step to assure asymptotic unbiased samples. Additionally, the authors propose to run the normalising flow without the MH step for faster sampling. The normalising flow is learned using trajectories of multiple structures using traditional MD simulation. By using multiple structures at the same time, the authors aim to make the method generalize to unseen structures. This is one of their major contributions. The work experimentally shows the success of the method on 3 different sizes of molecules; 1) Alanine Dipeptide, 2) 2 amino acid peptides and 3) 4 amino acid peptides. The experiments show that -Timewarp closely matches the samples from MD -Timewarp exploration (without MH step) provides estimates of the free energy close to MD, but less accurate then MCMC (with MH step).

Strengths

NOTE: This is an updated review from a review round of an earlier conference. The review is updated to reflect the changes made by. **Originality:** The paper discusses a very interesting new intermediate problem in ML for Chem/Bio that lies in between the two more often studied problems of Force Field learning and Conformer generation. I've commented on this below. **Quality:** The paper provides a good introduction to the problem at hand and as such sufficiently motivates the work. The proposed solution is interesting and makes good use of existing methods for similar problems without overcomplicating it. Experiments are in depth and well executed. **Clarity:** The paper is well written and structured. Figures are of high quality. **Significance:** The paper has the potential to be very significant. The problem introduced is novel and important.

Weaknesses

- Unfortunately, there's no code available and the datasets are not open sourced. - Related work could be improved with more references to work from the computational chemistry literature on enhanced sampling(ie. topic such as Transition Path Sampling and Bias Potential methods).

Questions

As I commented above, from my understanding, the problem discussed and addressed in this paper lies somewhere in between force field learning and conformer generation. The transition step length tau in this case determines which of the two the problem is closer to. If tau is very small, the problem at hand is similar to force field learning. With tau very large, the conditional distribution becomes similar to learning the Boltzmann distribution. Considering this, I'm surprised the work does not contain a study of the influence of tau on the speedup seen by the method over MD. Could the authors comment on this? For example, why was the specific value for tau used in the paper chosen?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors have sufficiently addressed the limitations.

Thank you for the response

I acknowledge having read the response by the authors and I remain positive about this work. I agree with the view that excluding the Metropolis Hastings step might be a feasible way forward. In addition, I wonder if perhaps there is more to learn/mimic from existing methods for enhanced sampling. Finally, I suspect that the work on the inclusion of explicit solvent, although technically feasible, is qualitatively different from the present work, as it would have to deal with changes in the loss functions to account for permutations and proper hydrogen bonding. I will look forward to the future work on this subject.

Score increase

Thank you for your response. With the proposed addition of a discussion on enhanced sampling and a further discussion on the role of tau in mind, I will increase my score.