Neural P$^3$M: A Long-Range Interaction Modeling Enhancer for Geometric GNNs

Summary

The accurate modeling of both short-range and long-range interactions in molecular systems is crucial for predicting properties like molecular energies and forces with high precision. However, traditional Geometric Graph Neural Networks (GNNs) fail to capture such interaction. The paper introduces Neural $\text{P}^3$M, a framework that enhances the capabilities of GNNs via meshing up the Ewald summation. Compared to previous Ewald-based GNNs, the Neural$\text{P}^3$M further enhances the model by incorporating mesh and atom representation in a trainable manner. The paper conducts extensive empirical studies and justifies the effectiveness of the proposed method.

Strengths

- Unlike the Ewald method, which primarily handles electrostatic interactions, Neural P3M is designed to enhance the modeling of both short-range and long-range interactions in molecular systems. - The Atom-Mesh interaction mechanism employed by Neural $\text{P}^3$ M, along with the Fast Fourier Transform (FFT), is generally faster than the traditional Ewald method, especially for larger systems. The efficiency of such a design is demonstrated by the comparison of the running time in Table 2. - The paper is generally well-written, with clear illustrations and tables.

Weaknesses

- The major contribution of the Neural $\text{P}^3$ M seems to lie in the representation assignment. However, the paper does provide enough ablation studies on the effectiveness of this component. It is unclear whether changing the Fourier Transform to FFT is the main reason for the efficiency. - The introduction of different concepts can be further improved. For example, the detailed introduction of the Ewald summation in Section 2 mostly is not explicitly ref in the following section. However, Equation 17, which is crucial for the Neural $\text{P}^3$ M, is not elaborated properly, especially for the influence function G and the connection to Equation 9. - Despite the parameterized assigning function, the representation assignment still assigns each atom to multiple mesh points when the number of mesh points increases. This would post drawbacks on the performance and forward time. The author could consider discussing the different choices of the distance cut-off, as this hyperparameter could differ across datasets.

Questions

1. Discuss the different choices for the distance cut-off. 2. Could the author explain more about the connection between Equation 17 and Equation 9 and the concept of the parameterizing strategy of influence function G? 3. Provide ablation studies on the proposed component and the discussion on the effectiveness of the FFT.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

It is interesting to involve mesh points to help capture the long-range interaction. The author could extend the idea to other molecules like protein or DNA.

Summary

This work introduces a long-range focused GNN that utilizes the combination of atom and mesh representations. The mesh framework in this work is trainable and unconstrained to the fragmentation algorithm. Results demonstrate superior performances across MD22, Ag, and OE62 datasets.

Strengths

- The need for long-range GNN methods is crucial for various molecular representation tasks - This work tackles this pertinent issue utilizing a novel combination of atom and mesh representation - The underlying method is well-described - The distinction of this approach to other recent and relevant approaches is well-described - The paper is well-written and easy to walk through - Strong results - Demonstrates improvement over vanilla VisNet on Ag dataset - SOTA results on most targets and molecules of MD22 and gives significant improvements on larger molecules - On OE62 dataset, NeuralP3M performs better on all architectures over other approaches. It’s also faster than Ewald (for most cases) due to FFT which is impressive.

Weaknesses

- Few of the results seem carefully designed to demonstrate improvements through this approach which makes the overall impact unclear. I've mentioned more concrete points related to this as questions and I can be convinced otherwise post rebuttal/discussion. - The anonymous link to the code is not in the paper.

Questions

- What’s the reason for choosing VisNet to add Neural P3M in Table 1 instead of Equiformer/MACE which has better baselines? - MD22 results seem to have very carefully designed hyperparameters (cutoffs and number of meshes in each direction) in Table 4 of Appendix E.2. Do you have thoughts on the feasibility of finding these hyperparams on datasets with larger diversity? Additionally, what do the results look like if you just have a uniform hyperparameter across all molecules? - Is it possible to add results with LSRM in Table 2 results as well? Was there a reason for not including that? - In Table 2, how does the GPU memory usage compare across these methods? Just out of curiosity, I wanted to know what were the max. batch sizes you were able to fit across these methods. That would most likely make throughput comparison with Ewald even more impressive (assuming a smaller batch size fit for Neural P3M) - How important is it to have both atom2mesh and mesh2atom modules? Do you get similar performances by just having one and maybe some speedup? Some of the prior literature in charge density models only have mesh2atoms message passing.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Yes, this work does mention its limitations and potential negative impact.

Summary

The paper introduces Neural P3M, a framework designed to enhance geometric GNNs by incorporating mesh points alongside atoms and transforming traditional mathematical operations into trainable components. The mesh representations offers discrete resolutions necessary for formulating long-range terms. The Neural P3M is also efficient due to the reduced computational complexity afforded by FFT. The paper starts with highlighting the importance of long-range terms, which are absent and inefficient in previous works. The paper then explain the preliminary of Ewald summation and the meshing up methods with detailed formulas. The novel methods and neural network blocks are presented in the next section and evaluated with various models and datasets. The experiment results show significant improvements. When integrated with ViSNet, Neural P3M achieves state-of-the-art performance in energy and force predictions across several large molecules, outperforming other leading models. The framework, combined with models like SchNet, PaiNN, DimeNet++, and GemNet-T, demonstrates enhanced performance and faster computation times compared to related works. Neural P3M provides a robust framework for enhancing geometric GNNs, enabling them to capture long-range interactions efficiently. The framework's adaptability to various molecular systems and its demonstrated performance improvements on key benchmarks make it a significant contribution to the field. The study also highlights areas for future research, such as optimizing the number of mesh points and exploring alternatives to FFT for modeling long-range interactions.

Strengths

1. Neural P3M effectively integrates mesh points alongside atoms, which allows it to capture long-range interactions more accurately than traditional GNNs. This enhancement addresses a significant limitation in current molecular modeling approaches, particularly for large molecular systems. 2. The framework is built upon well-established principles such as Ewald summation and P3M methods. This theoretical grounding lends credibility to the approach. 3. Neural P3M is designed to be a versatile enhancer that can be integrated with a wide range of existing geometric GNN architectures, including SchNet, PaiNN, DimeNet++, and GemNet-T. This compatibility ensures that the framework can be widely adopted and used in different contexts. 4. Neural P3M reduces the computational complexity of long-range interaction calculations, making it feasible to handle large-scale systems efficiently. The framework also exhibits a faster computation time than Ewald MP. 5. Its theoretical soundness, empirical success, and detailed implementation make it a valuable contribution to the field of molecular modeling.

Weaknesses

1. The framework's reliance on complex mathematical operations and integration of mesh-based methods with GNNs can make implementation challenging. Researchers and practitioners may require significant expertise in both GNNs and numerical methods to effectively utilize Neural P3M. It is better to remove some unnecessary equations in Section 2 and 3, or move them to the appendix or references. 2. Additionally, you can present the P3M Blocks by some pseudocode. 3. The framework's need to handle both atomic and mesh representations simultaneously may lead to increased memory usage, which could be a bottleneck for handling large datasets or systems with limited hardware capabilities. You can present the GPU memory usage while training model with or without Neural P3M. 4. There are some trivial mistakes in Table 2. If the higher Rel. is better, the up arrows should be used. The best runtime should also be highlighted.

Questions

1. Does Neural P3M have some restrictions of geometric GNN models? Or it can be combined with most GNNs? 2. What are the throughputs of Short-Range Block and Long-Range Block? Is the Long-Range Block slower than Shaor-Range Block? Which block is the main bottleneck or is it possible to improve the performance? 3. The variants of the same model usually share some common layers or blocks. Is it possible to reuse or frozen some common layers in a pre-trainded model, and fine-tune the newly added blocks, such as Long-Range Block? So the training performance will be further improved.

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

1. Distributed training is important and efficient while training the general models with large datasets. However, the distributed training of Neural P3M is not evaluated, and the proposed models and results are now limited to one GPU. 2. Profiler results and the number of parameters in each variant are not presented.

Summary

This work proposes Neural P3M, a framework that enhances geometric GNNs by integrating mesh points and leveraging Fast Fourier Transform (FFT) for efficient computation of long-range interactions. The framework includes short-range and long-range interaction modeling and enables the exchange of information between atom and mesh scales. Neural P3M improves the prediction of energies and forces in large molecular systems, achieving good performance on benchmarks like MD22 and OE62.

Strengths

The proposed framework is capable of being incorporated in short-range geometric GNNs, although distinct integration strategies are needed due to the varying inputs and outputs of different models. The improvements in benchmarks are promising, which demonstrates the power of the proposed model. The paper is well-structured and clearly written, making it accessible to readers without knowledge of the related concepts like Ewald summation.

Weaknesses

1. I feel the overall impact and novelty of this work are limited, given that several works have adopted Ewald summation in geometric GNNs. This work enhanced this concept by introducing FFT for accelerated Ewald summation, which is a common way in traditional simulations and is also already mentioned as a possible direction in [1]. 2. The experimental part is not comprehensive. The experiment on MD22 doesn't have a comparison to [1], and LSRM is not being compared on OE62. Besides, it is important to also compare the memory consumption between different approaches when dealing with long-range interactions, since in many real-world problems we hope to capture long-range interactions in large molecular systems. [1] Ewald-based long-range message passing for molecular graphs. ICML 2023

Questions

1. What do the "Embeddings", "Cutoff", and "SchNet-LR" mean in Table 2? 2. What is the computational complexity of the proposed method?

Rating

3

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors have addressed the limitations and potential societal impacts.

Response

I thank the authors for their rebuttal. Overall, all my raised concerns have been addressed. As a result, I've increased my score.

Gratitude for Your Review

We really appreciate your recognition of our work and your decision. Thank you for the effort and time you put into your review!

We are looking forward to your reply

Dear Reviewer M1Fm, Thank you for your insightful feedback on our manuscript. **As the deadline for discussion nears**, we wish to remind you that we have provided **a comprehensive response** to address the concerns you raised. With respect to the novelty of our work, we would like to highlight that the other reviewers have recognized our contributions and novelty. It seems there may have been a **misunderstanding** regarding this during the initial review, and we highlighted our novelty in the previous rebuttal. Furthermore, in response to your comments on our experiments, we have made the necessary enhancements. If you believe we have resolved the issues, we would be grateful if you could reconsider your score. We welcome any further questions or discussions you may wish to have. Warm regards, The Authors

We are looking forward to your reply

Dear Reviewer 7zWV, We are thankful for your valuable feedback and the recognition you have given our manuscript. **As the deadline for discussion nears**, we would like to gently remind you of the thorough response we have crafted to address the issues you highlighted. We have carefully addressed **all the points you raised in the Weaknesses**, and have conducted **additional experiments** to answer your interesting questions. We find your questions heuristic and are happy to further engagement on these topics. If you believe we have resolved the issues and adequately answered your questions, we would be grateful if you could reconsider your score. We welcome any further questions or discussions you may wish to have. Warm regards, The Authors

Conclusion for Author-Reviewer Discussion Period

Dear Reviewers, We sincerely thank you for your insightful and constructive feedback throughout the review process. Your thoughtful comments have helped make our work more comprehensive and concrete. We are grateful for your recognition of our efforts, contributions, and novelty in our Neural P$^3$M framework, and we are glad that our rebuttals with additional 28 experiments and profiling results have successfully addressed your questions and comments. We especially enjoyed the interactive discussions during the rebuttal period and we always look forward to the possibility of engaging in a productive dialogue with the reviewers. We are grateful to three reviewers for their active engagement in the discussion period and for their unanimous recognition of our contributions. We will make sure to add the additional experiment results and ablation studies in our revised manuscript. Again, we thank you for your valuable input and for helping us bring this research to fruition. Warm Regards, The Authors