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

We thank you for your recognition of our work’s contribution and clear organization. We will address your questions and concerns as follows.

Weakness 1, Ablation Study of Atom2Mesh & Mesh2Atom and Efficiency of FFT

We provide additional ablation studies on the impact of Atom2Mesh & Mesh2Atom modules in the following table under settings for training variants of SchNet on OE62 dataset. It can be shown that both modules indeed contribute to the final performance in a synergic fashion (Introducing just a single module results in poorer results).

Regarding efficiency, we would like to highlight that our approach, which incorporates mesh points, draws inspiration from the FFT applied to Ewald summation. In practice, we employ the FNO as the model, which allows for efficient computation. While the inclusion of mesh nodes does introduce a larger number of parameters, this is offset by the significant reduction in runtime compared to Ewald MP. This improvement is largely due to the efficiency afforded by the FFT. For a more detailed analysis, please refer to our profiling results presented in Table 3 in the Global Rebuttal and Table 2 in our manuscript.

Weakness 2, More Explanation of the Parameterizing Long-Range Interaction

Sorry for the lack of clarity. We elaborate more on our choice of using FNO for capturing the long-range interactions. Eq.9 indicates that the energy of long-range interactions can be regarded as the result of the convolution between the influence function $G$ and the charge density $\rho_M$. According to the convolution theorem, we can accelerate the convolution in the Fourier domain by point-wise multiplication. If we consider the charge density $\rho_m$ as a representation of mesh $m^l$ and parameterize $G$ directly in the Fourier domain $\tilde{G}$, we can obtain the latter part of Eq.17. As for the first part, it is merely a gated connection.

Weakness 3, Ablation Study of Assignment Cutoff

We provide additional ablation studies on the impact of the assignment cutoff distance as follows under settings for training SchNet on OE62 dataset. For the performance (69.1) reported in the paper, we use a combination of radius graph and k-NN graph, setting the maximum number of neighbors to 5, which generally results in minimal multiple assignments. To test the effect of multiple assignments, we use only the radius graph in this ablation experiment, and we observe that all the experiments perform slightly worse due to multiple assignments. However, an appropriately chosen cutoff (4 or 5 Å) still yields relatively optimal results. Notably, the results do not deteriorate further as the assignment cutoff increases. We hypothesize that this may be because a larger assignment cutoff creates a larger neighborhood environment, making it easier to learn the assignment function with the fixed number of meshes, thereby mitigating the multiple assignments problem.

Questions 1

See Weakness 3 for details.

Questions 2

See Weakness 2 for details.

Questions 3

See Weakness 1 for details.

Limitation 1, Extending to DNA/Protein

Thanks for your insightful suggestions. We acknowledge that we mainly focus on testing the effectiveness of our Neural P$^3$M framework in this work so we primarily experimented with molecules with abundant datasets and baselines. We noted that, currently, energy/force datasets for large biomolecules are scarce, as they often require expensive ab initio calculation. We will leave this scaling up as one of the future directions of this work.

We thank your recognition of our work’s novelty, contribution and clear organization. We will address your questions and concerns as follows.

Weakness 1, Evaluation & Improvement

We resolve the issues point-to-point in the Questions Section.

Weakness 2, Anonymized Code

We intend to make our code publicly available upon acceptance of our manuscript. For the purpose of this review process, we have provided an anonymized version of the code https://anonymous.4open.science/r/Neural_P3M-1552.

Question 1, ViSNet as the Baseline

We selected ViSNet as our base model for MD22 because LSRM demonstrated better performance using ViSNet than Equiformer across molecules. This allows us to make a direct and equitable comparison between our results on ViSNet and those reported by LSRM. Additionally, we present further results for Equiformer on MD22 dataset. The incorporation of Neural P$^3$M has enhanced the performance of Equiformer, even surpassing that of ViSNet (shown in the following table). Given the constraints of resources, we are presenting results only for the two largest molecules on MD22 at this moment. We will provide a complete evaluation if our work gets accepted.

Question 2, Choice of Mesh & Cutoff

As the molecules in MD22 range from peptides to nanotubes with diverse molecule sizes, each individual molecule is a separate dataset. We chose the number of mesh points and the cutoff distance accordingly in a pre-defined and consistent manner. Actually, we didn't tune these hyperparameters due to the vastness of the search space. In practice, we typically set the assignment cutoff distance at 4.0 or 5.0 Å, ensuring that the product of the number of mesh points and the cutoff is roughly equivalent to the cell size in each dimension. For example, the tubular nanotube would need more mesh points along its longest dimension $x$. For a dataset with larger diversity containing different molecules, OE62 is such an example, for which we have demonstrated the performance improvement of Neural P$^3$M using the average cell size following the same manner.

Question 3, LSRM on OE62

We initially did not include this result because of the diversity of OE62, as the fragmentation algorithm employed in LSRM isn't suitable for all molecules in OE62. Nonetheless, by filtering some molecules (validation:285, test:293, train:2387) and utilizing a marginally different dataset, we have also presented the performance of LSRM on OE62 dataset (See Table 1 in the Global Rebuttal). These additional results suggest that LSRM surpasses the baseline, yet it does not perform as well as our Neural P$^3$M.

Question 4, GPU Memory Usage

See Table 3 and its discussion in the Global Rebuttal for more details.

Question 5, Ablation Study of Atom2Mesh & Mesh2Atom

We provide additional ablation studies on the impact of Atom2Mesh & Mesh2Atom modules in the following table under settings for training variants of SchNet on OE62 dataset. It can be shown that both modules indeed contribute to the final performance in a synergic fashion (Introducing just a single module results in poorer results). The cost of two modules is very little and the bottleneck is in the short-range model, so there is no significant speedup.

We thank you for your recognition of our work’s novelty and contribution to the field of molecular modeling. We will address your questions and concerns as follows.

Weakness 1, Mathematical Details

Thank you for your suggestions. We have relocated some non-essential content from Sections 2 and 3 to the appendix to better organize the paper. Additionally, we intend to release the code upon acceptance of our work to make implementation easier for all researchers.

Weakness 2, Pseudocode

Thank you for your suggestions. We have included pseudo-code in the appendix to further illustrate our method. However, due to the constraints of OpenReview, it is not feasible to display it directly. For your convenience, we have also provided anonymized reference code (https://anonymous.4open.science/r/Neural_P3M-1552) to facilitate better understanding of our work.

Weakness 3, Computational Complexity

Thank you for your suggestions. We detail the memory usage in Table 3 in the Global Rebuttal. The bulk of the memory usage is still attributed to the short-range modules—for instance, 16719 MB versus 19945 MB in GemNet. This level of memory consumption is considered acceptable in light of the performance gains achieved.

Weakness 4, Typos in Table 4

Thanks for pointing these out. We have fixed them in the revised manuscript.

Question 1, Restriction on GNN

We have integrated Neural P$^3$M with a wide range of geometric GNNs in our experimental settings. This included both classic GNNs like SchNet and DimeNet and more recent and advanced equivariant architectures like PaiNN, GemNet, Equiformer, and ViSNet. In this way, our framework is a general one that can be integrated into most geometric GNNs.

Question 2, Computational Bottleneck

We have provided the empirical evaluation of the runtime in the following table. The runtime for the long-range interactions remains consistent, whereas the short-range interactions are the primary bottleneck, particularly as the model for these interactions grows in complexity. Efforts to reduce computational demands, such as implementing the density trick to reduce the cost of many-body interactions, have been investigated in models like ViSNet [1] and MACE [2]. Additionally, simplifying the complexity of the CG-product is another avenue being explored to speed up short-range interaction models [3].

[1] Y, Wang T, Li S, et al. Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing[J]. Nature Communications, 2024, 15(1): 313.

[2] Batatia I, Kovacs D P, Simm G, et al. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields[J]. Advances in Neural Information Processing Systems, 2022, 35: 11423-11436.

[3] Passaro S, Zitnick C L. Reducing SO (3) convolutions to SO (2) for efficient equivariant GNNs[C]//International Conference on Machine Learning. PMLR, 2023: 27420-27438.

Question 3, Freezing Common Layers

Following your recommendations, we evaluated the performance of our base GNN when it is first pre-trained and subsequently fine-tuned under settings for training variants of SchNet on OE62 dataset. The results are presented in the following table. We observed that freezing the pre-trained weights results in a notable decline in performance. Conversely, permitting fine-tuning of the weights yields a marginal improvement. This could be attributed to the possibility that the energy is an amalgamation of long-range and short-range interactions, suggesting that the short-range representations, when trained directly with energy labels, may not be ideally suited for direct integration with long-range components. Utilizing pre-trained short-range models is an intriguing area for future research and is something we plan to explore further.

Limitations 1, Distributed Training

Our framework is built on top of PyTorch Lightning, which allows for straightforward extension to Distributed Data Parallel (DDP) mode. We have conducted experiments under the setting for training SchNet + Neural P$^3$M on the large OC20-2M dataset with 4 GPUs, and our results demonstrate significant improvements, underlining the robustness of our framework.

Limitations 2, Profiler results

See Table 3 and its discussion in the Global Rebuttal for more details.

We thank you for your recognition of our work’s superior performance and clear organization. We will address your questions and concerns as follows.

Weakness 1, Novelty & Contribution of Neural P$^3$M

While it's true that FFT is commonly employed in traditional chemical computations, as discussed in Section 2. However, incorporating such a technique as a learnable component within a Geometric GNN framework is far from straightforward. To the best of our knowledge, we introduce the novel mesh concepts for energy and force prediction, marking one of our major contributions. In terms of model architecture, we have also introduced the innovative learnable Atom2Mesh and Mesh2Atom modules, which are designed to enhance the exchange of information between short-range atomic and long-range mesh representations. With the integration of mesh nodes, FFT emerges as a natural choice for modeling long-range interactions. However, we could also opt for alternative networks such as transformers or 3D CNNs. All of these carefully designed components set Neural P$^3$M apart from existing Ewald-based models like Ewald MP. Thus, our framework not only expands upon Ewald-based approach but also paves the way for a new class of mesh-based methodologies in the modeling of long-range interactions in 3D molecular structures.

Weakness 2, Additional Experimental Results

Thanks you for your suggestions. We've implemented LSRM in our experiments on OE62 dataset as well as in Ewald MP applied to ViSNet on MD22 dataset (across a total of 11 experiments). Details can be found in Tables 1 and 2 in the Global Rebuttal. While both Ewald MP and LSRM outperform the baseline, Neural P$^3$M consistently delivers superior performance in the majority of the test baseline models (OE62) and molecules (MD22).

For the memory consumption, please see Table 3 and its discussion in the Global Rebuttal for more details.

Question 1, Concepts in Table 2

The experimental results of different variants of the base models came from the original Ewald MP [1] paper. “Embedding” indicates a model with a larger embedding dimension, “Cutoff” for a larger cutoff distance, and “SchNet-LR” for models with the pairwise LR block. These variants are existing improvement methods for the base model.

[1] Ewald-based long-range message passing for molecular graphs. ICML 2023

Question 2, Computational Complexity

The computational complexity for the long-range interactions scales with $O(K\log K D^2)$, where $K$ represents the number of mesh points and $D$ is the dimensionality of the hidden layers. The complexity of the short-range interactions is contingent upon the chosen GNN architecture, typically scaling with $O(|\mathcal{E}| D^2)$, where $\mathcal{E}$ denotes the set of edges. The overhead introduced by the Atom2Mesh and Mesh2Atom modules is negligible when compared to the computational demands of the long-range and short-range components. Notably, the computational bottleneck is primarily determined by the selected number of mesh points and the volume of atom edges. In practical applications, we observed that setting $K$ to $3^3=27$ yielded a significant enhancement in performance relative to the baseline model.