Paper: A Multi-Grained Group Symmetric Framework for Learning Protein-Ligand Binding Dynamics (ICLR-2024)

Reviewer comment (Theory):
Problem formulation. Although learning a neural ODE for long-timescale dynamics is an interesting idea, the deterministic problem formulation seems inappropriate as MD simulation itself is not necessarily deterministic. Thermostats and barostats typically introduce stochasticity, and even in their absence the removal of water molecules (which are explicit in the MISATO simulations) injects intrinsic uncertainty into the modeling problem. Hence, the dynamics are fundamentally stochastic and I am skeptical that any neural ODE can faithfully capture the long-timescale dynamics of these protein-ligand systems.

Author response:
**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.


|                     | 100-MAE            | 100-MSE            | 100-Stability       | 1k-MAE            | 1k-MSE            | 1k-Stability       | All-MAE            | All-MSE            | All-Stability       |
|---------------------|----------------|----------------|-----------------|----------------|----------------|-----------------|----------------|----------------|-----------------|
| VerletMD            | 90.326         | 56.913         | 86.642          | 80.187         | 53.110         | 86.702          | 105.979        | 69.987         | 90.665          |
| GNN-MD              | 7.176          | 4.726          | 35.431          | 7.787          | 5.118          | 33.926          | 8.260          | 5.456          | 32.638          |
| DenoisingLD         | 7.112          | 4.684          | 29.956          | 7.746          | 5.090          | 18.898          | 15.878         | 10.544         | 89.586          |
| NeuralMD-ODE | **6.852** | **4.503** | **19.173** | **7.653** | **5.028** | **15.572** | **8.147** | **5.386** | **17.468** |
| NeuralMD-SDE  | 6.869          | 4.514          | 19.561          | 7.665          | 5.037          | 16.501          | 8.165          | 5.398          | 19.012          |

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

Reviewer comment (Evaluation):
Performance and baselines. Even in the single-trajectory setting, the MAE for the ligand coordinates seems rather large, in the range of 2-6 angstroms. This is perhaps due to the suboptimal problem formulation already discussed. However, the numbers provided are not meaningful because we do not know what the RMSF is in these simulations. The RMSF is the best result that can be achieved by a single static structure and is an essential missing baseline. The MAE in the multi-trajectory setting is even larger (7 angstroms) and suggests a complete dissociation of the ligand from the binding pocket. Finally, the stability metrics are only marginally better than DenoisingMD and do not represent a qualitative resolution of the problem.

Author response:
**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.

Reviewer comment (Theory):
Mathematical errors. Although the paper places much emphasis on SE(3)-equivariance, the construction of the atom-level and residue-level frames appears to be non-equivariant. Specifically, the cross product is not translation equivariant since  
$$(x_i + t) \times (x_j + t) = (x_i \times x_j) +  (x_i - x_j) \times t$$  
is not a function of $(x_i\times x_j)$ and $t$. Noticeably, the appendix claims the cross-product to be SE(3)-equivariant but only establishes rotation equivariance. Further, while I agree that the local frame construction described in the appendix is equivariant, this is different from what is done in the main text, since there is no nearest-neighbor atom $x_k$ and absolute positions $x_i, x_j$ are used.

Author response:
**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).

Reviewer comment (Writing):
The architectural details in Eq 5 and Fig 4 are very unclear, with several ambiguous uses of the dot product. Please label the embeddings with the dimensionality (so it is clear if they are scalars or vectors) and define symbols before they are used (for example $h_i$). It is also not clear what kind of MPNN is used and what is the meaning of the number of "layers."

Author response:
- Clarifications on ML terminologies.
    - In this work, our usage of the dot product is the same, and you can check [this wiki page](https://en.wikipedia.org/wiki/Dot_product).
    - 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.

Reviewer comment (Evaluation):
Please clarify if by MAE you mean RMSD or some alternative definition of positional error (with or without a factor of $\sqrt{3}$).

Author response:
- MAE is short for ``mean absolute error’’. We have shown its equation in Sec 3.3 below Eq (8) when it first shows up.

Reviewer comment (Reproducibility):
"We keep the same backbone model (BindingNet) for energy or force prediction for all the baselines." If so, please provide significantly more details about how these baselines were retrained.

Author response:
- 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).

Reviewer comment (Experiments):
If this work considers the semi-flexible, how did you deal with the protein movement that is present in the MISATO MD trajectories?

Author response:
- 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.

Reviewer comment (Reproducibility):
What is your integration timestep? If it is adaptive, please provide more details about how many timesteps are required for a typical 8ns simulation.

Author response:
- The timesteps in the current experiment are fixed (not adaptive) so as to fit the MISATO dataset.

Reviewer comment (Experiments):
In the single-trajectory setting, how is the temporal division carried out?

Author response:
- Training on the first 80 snapshots, and testing on the last 20 snapshots. The 80th snapshot is given in inference.

Reviewer comment (Experiments):
This work is based on the first large-scale dataset with binding dynamics, which may have its own limitations.

Author response:
**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.

Reviewer comment (Experiments):
The paper does not provide a direct comparison with other state-of-the-art methods in terms of computational efficiency.

Author response:
**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.

| PDB ID          | 5WIJ   | 4ZX0   | 3EOV   | 4K6W   | 1KTI   | 1XP6   | 4YUR   | 4G3E   | 6B7F   | 3B9S   | Average |
|-----------------|--------|--------|--------|--------|--------|--------|--------|--------|--------|---------|
| VerletMD        | 12.564 | 30.320 | 29.890 | 26.011 | 19.812 | 28.023 | 31.513 | 29.557 | 19.442 | 31.182 | 25.831  |
| GNN MD          | 22.639 | 41.083 | 26.219 | 26.798 | 22.435 | 35.406 | 25.370 | 36.677 | 42.018 | 40.339 | 31.898  |
| Denoising LD    | 37.047 | 39.764 | 38.232 | 32.222 | 22.181 | 34.501 | 41.498 | 38.003 | 42.261 | 41.867 | 36.758  |
| NeuralMD (Ours) | 33.164 | 39.415 | 31.720 | 31.909 | 24.566 | 37.135 | 39.365 | 39.172 | 20.320 | 37.202 | 33.397  |

Reviewer comment (Writing):
The article briefly mentions protein-ligand binding dynamics but does not further explain their importance and how such dynamics can be modeled. In addition, the article does not adequately discuss the interactions between proteins and ligands and how these interactions can be incorporated into the simulations.

Author response:
**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.

Reviewer comment (Evaluation):
The paper does not discuss the potential limitations or drawbacks of the NeuralMD framework.

Author response:
**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.

| PDB ID | Metric    | VerletMD | GNN-MD    | DenoisingLD | NeuralMD-ODE | NeuralMD-SDE |
|--------|-----------|----------|-----------|-------------|--------------|--------------|
| 5WIJ   | MAE       | 9.618    | 2.319     | 2.254       | 2.118        | **2.109**    |
|        | MSE       | 6.401    | 1.553     | 1.502       | 1.410        | **1.408**    |
|        | Stability | 79.334   | 45.369    | 18.054      | **12.654**   | 13.340       |
| 4ZX0   | MAE       | 21.033   | 2.255     | 1.998       | **1.862**    | 1.874        |
|        | MSE       | 14.109   | 1.520     | 1.347       | **1.260**    | 1.271        |
|        | Stability | 76.878   | 41.332    | 23.267      | **18.189**   | 18.845       |
| 3EOV   | MAE       | 25.403   | 3.383     | 3.505       | 3.287        | **3.282**    |
|        | MSE       | 17.628   | 2.332     | 2.436       | 2.297        | **2.294**    |
|        | Stability | 91.129   | 57.363    | 51.590      | **44.775**   | 44.800       |
| 4K6W   | MAE       | 14.682   | 3.674     | 3.555       | 3.503        | **3.429**    |
|        | MSE       | 9.887    | 2.394     | 2.324       | 2.289        | **2.234**    |
|        | Stability | 87.147   | 57.852    | 39.580      | 38.562       | **38.476**   |
| 1KTI   | MAE       | 18.067   | **6.534** | 6.657       | 6.548        | 6.537        |
|        | MSE       | 12.582   | 4.093     | 4.159       | 4.087        | **4.085**    |
|        | Stability | 77.315   | 4.691     | 7.377       | 0.525        | **0.463**    |
| 1XP6   | MAE       | 13.444   | 2.303     | 1.915       | **1.778**    | 1.822        |
|        | MSE       | 9.559    | 1.505     | 1.282       | **1.182**    | 1.216        |
|        | Stability | 86.393   | 43.019    | 28.417      | **19.256**   | 22.734       |
| 4YUR   | MAE       | 15.674   | 7.030     | 6.872       | **6.807**    | 6.826        |
|        | MSE       | 10.451   | 4.662     | 4.520       | **4.508**    | 4.526        |
|        | Stability | 81.309   | 50.238    | 32.423      | **23.250**   | 25.008       |
| 4G3E   | MAE       | 5.181    | 2.672     | 2.577       | 2.548        | **2.478**    |
|        | MSE       | 3.475    | 1.743     | 1.677       | 1.655        | **1.615**    |
|        | Stability | 65.377   | 16.365    | 7.188       | **2.113**    | 2.318        |
| 6B7F   | MAE       | 31.375   | 4.129     | 3.952       | 3.717        | **3.657**    |
|        | MSE       | 21.920   | 2.759     | 2.676       | 2.503        | **2.469**    |
|        | Stability | 87.550   | 54.900    | 16.050      | **3.625**    | 22.750       |
| 3B9S   | MAE       | 19.347   | 2.701     | 2.464       | **2.351**    | 2.374        |
|        | MSE       | 11.672   | 1.802     | 1.588       | **1.527**    | 1.542        |
|        | Stability | 41.667   | 43.889    | 8.819       | **0.000**    | **0.000**    |

Reviewer comment (Evaluation):
The paper does not delve deeply into the practical implications or real-world applications of the proposed method.

Author response:
**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.

Reviewer comment (Experiments):
The paper states that the speed of this method is superior to standard numerical MD simulation methods, but this method is only compared with one method and is not compared with ML-based MD simulation methods. Please explain the reasons for this and compare it with more methods to prove the efficiency of the proposed method.

Author response:
**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.

| PDB ID          | 5WIJ   | 4ZX0   | 3EOV   | 4K6W   | 1KTI   | 1XP6   | 4YUR   | 4G3E   | 6B7F   | 3B9S   | Average |
|-----------------|--------|--------|--------|--------|--------|--------|--------|--------|--------|---------|
| VerletMD        | 12.564 | 30.320 | 29.890 | 26.011 | 19.812 | 28.023 | 31.513 | 29.557 | 19.442 | 31.182 | 25.831  |
| GNN MD          | 22.639 | 41.083 | 26.219 | 26.798 | 22.435 | 35.406 | 25.370 | 36.677 | 42.018 | 40.339 | 31.898  |
| Denoising LD    | 37.047 | 39.764 | 38.232 | 32.222 | 22.181 | 34.501 | 41.498 | 38.003 | 42.261 | 41.867 | 36.758  |
| NeuralMD (Ours) | 33.164 | 39.415 | 31.720 | 31.909 | 24.566 | 37.135 | 39.365 | 39.172 | 20.320 | 37.202 | 33.397  |

Reviewer comment (Experiments):
The title of the article mentions a "multi-grained group symmetric framework", but not enough experimental results were provided to prove its effectiveness.

Author response:
**Effectiveness of multi-grained**

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

Reviewer comment (Writing):
The methods section of the article mentions "BindingNet model that satisfies group symmetry using vector frames and captures the multi-level protein-ligand interactions" but does not explain in detail how this model captures the multi-level interactions. Could you provide more details or examples to illustrate this point?

Author response:
**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.

Reviewer comment (Experiments):
The proposed method only is evaluated on one dataset, which might be specific to the used dataset. The authors need to evaluate their method on other datasets and compare it with more methods.

Author response:
**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.

Reviewer comment (Reproducibility):
The article does not describe in detail the specific methods used for the experiments; for example, in the case of protein-ligand binding dynamics simulations, no specific information on the simulation software used, simulation conditions, model parameters, etc. is mentioned. This makes it difficult for the reader to understand and assess the feasibility of the experimental methods.

Author response:
**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.

Reviewer comment (Experiments):
This paper lacks ablation studies illustrating the benefit of different components.

Author response:
**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.

Reviewer comment (Evaluation):
The description of evaluation metric is not crystal clear (see questions below).

Author response:
**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.

Reviewer comment (Experiments):
Can you present ablation studies by replacing your BindingNet architecture with other existing geometric architectures such as EquiFormer and EGNN?

Author response:
**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.

Reviewer comment (Experiments):
Can you conduct ablation studies to understand the benefit of multi-level protein-ligand representation? I am not sure if having multi-level representation is helpful.

Author response:
**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.

Reviewer comment (Evaluation):
For the MAE and MSE metric, is the MAE/MSE calculated over the whole trajectory (every time step) or only the final step?

Author response:
**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.

Reviewer comment (Experiments):
One potential limitation of this work is the dataset. Currently, authors are using the MISATO dataset, a binding simulation dataset with a large timescale. However, NeuralMD is agnostic to the time interval, and it can also be applied to binding dynamics datasets with time interval as a femtosecond.

Author response:
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.

Reviewer comment (Evaluation):
Authors qualitatively show that NeuralMD reaches more stable binding predictions. Is there a way to show quantitatively as well.

Author response:
Yes, we followed the *stability metric* from previous work, and reported those quantitative results in Tables 1 & 3.