A Multi-Grained Group Symmetric Framework for Learning Protein-Ligand Binding Dynamics

Changed title

We are happy to change the title from group symmetric to SE(3)-equivariance. And we want to point out that the SE(3)-equivariant model for interaction, especially the multi-granularity, is indeed one core contribution of our work. We spent the whole Sec 3.1 and 3.2 discussing the details.

Answers to the Questions

• Clarifications on ML terminologies.
  • In this work, our usage of the dot product is the same, and you can check this wiki page.
  • We only have scalars and vectors in the input and output, the representations (e.g., $h_i$) are invariant/equivariant representations. We have added the dimension of each variable in Sec E in the revised manuscript.
  • The number of layers means how many times we repeat the MPNN layer (we omitted the index of layer for simplicity).
  • The MPNNs are defined in Sec 3.2 explicitly:
    • The MPNN in BindingNet-Ligand corresponds to Eq 5.
    • The MPNN in BindingNet-Protein corresponds to the last line of paragraph Backbone-Level Protein Modeling.
    • The MPNN in BindingNet-Complex corresponds to Eq 6.
• MAE is short for ``mean absolute error’’. We have shown its equation in Sec 3.3 below Eq (8) when it first shows up.
• When we mean we keep the same backbone, that means we use the BindingNet to provide atom-wise energy/force prediction. This will be fed into the Gaussian distribution in GNN-MD or scoring function in DenoisingLD to learn/estimate $p(x_{t+1}|x_t)$. If readers are interested in these works, we refer to the Related Work paragraph in our manuscript (with references) for an in-depth read. BTW. In ML papers, we often refer the baseline papers to the audience for them to get the details, and leave the main content for our paper’s novel algorithm (Sec 3 in our paper).
• We had this in the first version. We have mentioned this under Ablation Study on Flexible Setting in Sec 4.1, and please check Sec F for more details.
• The timesteps in the current experiment are fixed (not adaptive) so as to fit the MISATO dataset.
• Training on the first 80 snapshots, and testing on the last 20 snapshots. The 80th snapshot is given in inference.

[1] Allen, Michael P., and Dominic J. Tildesley. Computer simulation of liquids. Oxford university press, 2017.

About performance and baselines

• First, we want to mention that our model is achieving almost consistently better performance, compared to all the baselines. This is definitely not a perfect solution, but it represents a novel scientific progress in the ML community.
• Then we want to emphasize that, one insight/goal of this work is to open a novel physics-inspired research line of ML for binding MD prediction, and please allow progressive updates on this.
• We also have certain questions about RMSF. (1) We do not understand the following sentence: The RMSF is the best result that can be achieved by a single static structure and is an essential missing baseline. (2) By RMSF, if you mean root mean square fluctuation, then it’s a metric, not a method/baseline. (3) RMSF, as a metric, tells the fluctuation of every single trajectory, thus, it is an unsupervised metric and is not easy to compare with. On the other hand, the stability (proposed by Fu. et al.) metric measures the fluctuation between the predicted trajectory and the ground-truth trajectory, and we want to argue that it is a more comparable metric that includes the fluctuation/stability.
• Comparison between NeuralMD and DenoisingLD. (1) According to Table 3, NeuralMD is better than DenoisingLD by 10%, 3%, and 70% on three datasets, respectively. We want to argue that this is a consistent and stable improvement. (2) Besides, DenoisingLD is overdamped Langevin dynamics, which is inappropriate for the general MD estimation unless under specific conditions. We leave more details in Sec G in the revised manuscript.

Thank you for raising your comments. There is one critical comment and some misunderstandings, and we hope we have addressed them in our rebuttal. If there are any other follow-up questions, we will be happy to address them.

Math is correct. BindingNet is SE(3)-equivariant because the mass center is removed.

Thank you for raising this concern. However, we mentioned that all the atom and backbone residue positions are reduced by the mass point of the molecular system in the first version, thus it is translation equivariance. We have described this in line 3 of Algorithm 1. Further for clarification, we have highlighted this in the revised version, especially in the proof (Sec C).

Correction on problem formulation

Thank you for raising this question.

• First on the solvent molecules. The ideal case is that the MISATO authors can provide them in the dataset. We contacted the authors and tried to extract the dynamics of solvent molecules, but such information is still missing in the dataset (we are waiting for the replies now).
• However, such missing information can be easily captured using our NeuralMD framework.
  • Current in Newtonian dynamics: $F = \text{BindingNet(protein-ligand)}$
  • Add random terms to mimic solvent molecules as Langevin dynamics: $F = \text{BindingNet(protein-ligand)} - \gamma m v + \sqrt{2m \gamma k_BT} R(t)$. We have this explained in the revised manuscript (the colored text at the end of Section 3.3), please feel free to check. For the rest part of the manuscript (e.g., Abstract and Sec 1), we can change it from ODE to SDE easily in the final version.
• The results of single trajectory and multiple trajectories are attached below. You can also find them in the revised manuscript.
• MISATO is under the NVT configuration, and it uses thermostats, then Langevin dynamics can capture such stochasticity [1].
  • Besides, these two (thermostats or barostats) are used on velocity to guarantee that the simulated trajectories (coordinates/positions ) follow the configuration. In this way, we are using NeuralMD to learn the coordinates/positions directly. By expectation, we want to claim that such an implicit bias can be learned by our models, especially after we introduce using the Langevin dynamics.
  • There are also variants of thermostats, e.g., the Nose-Hoover thermostat is deterministic, while the Langevin thermostat is stochastic. Thus, it highly depends on the dataset generation process. To verify it, we have sent emails to the authors of MISATO, and still waiting for their response.

Thanks for the detailed response.

(1) Thanks for clarifying the SE(3) equivariance. However, the Appendix refers to a nearest-neighbor algorithm $x_k$ to construct equivariant frames, deviating from the exposition in the main text. This misalignment should be remedied. I'm also unconvinced by this kind of equivariance accomplished by removing the center of mass as we lose the ability to locally learn equivariant features.

(2) I appreciate the attempt to justify the ODE formulation; however, these are unconvincing as the evaluation metrics are still based on the assumption that one wants to learn an MAE-minimizing simulator. The authors have claimed to capture stochastic dynamics by learning the best deterministic simulator and adding random noise after the fact. This is very different from learning stochastic dynamics.

(3) The RMSF gives an idea of the intrinsic variance of the trajectories and hence the MAE obtained by the best dummy predictor which always predicts a static frame for all timesteps.

Because the problem formulation remains unconvincing and the experimental results remain unimpressive, I will keep the current score.

Hi reviewer oR9H, thank you for the follow-up. We're pleased to delve deeper into addressing your concerns.

Q: SE(3)-equivariance

1. Thank you for the confirmation. The proof and the main body can match. If you read through Sec C, then you may find that $x_k$ acts as the anchor point. This can hold for the atom nearest to the center of $(x_i, x_j)$, or by the backbone structure of the protein, and the proof still holds. Now in the latest version, we just make $x_i, x_j, x_k$ in Eq 11 exactly the same as the one in Eq 3. The proof of other frames, as highlighted in green in the revised version, is straightforward.

2. Also, we didn’t claim that we follow a locally equivariant approach for building our method. We are wondering why you insist that we should locally learn equivariant features? Making our method local is quite straight, we only need to move the local mass center of a subgraph (vectors in each frame) to zero, but the motivation for doing so is unclear to us.

Q: Clarification of algorithm and metric

Thank you for asking these questions. There are some gaps in the discussion, and we are happy to clarify further.

• Both the learning and inference of NeuralMD-SDE are stochastic, and it is learning stochastic dynamics. We added the random noise and then learned the stochastic simulator. The previous revision contained the words 'obtain trajectories' which is confusing, and we have modified it to 'train and sample trajectories`. Please check Sec 3.3 in the latest version, From Newtonian dynamics to Langevin dynamics.
• We are not certain about what this sentence means: however, these are unconvincing as the evaluation metrics are still based on the assumption that one wants to learn an MAE-minimizing simulator. Yet, we have tried our best to provide our understanding to you, as follows.
• All of the works along this research line use MAE as the loss function (related works section), and we added stability metric, following Xiang et al.’s most recent benchmark paper on ML for MD simulation. If you think these are unconvincing, can you provide more insightful explanations on why this is so? This would be appreciated, and then we can know how to solve this issue.

Q: About RMSF

First thank you for agreeing that RMSF is a metric, not a baseline. In the previous round of discussion, we explained why using RMSF is not ideal as a single-trajectory metric, and why using stability is a good replacement. Feel free to check About performance and baselines above.

More binding datasets and more methods

• Thank you for raising this question. First, we want to highlight that we are not just using one dataset. As also confirmed by your comments, we include 10 single binding trajectories and 3 datasets on multiple trajectories. These trajectories are from PDB, the most comprehensive dataset for protein (binding), thus, we aim to assert that the experiments conducted in this study are indicative and representative of a substantial volume of binding MD data.
• Then as also confirmed by your comments, the dataset is one limitation in AI for binding MD estimation, and MISATO is the first large-scale binding MD dataset released in May 2023. If you can provide other existing public large-scale binding MD datasets, we are happy to try them in future work.
• For baselines, if the reviewer can list the specific papers, we are happy to compare.

Model difference with existing works

The other models for binding MD trajectory prediction can be categorized into two venues:

• No incorporation of physics rules, like GNN-MD.
• Inappropriate incorporation of physical rules, like overdamped Langevine dynamics in the DenoisingLD papers. We provide more details in Sec G of the revised manuscript. Feel free to check.
• Thus, in comparison, we can better incorporate the physical rules (Newtonian dynamics and the newly added Langevin dynamics), and such inductive bias can help better model the binding MD.

Dataset simulation illustration is out of the scope of this work

We confirmed that MISATO uses the Amber20 software suite and NVT, and we provide more dataset statistics (number of atoms and number of residues, etc.) in Sec D. However, we want to highlight that doing simulation is beyond the scope of our work, because we are not the authors of MISATO. The credit for the binding MD dataset generation goes to the MISATO paper.

Potential limitation One potential limitation, which has been raised by reviewer oR9H is that we didn’t consider the solvent molecules in the MD system. This is because the public dataset does not include this information. Meanwhile, we are still able to handle such a limitation by adding random terms into the predicted forces, i.e., the Langevin dynamics. We have added the results of single trajectories below. Further in the revised manuscript, we added one paragraph in Sec 3.3 for a detailed explanation and the results of multiple trajectories in Sec 4. We can observe that the performance of ODE and SDE are very close to each other.

Practical implications

The experiments of binding MD estimation on MISATO (a subset of PDB) have practical implications. If we want to extend this to other real-world applications, that means we need other binding MD datasets, and we are happy to check if the reviewer can provide specific references.

Effectiveness of multi-grained

The proof of this is trivial because, without the multi-grained modeling, the model cannot fit into the GPU memory.

Multi-level interactions

So we have the multi-grained modeling discussed in Sec 3.1 & 3.2.

• In Sec 3.1, we have shown three levels of vector frames in Eq 2, 3, 4.
• In Sec 3.2, we have shown how to do MPNN on these three vector frames, respectively.
• Finally, the atom forces are used to get the trajectory following Newtonian dynamics, $a = F/m$.
• This is also illustrated in Figure 4.
• We added more details in Sec E of the revised manuscript.

Thank you for acknowledging our work as the first to explore large-scale binding MD datasets, solid performance, and efficiency. We believe that we have attempted to solve your concerns during the rebuttal and in the revised manuscript. We have listed the details below.

Dataset limitations

The binding dynamics dataset issue is out of the scope of our work. We also want to mention that this is actually the most critical bottleneck that prevents ML researchers from developing tools for domain-specific tasks. If there exist any other public binding MD datasets, we are happy to use them.

Efficiency with other DL models

Thank you for raising this. We have now added them below (FPS, corresponding to Table ). The main bottleneck of using ML for binding MD simulation is on the backbone model, and since we are using the same backbone model in our experiments, i.e., the running time of these DL models is on the same order of magnitude. That’s why we didn’t highlight the time comparison among different DL models in the first version.

How to model the protein-ligand interactions

This is an important module, and we have shown it in Sec 3.1 and 3.2 in the first version. To be more specific:

• In Eq 4 (Sec 3.1), we illustrated how the vector frame ($F_{complex}$) is constructed for protein-ligand interactions.
• In Eq 6 (Sec 3.2), we showed the MPNN module for protein-ligand binding, based on $F_{complex}$.
• We also showed the pipeline in Figure 4 & 5 (appendix).
• We discussed how the interacted forces can be used for the binding MD estimation in Eq 6 and the following paragraph (Sec 3.2). So the output is atom-level forces, which will be fed into the Newtonian dynamics (Eq 8) for trajectory estimation.
• We added the model architecture in more detail in Sec E in the revised version.

We appreciate your recognition of our work on physics-inspired architecture, comprehensive evaluation, and efficiency. Your primary concerns are about backbone models and ablation studies. We believe that the following explanations adequately address the points of concern you raised.

Ablation studies on different components

Thank you for raising this question. We assume that your question is about if we can test the three components of BindingNet separately, as shown in Figure 4. The answer is no, because separating them would destroy the physical rules in MD. Let us explain in more detail:

• The BindingNet-Ligand models the interactions within the ligand itself.
• The BindingNet-Protein and BindingNet-Complex learn the atom-level forces induced by the proteins.
• By combining these modules, we can learn the atom forces from all the other atoms, which is the physical-inspired MD estimation.
• On the other hand, we also added an ablation study on introducing stochastic terms induced by the solvent molecules (they were considered in the dataset generation, but not provided to the public). We have attached the results in the revised manuscript. Please feel free to check it.

Can we replace BindingNet with Equiformer and EGNN, and why multi-level?

These two questions are highly connected, so we answer them together. In practice, the multi-grained modeling is appealing from the engineering perspective. If we model the large molecular systems, even for one ligand-protein complex, there are on average 1k residues (around 20k-30k atoms) in the MISATO dataset. Thus, modeling such a big molecular system is challenging for GPU memory. This means:

• We cannot replace BindingNet with Equiformer or EGNN.
• The multi-grained (or multi-level) modeling is motivated by reducing the computational cost, and without it, the GPU memory cannot hold for such a large molecular system.
• Actually, even in our current system, we are taking batch size with 2 so as to fit the GPU memory.

Clarification on evaluation

Yes, both the MAE and MSE are calculated over the whole trajectory snapshots in the test set, not only the final snapshot. We have this highlighted in Sec 4.

Hi reviewer 582T,

Thank you for the follow-up. We are happy to address your questions in more detail.

1. We list the main reasons why EGNN is not a good fit below.

   • [Memory] We agree that EGNN is efficient, but the bottleneck here is the computational memory. (The original EGNN is connecting all atoms)
   • [Math correctness] EGNN is E(3)-equivariant. However, when the protein is fixed, the modeling of the ligand needs to be SE(3)-equivariant and reflection-antisymmetric. We explained more details in Sec C.
   • [Performance] Further on the performance, the Geom3D benchmark has provided around 50 geometric tasks, showing how the SE(3)-equivariant models outperform EGNN / GVP.

2. Thank you for raising this question, and let's explain more insights to you.

   • If we focus on the single-trajectory setting (Table 1), then you can see that the MAE and MSE are comparatively small. This is also the setting that most baseline papers are working on.
   • In this paper, we introduce another challenging task of generalization among multiple trajectories (Table 2), this is where you see that the MAE & MSE are much larger. Notice that here, the MAE & MSE are the averages of all the snapshots and all the testing trajectories.
   • Thank you for raising this question on RMSF. RMSF is the metric of a single trajectory, and if we want to compare the fluctuation between the predicted trajectories and ground-truth trajectories, then that is the stability metric. In other words, we have already included the RMSF on prediction-and-true trajectories.

Hope this answers your questions, and we are happy to provide more insights.

I would like to thank the reviewers for their response. I have two comments:

1. I would probably disagree that one cannot replace BindingNet with EGNN because we have been training EGNN internally on the MISATO dataset and it runs pretty fast. Due to the lack of baselines, it is a bit unclear whether BindingNet is superior than existing geometric encoders.

2. I initially thought that MAE/MSE was calculated at the last snapshot. If it was calculated over the whole trajectory (every snapshot), I would imagine the MAE/MSE to be substantially lower (currently MAE is around 8.1 for MISATO-all) because a ligand wouldn't move that much during MD simulation. What's the MAE if you just use the ligand pose at time 0? (meaning, the ligand doesn't move at all during MD simulation). This is probably related to RMSF metric the other reviewer brought up

Hi Reviewer 582T,

Thank you for the follow-up!

1. Yes, we are definitely going to explore more SE(3)-equivariant models in the future. BTW. We also share some insights that another motivation for using vector frame modeling is for certain tasks (e.g., protein representation), the vector frame modeling outperforms other methods by a large margin. You can check this CDConv paper in case interested.

2. Thank you for the suggestion. So we called this baseline Void MD since we are just taking the first frame as the prediction for the rest frames for each trajectory. We got some interesting results, as shown below. As you can see, our NeuralMD is still reaching the best performance, while all the existing/published baselines are worse than Void MD. We also acknowledge the marginal performance improvement, but we want to point out the challenge of this task, and we reach a consistent improvement. This is a promising sign to prove the better generalization ability of NeuralMD, especially compared to the published baselines.

We hope this answers your questions, and any re-evaluation is appreciated.

Regards,

Authors of NeuralMD

We appreciate your positive comments on our work. Your primary concerns are about the dataset and evaluation, and we have explained them below.

1. [regarding the dataset limitation] We agree that the dataset is one of the potential limitations here; as highlighted in the conclusion section, making more datasets available may need help and support from the whole community. Creating such a dataset not only is very costly, but also requires sufficient domain expertise from different domains, which makes it challenging. As stated in the paper, MISATO dataset is the only publicly available large-scale dataset, and that is why we empirically showed the potential of our NeuralMD method using this dataset. Regarding its effect on femtosecond MD, we unfortunately can’t perform any empirical studies at this stage, although we theoretically show that our method is agnostic to the MD time interval.

2. [regarding quantitative evaluation] Yes, we followed the stability metric from previous work, and reported those quantitative results in Tables 1 & 3.

Hi Reviewers K9WG, 582T, ZPQv, oR9H,

Since the deadline is approaching today, should there be any remaining unclear statements, please don't hesitate to inform us. We'll make every effort to address your concerns.

Regards,

Authors of NeuralMD