Universally Applicable And Tunable Graph-Based Coarse-Graining For Machine Learning Force Fields

We thank you for your detailed review of our manuscript and the well-explained summary of our work.

Weaknesses

I think... The extension of Webb et al. is only one of our contributions. Our main contribution is a significant step towards a truly generalizable CG-MLFF model, i.e., one that is transferable across different systems of diverse (organic) chemistry. Such an approach has not been presented before. We present a complete pipeline for: (1) the generation of a chemically diverse dataset (2) a versatile and tunable CG mapping, and (3) the application of a state-of-the-art MLFF architecture (MACE).

The authors... We will add a statement of explanation close to line 150 to avoid the confusion.

In line... The 1st eigenvector of the graph Laplacian contains information about the contribution of each node (i.e., atom) to the overall connectivity. The more connected an atom, the higher its contribution. Considering not only the atom’s neighbors but also its neighbor’s neighbors, and so on. Assigning high priorities to well-connected atoms, ensures the overall bond topology is maintained when coarse-graining.

Equation... For an intuitive explanation, first, it is important to understand that there exist many valid thermodynamically-consistent force mappings given a CG mapping. By coarse-graining, there is loss of information about the all-atom forces, however, some mappings add more statistical noise than others. Kramer et al. derives that a minimal CG force magnitude correlates with the minimal noise in CG forces.

What.... A more detailed discussion of this will benefit the paper. The heuristics include readily computable features of the chemical bonds, which could be relevant for prioritizing the atoms attached to them; Bond distance; bond polarity, and essentially “bond mass”. We are open to further suggestions for future versions of the model.

The 1/... Modifications at each step of simulation would lead to unsmooth behavior. Our coarse-graining remains constant, and based on the bond distances of the minimum energy structure at the start of simulation (roughly equal to the expected value of the bond distance during the simulation). During training, input structures are not optimized, however, many configurations of the same bond are sampled in the training set (with the same expected value), which should regularize the model to not overfit one specific distance for a given bond (if in reality, it fluctuates during a simulation).

In line... It is the number of elements, the proportions of each element has also been tested, leading to worse results. We understand the concern, however here the CG strategy will group between 1 and 10 atoms to form a bead. If moving to more coarse-graining, the models may be revisited.

Line... This subset is only used for tuning the CG parameters. 100 fragments -> 3000 data points, each fragment comes with all conformations. Training on “entire systems”, though not as many compared to the full force field training later in the pipeline. Taking more fragments for the CG tuning could improve model, but this has not been investigated rigorously in this work.

Line 34... For the force field training, we use the same loss as regular training of the ML force field, the same as in the original MACE code, i.e., the MAE of the forces.

Line... It is the Pearson correlation between the predicted and ground truth forces.

Fig... The presented curve does not show the data on the training set but rather for the validation set. This is mentioned in the caption. We will rephrase to “validation set metrics for the standard and tuned model during training.”.

Line... This was due to time constraints. However, the model has finished training, and the results are qualitatively similar, hence, the training had already almost converged at that point.

Line 428... The tuned model may slightly overfit to the training data leading to less stable MD simulations for unseen systems. Little improvement in MD stability under tuning may also be due to the shared optimization of the average forces metric during training and CG model tuning. Ultimately MD stability is related to inaccurate individual predictions along a trajectory. This is a well-known issue not only for CG force fields, but also for all-atom ones, and the principal reason.

Figure... In comparison between models, Figure 4 demonstrates more physically reasonable simulations from the tuned compared to standard model. The energy barriers between different conformations are high for the standard model, it is stuck in one configuration. In all other simulations, especially the GFN-FF reference, we observe multiple conformations being visited, which is also observed for the tuned model, demonstrating more accurate recovery of relative energies of conformations.

We will address all these elaborated points in the updated manuscript.

Thank you for your detailed comments on our manuscript, we'll respond to each of your numbered weaknesses and questions below:

Weaknesses

1. This is a very good question and we agree that this part of the paper is lacking clarity. We optimize our CG mapping parameters with respect to the DFT-aggregated forces, or more precisely, their average magnitude across a small sample of the dataset. As we show in section 3.1 and explain in the Methodology section, this is a very good proxy for how well the MLFF model trains on the CG data. The MLFF-predicted forces would be much less informative at that stage as they are basically random, because the model has not been trained yet at that stage. We propose to add the three sentences above for clarity to section 3.1. Furthermore, we would like to point out that the extension of the CG mapping strategy by Webb et al. is only one part of our key contribution to the field. We agree that the data presented in our paper is not sufficient to prove for all upcoming use cases that the tuned approach will have significant superiority over the standard one. However, the main contribution of our work is that we apply this CG mapping strategy to build a transferable CG-MLFF model that is not restricted to very specific systems like proteins (where the CG mapping can be hardcoded). Furthermore, we apply the MACE architecture in the CG context which has not been done previously.

2. We agree with this criticism partly in the sense that more evaluations and metrics are required to demonstrate high stability and robustness of a CG-MLFF approach. However, we believe that this is outside the scope of our proof-of-concept study presented in the manuscript. For our responses on some of the details of the review above, please see our responses below to the specific questions the reviewer posted.

3. We refer the reviewer to our response to Weaknesses for Reviewer iQM8.

Questions:

1. This question has been addressed in the response to the weaknesses above. We agree that more clarity in the description is needed and it will be added.

2. We agree that additional stability metrics would be beneficial to validate model reliability and we have looked into many different metrics as well as the atom-level stability. However, the main issue with many of these are that they are defined for all-atom simulations, where for example, the correctness of the neighborhood of an atom can be easily determined at each time step. For CG systems, we attempted to come up with a generalization of such metrics but did not succeed in a consistent way as comparison with Martini trajectories demonstrated. Of course, we remain interested in dedicating more time into developing such evaluation scores for CG MD trajectories in the future.

3. We agree that a quality assessment of the MD trajectories by comparison to either Martini or a classical all-atom force field is of high interest, and hence, with have included such a comparison for one of our test systems, the IPP lactotripeptide, as depicted in Fig. 4 and 5. We have selected this system as it exhibits very high MD stability with our model for long simulation times (several nanoseconds) which are required for such comparisons of RMSD and TICA plots. For other systems, we have either not observed this outstanding stability or, as in the presented case of the RNA systems, observed other unphysical artifacts in the simulations. We would like to point out that we do not claim to present an application-ready superior CG model in our work, but instead present a proof-of-concept for a very novel type of model – a transferable CG-MLFF that can be applied across different (organic) chemical systems and is trained on a dataset with large chemical diversity. We see this as a significant step towards a CG-MLFF model that is usable in real-world applications. However, we also want to be transparent about the limitations of the current model which aims at generalizing across chemical systems in a way that has not yet been reported in ML-based CG modeling.

4. As transparently explained in the manuscript, there were a few more test systems we looked at in this context. Either we did obtain stable simulations, but very uninteresting sampling of conformations even for our reference methods (GFN-FF, Martini), or systems were not stable enough for long periods of simulation time. We agree that finding more interesting systems for comparison of energy landscapes would be beneficial in the long term, but we also believe that the scope of this paper was providing a proof-of-concept that generalizable CG-MLFFs are possible and for this scope we believe that the example presented in the manuscript is sufficient.

5. We would like to ask the reviewer if they could explain this observation a bit more to us. We are unsure about what this observation is referring to.

6. Please also refer to response to Weaknesses 3.

Apologies for the delayed response, and thank you for your follow-up. However, I will still keep my initial score.

In response to the question:

We would like to ask the reviewer if they could explain this observation a bit more to us. We are unsure about what this observation is referring to.

I recommend referring to Figures 2 and 5 in [1]. This paper provides an excellent example of how to compare two different force fields based on the landscape derived from simulated trajectories.

[1] Charron, Nicholas E., et al. “Navigating protein landscapes with a machine-learned transferable coarse-grained model.” Nature Communications, 2023.

Thank you for your detailed comments on our manuscript, we'll respond to each of your numbered weaknesses and questions below:

Weaknesses:

The authors should clarify the novelty of their work compared to prior work using potentials like Allegro (Voth et al, see citation above) or FLARE (Kozinsky et al, see citation above) and show quantitative comparisons with these prior works. This is especially the case since those prior works demonstrated quantitative agreement between all-atom and coarse grained observables for the simulations of interest.

The main goal of this paper was to show what a general strategy for coarse-graining, the most state-of-the-art architecture for force field prediction and a new dataset generation method allowing to include larger fragments, could lead to in terms of a general strategy to perform coarse-grained MD.

It is one of the few (if not the only) paper that explores the training and generalizability of a ML-based coarse-grained force field for a wide range of chemistry (RNA, lipid, proteins) and for a wide range of molecule sizes. Furthermore, the paper doesn’t just validate the force field on a validation set, it tests the force field’s usefulness by applying it in MD simulations. Papers that do run MD simulations for validation, often just run on systems they have been specifically trained on instead of showing generalization to new systems. To the best of our knowledge, Allegro has neither been employed in the CG context nor in the context of being a generalizable force field yet. However, we agree that investigating this architecture in future work is an interesting direction. Regarding FLARE, we have a similar view.

We agree that an evaluation on the original molecules from Webb et al. comparing the original CG scheme to our tuned one could be helpful in addition to our training comparison. However, as the molecules tested in Webb et al. are not of the same type as the ones we trained on in our work and since they trained specifically on these individual systems instead of in a transferable fashion, a comparison of the CG methods would be difficult unless the same force field would be employed.

Regarding prior CG work with MD: except for coarse-grained chignolin, data are not available or deal with material science, which is outside of the ability of our model. For chignolin, our simulation was not stable enough to do a direct comparison. That being said, we will continue to search for more reasonable comparisons in future work as we agree that these are strongly required to assess a CG force field model. Moreover, in most prior work, specific force fields have been trained for those specific molecules, which we reemphasize makes the comparison difficult.

We should definitely add more all-atom vs. CG comparisons in the future, possibly using MACE-OFF for small-enough molecules in our test set and classical force field for larger molecules. However, we emphasize that we explicitly added a comparison of our tripeptide test system with the all-atom force field GFN-FF. We also did a comparison of this test case with MARTINI, but we also agree that more such comparisons must follow in the future, but were not possible within the scope of our current model. We also thank the reviewer for their useful paper recommendations.

Questions

2. Ablations are needed to show the improvements of MACE vs SchNet and original (Webb et al) vs the CG implementation here would be very helpful. Alternatively, quantitative comparisons between this work and prior work would also establish the improvements in this work.

We consider the comparison of performance of different force field architectures in a CG context to be out of scope for this paper. As explained above, quantitative comparison to prior work is difficult but additional all-atom vs. CG comparisons and CG-MLFF vs. MARTINI can be explored.

1. Quantitative agreement and thermodynamic consistency should be shown for reference systems in prior work. Qualitative comparisons of energy or RMSD for a fixed time window is not sufficient validation of a coarse grained simulation.

Same as above.

Thank you for your detailed comments on our manuscript, we'll respond to each of your numbered weaknesses and questions below:

Weaknesses:

1. In the field of machine learning force fields (MLFFs), we generally distinguish between system-focused and transferable MLFF approaches. In the former, the model is just trained on a single molecule in different conformations (a single potential energy surface) with the goal to later obtain accurate simulations for this molecule. On the contrary, transferable approaches train on a dataset with a variety of systems with the goal to be applicable to unseen molecules of similar chemistry. In the field of CG-MLFFs, the large majority of published work has been done in the system-focused setup which is common for proof-of-concept studies. One of the reasons for this is that a generally applicable CG mapping is not trivial to define, however, as a result we are lacking understanding of whether a truly transferable CG-MLFF approach is even feasible. Due to the difficulty of defining a general CG mapping, the few published transferable CG-MLFF approaches are restricted to very specific types of systems, typically proteins where the CG mapping can be fixed for each amino acid in a simple rule-based way. Our presented approach is novel in the aforementioned sense. Our CG mapping strategy works for any (organic) chemical system, which is reflected not only in the fact that our dataset contains lipids, RNA and proteins, but also that due to the fragmentation-based dataset generation we train on a diverse set of organic fragments (including molecules from the PubChem database). This is the generalizability that we claim at the end of our Introduction. Transferability across system sizes is another important topic for MLFFs. The ability to train on system sizes in a the range of 20 to 200 atoms per system, but applying the learnt model to systems of 1000s of atoms. This is an MLFF-related challenge and not related to the coarse-graining problem. For all-atom force fields, this is already a difficult task and for CG-MLFF, to the best of our knowledge, it has never been shown possible. We show in our paper that our model is able to do this by running an MD on the 2Z75 system of >3000 atoms – however, we acknowledge that a detailed analysis of robustness across a large number of such systems was beyond the scope of our work. We propose to add a paragraph to highlight our core contribution more prominently in our introduction

2. Our paper is both (i) applying the CG mapping strategy of [1] in a new context that has not been done before (transferable MACE-based CG-MLFFs), and (ii) extending the method proposed in [1] by adding tunability. It is important to highlight that the combination of these two aspects is our key contribution to the applied ML community. First and foremost, we present the first transferable CG-MLFF approach that is not limited to a single type of system like proteins for which the CG mappings can be hardcoded. This feature is enabled by the versatile CG mapping strategy proposed in [1]. Furthermore, we add two additional contributions, (a) the extension of [1] by adding the tunability of the spectral decomposition (multiple Laplacians) and (b) using MACE as the force field architecture which has not been applied to the CG use case before. We propose to highlight our core contribution more prominently in our introduction and abstract.

3. As outlined above, our approach contains several new features compared to previous CG-MLFF work. As a result, we decided to concentrate on specific aspects of the pipeline for ablation studies, most prominently, comparing the standard CG mapping strategy close to the one reported in [1] with our tunable extension. Adding a comparison of MACE with other equivariant potentials like the ones you mentioned are definitely also interesting but were not in our focus for this work. Unfortunately, as model training takes multiple weeks on our large dataset, we are not able to add such comparisons to the paper at this moment but would be interested in revisiting this aspect in future work. We also propose to add a sentence explaining this to the paper in section 2.1.

Questions

1. The presented curve does not show the data collected on the training set but rather for the validation set. This is mentioned in the caption, but we assume the word “Training curves” has led to confusion. We would therefore like to rephrase it to “Validation set metrics for the standard and tuned model during training.”.

2. The same plot for the standard model is shown in Figure 8 in Appendix A.7 due to the space constraints in the paper. We propose to change the paper to mention this not only in the text but also in the figure caption. This should make it easier to find. We are not aware of any straightforward way to quantify the agreement of TICA plots that has been applied in the literature before. Are you aware of any metric that we could compute?