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

We sincerely thank your constructive comments and your recognition of our work. We address your concerns below accordingly.

• 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. ... Overall all those figures and metrics are complementary and I believe for a few samples to have all of those would help a lot.

Thank you for your valuable suggestions! In response, we have visualized comprehensive metrics, including Figures 4 and 5 in the main text and the newly added Figure 7 in the appendix. These visualizations of test peptides in different length scales intuitively demonstrate the excellent performance of FBM on OOD data and a great alignment with MD simulations.

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

That's a good question for discussion. In fact, to fit the true Boltzmann distribution, we should aim for an "appropriate" rather than a "larger" guidance strength. On the one hand, a larger guidance strength provides stronger physical constraints, theoretically making the model more inclined to generate reasonable conformations, which aligns with the conclusions shown in Table 6. On the other hand, an excessively large guidance strength suppresses the variability of the generative model, causing the distribution to shift toward the equilibrium Boltzmann distribution while neglecting the distribution $q_t(x_t)$ learned from time-correlated data pairs. As a result, the model may perform worse in terms of distribution similarity. This might explain the issue you raised.

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

Thank you for your instructive suggestion and we've now conducted a tiny experiment on the smallest protein, Chignolin, and the experimental details can be found in section F.3, Table 7 and Figure 8, where FBM shows a close match to those of MD in most cases, demonstrating its usefulness and scalability to more complex molecular systems.

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

Thanks for your concern and it deserves further discussion. We believe that the statements "the distributions at both ends of the bridge are ideally Boltzmann distributions" and "bridge matching performs point-to-point correspondence" are not contradictory.

Firstly, given a temporal interval $\tau$, we select multiple data pairs $(x_0,x_1)$ from the MD trajectories, which are temporally correlated and thus not independent. However, this does not affect the fact that the marginal distributions of their joint distribution are ideally Boltzmann distributions, i.e., $\int{q(x_0, x_1)dx_0}=\int{q(x_0,x_1)dx_1=\exp(-kU(x))}$.

Secondly, as a Markov process, bridge matching actually learns the expectation of the posterior distribution $q_t(x_0,x_1|x_t)$ given the state $x_t$ at diffusion time $t$ from the joint data distribution $q(x_0, x_1)$. The equation converges to a delta distribution as $t\to{0}$ or $t\to{1}$ only when the bridge is pinned down to two given starting and ending states (i.e., during the training process). Therefore, in the ideal case, the distribution learned by bridge matching at $t=0$ (or $t=1$) is the expectation of the delta distribution, which is "luckily" identical to the marginal distributions $q(x_0)$ and $q(x_1)$. Taking all the above into consideration, bridge matching is compatible with the assumption of the stationary distribution in MD simulations.

• Experiments on peptides alone are too limited and should at least be performed on the smallest “protein” (backbone), Chignolin.

Thanks for your suggestion! We have now conducted a tiny experiment on Chignolin, and the experimental details can be found in section F.3, Table 7 and Figure 8, where FBM shows a close match to those of MD in most cases, demonstrating its usefulness and scalability to more complex molecular systems.

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

[4] 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.

[5] Lindorff-Larsen K, Piana S, Dror R O, et al. How fast-folding proteins fold[J]. Science, 2011, 334(6055): 517-520.

[6] Anderson B D O. Reverse-time diffusion equation models[J]. Stochastic Processes and their Applications, 1982, 12(3): 313-326.

• The concept of force-guided training was discussed in [3], with the main difference being that [3] uses SE(3) diffusion.

We appreciate the reviewer’s observation that the concept of force-guided training draws inspiration from [3]. This connection was honestly and explicitly acknowledged in our Related Work section. However, we would like to emphasize that our adoption of the bridge-matching model as a replacement for the diffusion model is both well-motivated and nontrivial.

From a problem-solving perspective, the bridge-matching model is better suited to our MD task. While the approach in [3], which employs a diffusion model, is effective for generating conformation ensembles, it does not account for the temporal correlation between them. Consequently, it cannot be utilized to study dynamic processes. For example, proteins may undergo folding and unfolding transitions under varying environmental conditions [4]. While conformation ensembles can capture different states, they cannot characterize the transitions between these states. To address this limitation and effectively learn transition probabilities between temporally correlated data pairs, we chose bridge matching. This approach generates molecular conformations conditioned on the starting conformations rather than relying on Gaussian priors, making it a more appropriate choice for the MD task.

From a methodological perspective, conducting force-guided training for the bridge-matching model is far from trivial. A core aspect of deep generative models is the learning of the generation process. For the diffusion model used in [3], this process (the backward process) is driven by the score function $\nabla\log{p_t(x_t)}$, which can be straightforwardly derived using Eq. (6) in [3]. Additionally, the immediate force field $\nabla_{x_t}\varepsilon_t(x_t)$ involved in the computation is directly obtainable via Eq. (5) in [3]. In contrast, in our work, the generation process for the bridge matching model is governed by the vector field $v'(x_t, t)$, which is significantly more challenging to derive. To address this, we should not only establish the relationship between the vector field, the score function and the immediate force field through Propositions 3.1, 3.2 and 3.3, but also rely on the assumptions of symmetric Brownian bridge and the detailed balance of the stationary distribution of MD simulations. While the final form appears relatively simple, the derivation process is intricate, as detailed in the proof provided in Section B of our paper.

Furthermore, the immediate force field involved in Proposition 3.1 has a more complex form compared to that of [3] (please compare Eq. (9) in our paper with Eq. (5) in [3]). This is basically because both ends of bridge matching are data distributions, while in diffusion only one end is a data distribution and another is Gaussian. Computing the force field necessitates the score $\nabla\log{q_t(x_t)}$, which, in turn, requires introducing the vector field of the reverse-time Brownian bridge (Eqs. (11–12)) and the use of Proposition 3.3.

As such, our force-guided bridge-matching approach is not a simple substitution for the diffusion model. Instead, it embodies theoretical insights and involves nontrivial derivations, making it a meaningful and innovative contribution to the field.

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

Overall, we acknowledge that our task (i.e., "forward simulation" or "time-coarsed dynamics") is consistent with that of [1] and [2], yet we did not claim that the task was first introduced by us. Similarly, we recognize that the concept of force-guided training is inspired by [3]. That said, applying force-guided training to our task poses unique challenges. We opted for the bridge-matching model over the diffusion model as it is better suited for MD tasks. To integrate the force field into the bridge matching framework, we have developed several insightful propositions and advanced techniques, which we believe represent meaningful theoretical and practical contributions.

We sincerely thank the reviewer for the comments. We respectfully disagree that our paper lacks novelty and is a straightforward combination of [1], [2] and [3]. The reviewer probably misunderstood our claims and contributions. We provide the reasons below.

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

Thank you for bringing these two references [1,2] to our attention. It seems there might be a misunderstanding regarding our claims. We do not assert that we are the first to perform MD simulations using "training with trajectory pairs" or "next-frame prediction". In fact, in the third paragraph of the Introduction, we explicitly acknowledge that prior works, such as Timewarp [4], have already employed this training strategy, albeit under a different terminology, namely "time-coarsened dynamics simulation." Our intended claim is that our method enhances these existing approaches by incorporating the underlying Boltzmann constraint and physics priors, which are not addressed in [1,2].

We also apologize for initially overlooking the two works mentioned [1,2], as they are both workshop papers. In response to the reviewer's feedback, we have now properly cited and discussed these works in the Introduction and Related Work sections of our revised manuscript. Since both methods take a coarse-grained approach (F$^3$low [1] focuses on protein backbones, and MDGen [2] treats backbone atoms as rigid bodies), we did not implement them in our experiments to ensure fair comparisons.

Dear reviewer 8Atg,

We sincerely appreciate your great support and insightful comments, and we are also very grateful for your increased evaluation of our work. Here we would like to address your additional concerns, which may help facilitate a better understanding and a more in-depth discussion.

• I would like to point out that flow matching and diffusion models are, to some extent, unified (meaning the score function and vector field are somewhat equivalent).

We fully agree with your statement that the vector field and score are equivalent to some extent, which have been proven by Proposition 3.3. However, the main difference arises from the fact that diffusion has a simple Gaussian prior at one end, while both ends of bridge matching are data distributions. This leads to the diffusion SDE process being controlled by a single variable (e.g., the score function), while bridge matching must be described by two variables (e.g., $v_{\mathrm{fwd}}$ and $v_{\mathrm{rev}}$ in our work). To be more specific, both the reverse-time diffusion and bridge matching can be unified with the following SDE process, where we use the notations in [1]:

$dx_t=b(t,x_t)dt+\epsilon(t)s(t,x_t)dt+\sqrt{2\epsilon(t)}dW_t,$

where $s(t,x_t)$ is the score function at diffusion time $t$, $b(t,x_t)=-\frac{1}{2}\beta_tx_t$ as in [2] while $b(t,x_t)=\frac{v_{\mathrm{fwd}}(t,x_t)+v_{\mathrm{rev}}(t,x_t)}{2}$ in our bridge matching framework. It is obvious that, though the score function $s(t,x_t)$ is considered in both methods, the drift term $b(t,x_t)$ has a quite non-trivial form in bridge matching. In fact, most of our effort has been dedicated to proving that, after incorporating the Boltzmann constraint, $b'(t,x_t)=b(t,x_t)$ holds (Proposition 3.2), which is not considered at all in diffusion. We believe this is the main challenge and contribution of incorporating the Boltzmann constraint into bridge matching.

• Additionally, in your rebuttal, the bridge matching works because $q(x_0)$ and $q(x_1)$ are treated as distinct delta distributions. However, my question assumes that $q(x_0)$ and $q(x_1)$ follow Boltzmann distributions. In this case, introducing bridge matching does not seem ideal. It might be more convincing to condition on $x_0$ and sample $x_1$ directly from the Boltzmann distribution.

Our rebuttal may have caused some misunderstandings and we would like to clarify it here. Our statement is, given the bridge pinned down to two endpoints $x_0,x_1$, i.e., in the training process, the learned forward vector field $v_{\mathrm{fwd}}$, which is equivalent to $q_t(x_1|x_t)$, will converge to $q(x_1|x_t)\delta(x_t-x_0)$ when the diffusion time $t$ converges to 0. Therefore, in the ideal case, the model learns the expectation of the delta distribution when $t=0$ after the whole training process. That is, $\mathbb{E}_{x_0}q(x_1|x_t)\delta(x_t-x_0)=q(x_1|x_0)$, where $q(x_1|x_0)$ is the transition probability estimated on the training dataset. Thus, if we assume "$q(x_0)$ and $q(x_1)$ follow Boltzmann distributions" in the ideal case, our model will learn the transition probability of the Boltzmann distribution accordingly. Thus we believe that the assumption of Boltzmann distributions and the bridge matching framework are indeed compatible.

By the way, we also believe that, in the case of randomly selecting data pairs from the original MD trajectories for training, the assumption of both ends following the Boltzmann distribution is unlikely to hold. In such non-ideal circumstances, introducing the Boltzmann constraint helps bring the generated distribution closer to the Boltzmann distribution, which is also the reason we incorporate the force guidance into the bridge matching framework.

We hope that our response helps to further address your concerns.

Best regards,

Authors of #4521

Reference

[1] Albergo M S, Boffi N M, Vanden-Eijnden E. Stochastic interpolants: A unifying framework for flows and diffusions[J]. arXiv preprint arXiv:2303.08797, 2023.

[2] Wang Y, Wang L, Shen Y, et al. Protein conformation generation via force-guided se (3) diffusion models[J]. arXiv preprint arXiv:2403.14088, 2024.

We sincerely thank your constructive comments and your recognition of our work. We address your concerns below accordingly.

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

Thank you for your instructive suggestion. As mentioned in the general response, we have now conducted a tiny experiment on the small protein, Chignolin. The experimental details can be found in section F.3, Table 7 and Figure 8, where FBM shows a close match to those of MD in most cases, demonstrating its usefulness and scalability to more complex molecular systems.

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

Thanks for the suggestion. We acknowledge that our initial time comparison with MD simulations was not very convincing. Therefore, we provide a more fair and reasonable comparative experiment in our revised paper. In particular, 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 also used by Timewarp [1]. The time 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.

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

Thanks for your advice! We've now adjusted the color scheme of TIC plots in Figure 3, 4 for higher contrast. Meanwhile, we add much more visualizations for comprehensive metrics in Figure 4, 5 and 7, including the distribution of Rg, residue contact map, cumulated valid conformations during inference, $C_{\alpha}$-RMSD over time compared to the initial states, etc. Through the visualization results of peptides of different lengths, we demonstrate that FBM exhibits excellent transferability, a strong capability to generate reasonable conformations, and is able to consistently reach the reference equilibrium distribution (1e28:C, 1n73:I, 1gxc:B), and there's still room for improvement when dealing with more complex molecular systems (1qj6:I).

[1] 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.