Force-Guided Bridge Matching for Full-Atom Time-Coarsened Dynamics of Peptides

Authors: Ziyang Yu, Wenbing Huang, Yang Liu
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: molecular dynamics, force-guided bridge matching, graph neural network

Abstract:
Molecular Dynamics (MD) is crucial in various fields such as materials science, chemistry, and pharmacology to name a few. Conventional MD software struggles with the balance between time cost and prediction accuracy, which restricts its wider application. Recently, data-driven approaches based on deep generative models have been devised for time-coarsened dynamics, which aim at learning dynamics of diverse molecular systems over a long timestep, enjoying both universality and efficiency. Nevertheless, most current methods are designed solely to learn from the data distribution regardless of the underlying Boltzmann distribution, and the physics priors such as energies and forces are constantly overlooked. In this work, we propose a conditional generative model called Force-guided Bridge Matching (FBM), which learns full-atom time-coarsened dynamics and targets the Boltzmann-constrained distribution. With the guidance of our delicately-designed intermediate force field, FBM leverages favourable physics priors into the generation process, giving rise to enhanced simulations. Experiments on two datasets consisting of peptides verify our superiority in terms of comprehensive metrics and demonstrate transferability to unseen systems.

Decision: reject
Reviewer scores: 5, 8, 8

Reviews and discussion:
<thread>
creation date:
1730105335673

thread_type:
official_review

<thread_object>
creator: reviewer_r7tw

<content>
summary:
In this paper the authors propose a model to produce relevant MD  trajectories from an initial structure, using bridge diffusion with a possibility to ensure that the recovered conformation statistic obeys a  Boltzmann distribution rather than blindly following the learned data distribution. They accomplish their goal by constraining the form of the conformation distribution to be the learned data distribution with each conformation being weighted by a learnable Boltzmann factor. So the constraints appeared from the beginning and are present in the trained model from the get-go: no need at prediction time to go through expensive re-weighting procedure, as it is part of the model. Hence their model relies on first learning the reverse and forward process that would generate the initial and final experimental conformations (conditional to those), then use this newly learned score function and use it within a dedicated network to calculate the associated Boltzmann weight (i.e. the conformation energy). At prediction time one can either just perform the simple diffusion, or decide to add the correction term to the simple diffusion-associated force field, which turns out to “simply” be the gradient of the energy term in the learned Boltzmann re-weighting coefficient: hence the name.

strengths:
The paper offers a rather intuitive and straightforward way to add some physics-derived constraints to the trajectory produced at prediction time. Moreover, their model is state of the art when taking into account a rather diverse set of evaluation metrics: from comparing model output collective variables at the distribution level, in terms of validity through checking bond clashes and breaking, or finally in term structure flexibility compared to ground truth. The simplicity of adding or not the correction term in a tunable way through a guidance strength hyperparameter allows for a direct ablation study which shows the usefulness of the correction (guidance) term.
The care brought into splitting the dataset (clustering and splitting by  cluster), indeed allows them to claim generalizability, even though I  didn’t really understood why AD seems to be treated a bit separately  from the peptide in the test set.
Last but not least, it also proposes a robustness/explanatory study by adding the results of their hyperparameter tunning and some effects of those hyperparameters.

weaknesses:
The paper is pretty strong and its only rather minor weakness has to do with easiness to have a direct interpretation/overall view of the metrics for the samples presented. Right now the metrics show their claims (their model is better) which is great, but they, I  believe, are hard to link to the model's performance and usefulness. 
For example, I think Figure 6 is nice because it shows a direct comparison with ground truth in a more obvious way than other figures. Hence the few following questions and propositions of improvement described in the next section will be about that.

questions:
Overall all those figures and metrics are complementary and I believe for a few samples to have all of those would help a lot. Indeed we have the different JS values for the full test set but having them also for the samples presented, as well as having directly the distribution comparison to the ground truth as in Figure 6 would be super helpful. I  believe it can be hard to understand what a JS value means except that it is bigger for the output of one model compared to the other. 
It could be also nice to have more of those figures for your best-picked trajectories but also for those that didn’t work as well: again just to have a feel about what could be driving those numbers.
Moreover the study of hyperparameters effect (table 6) is a great add to show the robustness of the results. I believe it could be used to show the effect of force guidance on the trajectories generated and their metrics:  mainly following for a single or a few peptides how their trajectory change when the strength of guidance increases. This is touched upon in section f1 where it is pointed out that validity and contact metrics indeed nicely follow the guidance strength, but I feel like this is less true/obvious for the other metrics.
Finally, it would also be nice to see how easily this is applicable in more real case scenarios, with some scaling laws with peptide size, etc…

Overall I would say that all of my remarks are minors, but I still encourage  the authors to follow them as much as possible as I believe it will ease the overall understanding of the model ability. I would also like to state that even though I was able to follow all the proofs and didn’t see anything that I would deem to be problematic, I don’t believe I  have the level of expertise truly required to judge the math present in this paper with certainty.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730425036495

thread_type:
official_review

<thread_object>
creator: reviewer_orib

<content>
summary:
The paper presents a novel model, FBM, which integrates force guidance into a bridge matching framework for time-coarsened dynamics. Traditional MD methods often struggle with achieving a balance between computational efficiency and accuracy, especially over long time scales. FBM addresses this challenge by directly targeting a Boltzmann-constrained distribution, effectively combining physical priors with generative modeling to capture full-atom time-coarsened dynamics.

strengths:
- The integration of force-guided priors within the bridge matching framework for MD simulations is innovative. This approach goes beyond standard generative modeling by incorporating Boltzmann-distributed physical priors, improving the model's ability to simulate realistic molecular transitions over extended time scales. 
- The authors conduct a comprehensive evaluation of FBM on two peptide datasets, demonstrating that it consistently outperforms baseline models across several metrics, including distribution similarity and validity of generated molecular conformations. 
- The paper is generally well-structured, with a clear presentation of the methodology and objectives. 
- The model’s ability to generate high-quality conformations with transferability across unseen systems is particularly valuable for applications requiring generalized and scalable MD simulations.

weaknesses:
- The experiments to evaluate the model’s performance are conducted mainly on datasets consisting of small peptides. It would provide a more comprehensive understanding of the model's performance to extend the evaluation to include larger peptides or more complex molecular structures.
- While the FBM model demonstrates promising results, a more extensive comparison with traditional MD simulations across additional cases would strengthen its credibility. It would be helpful to provide more detailed time comparisons between FBM and MD simulations.

questions:
The TIC plots in Figure 4 are not entirely clear in demonstrating the advantages of FBM. It might be more interpretable for readers to adjust the color scheme, explore alternative visualizations, or include additional quantifiable metrics.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730447713186

thread_type:
official_review

<thread_object>
creator: reviewer_8atg

<content>
summary:
The paper proposes a conditional generative model called Force-guided Bridge Matching (FBM) to improve Molecular Dynamics (MD) simulations. Traditional MD simulations, while accurate, are time-consuming due to the need for small time steps. FBM leverages a physics-informed approach to time-coarsened dynamics by incorporating an intermediate force field to guide the generation process. This allows the model to better approximate the Boltzmann distribution, essential for realistic molecular conformations, and improves simulation efficiency without sacrificing accuracy.

strengths:
1. The paper is clearly written and easy to follow, also this task is indeed important.
2. The experiments are effectively support the conclusions.
3. The code and detailed hyperparameters are provided, making it easy for others to reproduce the results.

weaknesses:
I believe the main weaknesses with this paper is its lack of novelty, as it appears to be a straightforward combination of [1], [2] and [3] in a slightly altered setting. It seems the author did not conduct a comprehensive literature review. My reasons are as follows:

1. Training with trajectory pairs was already introduced in [1]; the only difference here is that [1] focused on protein backbones and employed SE(3) FM. In addition, next frame prediction is also a subtask in [2].
2. The concept of force-guided training was discussed in [3], with the main difference being that [3] uses SE(3) diffusion.

In essence, this paper combines aspects of both [1] ([2]) and [3] without properly referencing [1], [2]. While it presents a solid engineering application, it lacks originality.

Additionally, I have another concerns, which I’ve outlined in the Questions section.

>[1] Li, S., Wang, Y., Li, M., Shao, B., Zheng, N., Jian, Z., & Tang, J. F $^ 3$ low: Frame-to-Frame Coarse-grained Molecular Dynamics with SE (3) Guided Flow Matching. In ICLR 2024 Workshop on Generative and Experimental Perspectives for Biomolecular Design.

>[2] Jing, B., Stark, H., Jaakkola, T., & Berger, B. Generative Modeling of Molecular Dynamics Trajectories. In ICML'24 Workshop ML for Life and Material Science: From Theory to Industry Applications.

>[3] Wang, L., Shen, Y., Wang, Y., Yuan, H., Wu, Y., & Gu, Q. Protein Conformation Generation via Force-Guided SE (3) Diffusion Models. In Forty-first International Conference on Machine Learning.

questions:
1. Using bridge matching may not be ideal, as all samples along the same trajectory are expected to follow the same Boltzmann distribution. Ideally, we would achieve **a constant transformation**; however, in practice, we are performing point-to-point correspondence (from one delta distribution to another). This approach, therefore, appears theoretically unsound. What are your thoughts on this?
2. Experiments on peptides alone are too limited and should at least be performed on the smallest “protein” (backbone), Chignolin.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732180958546

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We thank all reviewers for their valuable comments. We have carefully revised our paper to address some common concerns listed below. Note that all the revisions in our updated paper have been highlighted in red.

- Responding to the concerns raised by Reviewer 8Atg about the originality of our work, we explain that this work not only provides solid engineering results but also presents a theoretical contribution by first integrating an intermediate force field based on the Boltzmann constraint into the bridge matching framework. We outline the necessity of using bridge matching to model the temporal correlations in dynamic processes, explain how bridge matching is compatible with the equilibrium distribution of MD, and discuss that, despite the simplicity of the FBM generation process, it is a non-trivial derivation that effectively combines multiple prior assumptions.
- Responding to the suggestions from Reviewer 8Atg, oRib, R7tW that experiments should be extended to include more complex molecular structures, we conduct a tiny experiment for the smallest protein, Chignolin. We fine-tune our model to the MD dataset of Chignolin curated by [1], then we evaluate our FBM on three other independent trajectories with different initial states. The details of experimental settings and results are reported in Section F.3, Table 7 and Figure 8 in our revised manuscripts. Here, we briefly present the results on three test trajectories, where FBM significantly outperforms FBM-base across multiple comprehensive metrics, showcasing a strong and stable generation ability for molecular dynamics。

| Index | Model    |   P{\small W}D  |   R{\small G}   |       TIC       |   R{\small AM}  | V{\small AL}-CA ($\uparrow$) | C{\small ONTACT} ($\downarrow$) |
|-------|----------|:---------------:|:---------------:|:---------------:|:---------------:|:----------------------------:|:-------------------------------:|
| e1s44 | FBM      | **0.315**/0.056 | **0.285**/0.080 | **0.544**/0.011 |   0.509/0.029   |        **0.691**/0.041       |         **0.161**/0.069         |
|       | FBM-base |   0.417/0.069   |   0.467/0.148   |   0.554/0.023   | **0.480**/0.013 |          0.465/0.057         |           0.249/0.066           |
| e59s7 | FBM      | **0.395**/0.017 | **0.400**/0.034 | **0.522**/0.015 | **0.443**/0.020 |        **0.780**/0.012       |         **0.184**/0.029         |
|       | FBM-base |   0.456/0.026   |   0.407/0.016   |   0.526/0.002   |   0.490/0.017   |          0.460/0.031         |           0.219/0.056           |
| e3s24 | FBM      | **0.305**/0.046 | **0.332**/0.089 | **0.524**/0.004 |   0.519/0.011   |        **0.628**/0.010       |         **0.123**/0.032         |
|       | FBM-base |   0.450/0.086   |   0.549/0.149   |   0.527/0.015   | **0.502**/0.019 |          0.490/0.040         |           0.278/0.104           |

- Reviewer oRib suggests that it would be helpful to provide more detailed time comparisons between FBM and MD simulations. We believe that during the simulation process, both the generation speed and sample quality of the model need to be fairly considered. Therefore, we chose the effective-sample-size per second of wall-clock time as the evaluation metric for efficiency, which was used by Timewarp [2]. The comparison of MD simulations and FBM is shown in Figure 5.c, where FBM achieves an efficiency improvement of around 10 times relative to MD over the entire test set.
- Reviewer oRib suggests that the color scheme should be adjusted for better comparison, and Reviewer oRib, R7tW believe more visualizations of comprehensive metrics would be super helpful. In our revised paper, we have adjusted the color scheme of TIC plots of Figure 3 and 4 for higher contrast. We also made significant revisions to Figures 4 and 5 that provides visualizations of comprehensive metrics, including: the distribution of the radius-of-gyration (Rg), the residue contact map, cumulated valid conformations during inference, $C_{\alpha}$-RMSD values along the trajectory compared to initial states, and the effective-sample-size per second compared to MD simulations. Meanwhile, the visualization of 3 other test peptides was illustrated in Figure 7, Section F.2, where we show FBM demonstrates super transferability to unseen peptides with different residue lengths.

We would give further responses to other specific concerns to each reviewer individually.

> [1] Pérez Culubret, Adrià; De Fabritiis, Gianni (2021). Chignolin Simulations. figshare. Dataset. https://doi.org/10.6084/m9.figshare.13858898.v1.

> [2] Klein L, Foong A, Fjelde T, et al. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics[J]. Advances in Neural Information Processing Systems, 2024, 36.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732567348245

thread_type:
official_comment

<thread_object>
creator: area_chair_6fcs

<content>
comment:
Dear reviewers,

Tomorrow (Nov 26) is the last day for asking questions to the authors. With this in mind, If you have not already done so, please read the latest update to the rebuttal provided by the authors earlier today, as well as the other reviews. If you still have outstanding comments or questions, please provide your updated view _accompanied by a motivation_ and raise any outstanding questions for the authors.

**Timeline**: As a reminder, the review timeline is as follows:
- November 26: Last day for reviewers to ask questions to authors.
- November 27: Last day for authors to respond to reviewers.
- November 28 - December 10: Reviewer and area chair discussion phase.

Thank you for your hard work, 

Your AC
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper received mixed scores, with two reviewers leaning to a clear accept (score 8), and one reviewer initially learning towards recommending to reject this paper (score 3). Among other things, the paper was thought to have a clear presentation, and claims that were validated with their empirical results, which included evaluations on alanine dipeptide and small dipeptides. As areas for improvement, points were raised on the need for validation on more than just small peptides and the need for timing comparisons. Furthermore, the reviewer who recommended rejection indicated that they viewed the work as a straightforward combination of existing concepts in prior work, and that the work did not show an advantage through this combination. 

During the rebuttal, the authors have added additional experiments for the small protein Chignolin, where they compare their model against MD only. Additionally, the authors included comparisons of effective sample size per second of wall-clock time as an additional metric for efficiency, which was introduced in prior work.

Based on this rebuttal and following discussions the two reviewers in favor of acceptance maintained their score, while the reviewer with the most concerns raised their score from a 3 to a 5, based on the remaining concern of limited novelty.

After the author-reviewer discussion ended, I read through the reviews again, as well as the discussions with the authors and the newly computed results and the revised paper. Upon going through the revised paper to look through the added results, I found that the results on alanine dipeptide raise questions about their correctness. In particular, the distribution of samples from the models in figure 3 do not seem to resemble the distribution represented by the contours based on kernel density estimation of MD trajectories, and the baseline methods in particular do not perform well at all. This discrepancy between model distributions and the MD-based distribution is striking since one of the baseline methods, Timewarp, has shown in their original paper (https://arxiv.org/pdf/2302.01170) to be able to draw samples of alanine dipeptide according to a distribution that closely mimics the distribution of samples from MD trajectories. See figure 3 in the Timewarp paper. This raises doubts for me on whether or not the results in this current submission were calculated correctly. 
After consulting with the reviewers, the reviewers agreed that the results raise concerns about correctness. Unfortunately, the range of angles used to compute Ramachandran plots of alanine dipeptide in this submission were not displayed, nor are samples from the reference MD simulation included in the Ramachandran plots. This makes it hard to determine if the discrepancy between the results reported in this paper compared to the Timewarp paper are influenced by different data or plotting settings. Moreover, we received input from one additional reviewer with significant expertise in this field, who was initially unable to review the paper but looked through the results and other reviews. In addition to sharing concerns about the alanine dipeptide results, they noted that the TICA plots in figure 4 show worryingly bad results, and that the energies in Figure 6 are also not good. Unfortunately, during the reviewer-AC discussion phase it was no longer possible to ask for clarification from the authors, and the presented results did not have enough details to be able to determine where the discrepancies with previously reported results come from. Given the overall consensus among reviewers and AC that the presented results raise concerns, I recommend rejecting this paper. 

I would like to thank the authors for submitting the paper and for providing additional results during the rebuttal. I hope the above explanation can be helpful in rewriting the manuscript for a future submission.