Neural P$^3$M: A Long-Range Interaction Modeling Enhancer for Geometric GNNs
Authors: Yusong Wang
Chaoran Cheng
Shaoning Li
Yuxuan Ren
Bin Shao
Ge Liu
Pheng-Ann Heng
Nanning Zheng
Venue: NeurIPS 2024 poster (2024)
Keywords: Molecule geometry modeling
Geometric GNNs
Long-range interactions
Primary area: machine_learning_for_other_sciences_and_fields
TL;DR: We introduce Neural P$^3$M, which enhances Geometric GNNs by integrating meshes with atoms and reimaging traditional mathematical operations in a trainable manner.

Abstract: Geometric graph neural networks (GNNs) have emerged as powerful tools for modeling molecular geometry. However, they encounter limitations in effectively capturing long-range interactions in large molecular systems. To address this challenge, we introduce **Neural P$^3$M**, a versatile enhancer of geometric GNNs to expand the scope of their capabilities by incorporating mesh points alongside atoms and reimaging traditional mathematical operations in a trainable manner. Neural P$^3$M exhibits flexibility across a wide range of molecular systems and demonstrates remarkable accuracy in predicting energies and forces, outperforming on benchmarks such as the MD22 dataset. 
It also achieves an average improvement of 22% on the OE62 dataset while integrating with various architectures. Codes are available at https://github.com/OnlyLoveKFC/Neural_P3M.

Neural M A Long-Range Interaction Modeling Enhancer for Geometric GNNs
Yusong Wang Chaoran Cheng"; Shaoning Li Yuxuan Ren Bin Shao Ge Liu", Pheng-Ann Heng', Nanning Zheng-| National Key Laboratory Human-Machine Hybrid Augmented Intelligence, National Engineering Research Center for Visual Information and Applications, and Institute Artificial Intelligence and Robotics, Xi'an Jiaotong University
University of Illinois Urbana-Champaign
3 Department of Computer Science and Engineering, The Chinese University of Hong Kong University of Science and Technology of China
Microsoft Research A14Science
wangyusoog22000ssuu xjtu. {snli24, pheng)@cse. cuhk. edu hk, binshao@miccm wttiaaamicroooo..ccommmm unzheng0mail. wtp.mmmmmiiiii..........ee edu. cn
www..edu.uu.uuu.....ooooo edu. cn, {chaoran7, geliu)@illinot udu
Abstract
Geometric graph neural networks (GNNs) have emerged as powerful tools for modeling molecular geometry. However, they encounter limitations in effec- tively capturing long-range interactions in large molecular systems due to the localization assumption of GNN. To address this challenge, we introduce Neu- ral M a versatile enhancer of geometric GNNs to expand the scope of their capabilities by incorporating mesh points alongside atoms and reimaging tradi- tional mathematical operations in trainable manner. Neural PdM exhibits flex- ibility across a wide range of molecular systems and demonstrates remarkable accuracy in predicting energies and forces, outperforming benchmarks such as the MD22 dataset. also achieves an average improvement of 22% on the OE62 dataset while integrating with various architectures. Codes available at http:://ghhub..oomm/////wvo////
P3M.
Introduction
Prevailing geometric graph neural networks (GNNs) have demonstrated remarkable capabilities in capturing the geometrie information inherent within molecular graphs. Not only do they accelerate the computational efficiency compared traditional Density Functional Theory (DFT) methods for molecules, but also hold the promise of achieving high-level accuracy in predicting crucial molecular
properties such as energy and forces 2 18 23 Despite their success in modeling small molecules, limitations still persist in extending these methods to larger molecular structures and systems governed by periodic boundary conditions (PBC). Current methods [16] excel in approximating the short- range interactions, which encapsulate interactions among local atom groups within defined distance cutoff, characterized by a rapid decay in real space. The primary obstacle lies in effectively capturing rangerange interactions within these complex systems.
Several attempts have been undertaken to incorporate long-range physical interactions into geometric GNNs. Early studies [99 21 combined physical equations, such Coulomb's 's law, with models
"quual contribution. Cooresppondigg author.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

tailored for short-range interactions. Conversely, recent advancements are steering towards the development of sophisticated models capable of learning long-range interactions directly from data. One such strategy is the spatial-based method, exemplified by LSRM [33 It utilizes specific fragmentation algorithms like BRICS [5] fragment molecules into discrete groups real space. The long-range interactions are thereby captured in hierarchical manner by facilitating message passing between the fragments and atoms. Another strategy the spectral-based method [12] 24], which treats the long-range parts in the reciprocal space following the concepts of Ewald summation [4]. The long-range parts exhibit a rapid decay instead in the reciprocal space, which enables efficient evaluation with frequency cutoff.
Following traditional computational chemistry, an intuitive direction would to mesh up the Ewald summation, harnessing fast Fourier transformation (FFT) for acceleration. While this poses non- trivial problem, rich of established works represented by Particle-Particee Particle-Me provide a solid foundation for such undertakings. In this work, inspired by the underlying unified concepts behind these FFT-accelerated methods, we propose novel perspective by integrating atom and mesh into neural networks. To be concrete, we reimage the traditional mathematical operations mesh-based methods in trainable manner, laying the foundation our new framework, termed Neural pdM M (Fig. 1b. Neural PdM is designed to be versatile enhancer, compatible with wide range of existing models. In contrast to LSRMM, Neural PdM framework remains unconstrained to any fragmentation algorithm. and hence enhances its flexibility across diverse molecular systems. Different from Ewald MP, Neural P'M explicitly incorporates mesh representations, thereby offering discrete resolutions necessary for formulating long-range terms. Additionally, it incorporates the exchange of information between short-range and long-range terms at the atom and mesh scales. Moreover, proposed framework exhibits theoretical efficiency surpassing that of Ewald MP due to the reduced computational complexity afforded by FFT.
(P*M) [[]]
Short-range
(Atom-Atom)
Neural P3M
Charge Assign (Atom-Mesh)
Long-range (Mesh-Mesh)
Figure 1: Illustration of Particle Particle Particle-Mesh (P*M) and its relationship with Neural P'M framework. The Atom2Atom block corresponds to the short-range term. The Aesm2 and Mesh2Atom block are similar to the charge assignment and back-interpan buck-interpolation. The Mesh2Mesh block corresponds to the long-range term.
We evaluate our framework on several benchmarks by integrating a variety of geometric GNNs. Neural PdM achieves the state-of-the-ar performance on the MD22 dataset [ and Ag dataset [16] when combined with ViSNet 23 It consistently demonstrates improvements in energy mean absolute errors (MAFs), achieving an average reduction of 22% on the OE62 dataset 20 summary, our contributions can be summarized as follows:

---

Framework. We propose a novel framework Neural P to capture short-range and long-range interactions at both atom and mesh scale.
Enhancement and Versatility. Neural P'M exhibits compatibility and significant improve- ments with short-range-centne methods on the Ag, MD22 and OE62 benchmarks.
Flexibility. Neural P'M is well-suited for diverse molecular systems without any constraints.
2
Preliminary
Ewald Summation Ewald summation a widely used technique in calculations of long-range interactions in periodic systems (7]. Specifically, consider the pair-wise electrostatic potential as (T.y)) = r/y|rayll l222- The total electrostatic potential energy E can be evaluated as the infinite summation pairs under the periodic boundary condition (PBC)
N Σ 1=1
NV
E
--
(1))) + n-n clle)lllddddd =
/f
Σ 2
1
(i(r) (1)
2
where (r) is charge density, c is the cell vector, and N is the number atoms in a cell. The ' summation is introduced to exclude the term j if and only if O. O (r) represents the potential generated by all particles excluding the particle i. A continuous partition function that delays rapidly with respect the distance is used to separate the short-range and long-range terms. One standard approach is to partition the contributions based on the error function erf:
(r) er (rl12)
(F) erl (3((11)))))
||||
|||12
where 33 is a fixed constant. We assume the charge density is described by the delta function as point charges, i.e. (r(r) 9;ō(r rz). With the rapid delay of the partition function, it is safe to assume convergence by only considering the interaction pairs within a specific cutoff distance as
Ex
1 Σ 2
J
N (1( (r)d" (
Σ " (rig) (1,))E
2
where & is the set of atom pairs within the cutoff distance. By the Parseval's theorem, the correspond- ing long-range term can be expressed the summation in the Fourier domain
N Σ 2
Err
/ (
Σ (mm))(m m²0
2V
where is the volume of the unit cell and ğ(m) = 4m /
are the Fourier transformed Green
function of the Coulomb potential 1 //||T|1, and ((m) = exp( charge density ç(m) is defined as
m / 3/4,2). The Fourier-transformel
N
p(m) = f
imr c =
dee im-r
(5)
The frequency vector can be truncated as the long-range term quickly converges in the Fourier domain. As the long-rang term introduces the ielf-inntraction energy, a correction term is also applied to the final potential energy as
N Σ / (r)d 2 1=1
N
4
(6)
VT
Meshing up the Ewald Summation The traditional Ewald summation method has a computational complexity of O(N2), which becomes impractical for large-scale systems. A common approach to accelerate the process is to employ FFT. Currently, a variety of mesh-based implementations are available. While they differ in their implementations, they share similar conceptual foundation (6)

---

Initially, point charges (particles) with their continuous coordinates, must be scattered onto grid-based charge densities (meshes). The charge densities on meshes are interpolated using charge assignment function W to ensure finite support for summation:
Viria / W(1, r) p(r)drr
1
N
PM (rp)
Wrip r;)
V gräd
where Vinia is the volume of the discrete grid ensure that PM is density. Once have discrete grid-based charge densities, we need to modify Eq accommodate discrete mesh points. According
the proof in Appendix B
Eq caan be rewritten the convolution in the real space:
N
1
E
2:00 (n.)
(r.
E 01g+
G * (r()
(8)
where G referred to influence function following Hockney and Eastwood 11 and the convolu- tion operation. The discrete approximation for E. can be expressed in corresponding manner as follows:
E
Σ (T)) (x)) TYEV
where, V is the set of mesh points. By altering the standard influence function to accommodate different charge assignment functions, one develop distinct algorithms. Subsequently, FFT is employed to accelerate the convolution process. Following the calculation of the energy, forces on particles can be determined by differentiation, either in the real space Fourier space. Alternatively forces can also be derived by differentiating on meshes and then applying a back inck-iintepooation technique to assign forces to partieles.
The adaptation of FFT to the Ewald summation has been quite enlightening. We will delve into a detailed examination of the correlation between our Neural PdM and these mesh-based techniques in subsequent section.
Method
We are interested in learning the energies and forces 3D molecules, potentially under the assumption of the periodic boundary condition. Specifically, consider 3D molecule represented as point cloud G = =x Heuu with atom coordinates x" and atom types we want to learn the molecule-level energy É()) and atom-level forces F()). Different from previous work [12] which utilizes the vanilla Ewald summation in the Fourier domain, our framework is mesh-based which provides discrete structural information and allows for information flow between long-range and short-range representations. Our fundamental concept is akin to these mesh-based methods mentioned in Section [ We use short-range blocks on atoms to capture bonded terms and non- non-bonded short-range terms while applying long-range blocks on meshes to handle long-range terms. We enable the transfer of information between atoms and meshes via the representation assignment. A pseudocode for the Neural M block is provided in Appendix D.1 to enhance understanding. We further elaborate on the Neural MdM architecture as follows.
3.1 Mesh Construction
Firstly, we construct meshes on which long-range interactions can be captured. In periodic systems such crystals, the cell is naturally delineated. For non-periodic systems, we adopt the approach used by prevalent quantum chemistry software, which involves padding the bounding box with a specified margin define the cell. Detailed information about the construction of the cell can be found in Appendix C The coordinates of mesh points X."j, can be described as:
14 + 1/2 c,
nJ + 1/2 Ny Cy
11k + 1/2
X
(10)
N₂
N
where = [cz; Cy, Cz is the cell vector and N., Ny, N₂ is the number of discretizations along each dimension. For convenience, we can regard meshes as point cloud with a single subscript for the index as {x"" hev.

---

Decoder
Geometiic GNN
Neural PM Block
Atom2Atom
Atom Aggregation
(c) Short-range Block
Neural PMM Block
Atom2Mesh
MeshZAtom
Encoder
Neural PM Block
(d) Long-range Block
m
Mesh Aggregation
Mesh2Mesh
m
CFny
Embedding
MLP
171
(a) Model Architecture
(b) Neural PM Block
(e) Repr. Assignment
Figure Overall framework architecture and details of each block. Geometric GNN models short-range interactions, Fourier neural operator (FNO) captures global long-range interactions, and continuous filter convolution (CFConv) exchanges information between two parts.
3.2 Embedding Block
Once coordinates of mesh points are established, we can proceed to construct short-range atomic radius graph and bipartite radius graph between atoms and meshes as follows:
E Shhurt
{ei ||x
xfil2
j E uU).
(11)
gussign
{ei : ||x - xx" 1l2 S asiin Vi E uu,jEV,
(12)
where 4 is the atom and V is the mesh set. Specifically, for periodic systems, the edges also obtained by considering possible cross- cross-boundar connections. The atom representation h& is initialized as:
h Embed(2;)
(13)
The initial mesh representation, denoted as mQ is obtained by averaging the representations of all neighboring atoms on the atom-mesh bipartite graph:
0 m
Σ
(14)
M(1)
JEM(S)
where M(i) represents the neighboring nodes connected mesh node within gassign The edge features in both E Shorrt and gassign can be expanded via set of radial basis functions (RBF):
f
pshort
RBF RBB (||x?
cassign
RBF ([|x;" x9112)
(15)
3.3 Neural M Block
Short-range Block The short-range block (FigZc)) updates the atomic representations using a graph neural network that is either Squiiaaqui invariant. This process can be generally expressed as follows:
7 = GNN(HL Eshon,, fshort)
(16)
We noted that the usage of radius graphs inherits the localization assumptions in geometric GNNs and any node only able to aggregation information from its direct geometric neighbors one short-range block. Therefore, we naturally interpret it as capturing the short-range contribution to the energy and forces. As this part involves only atoms, we call such module Atom2Aton which corresponds to the particle-partic part (short-range term) in the P.M.

---

Long-range Block The long-range block dd) updates mesh representations globally. Re- calling E the key aspect is devise the influence function and utilize FFT along with the convolution theorem for efficient computation of the convolution. Within our framework, we pa- rameterize C' directly in the Fourier domain, and the updated mesh representations can be described as:
n c o (Whonem (G))(.)))(())) (m'))
(17)
where F,F-1 are the Fourier transform and inverse Fourier transform on the discretized mesh, respectively. a is the activation function. Wilong and G are the learnable weights that parameterize the operator in the real space and Fourier space. If we consider our formulation coincides with the Fourier neural operators (FNOs) on the discretized continuous function. Similarly, as the long-range block only involves interactions within meshes, we call it Mesh2Mesh.
continuous function v(m),
Representation Assignment The representation assignment block (Fig 2 e)) allows for information flow between atom representations and mesh representations, effectively mixing short-range and long- range terms to obtain more comprehensive descriptor of the molecule. By parameterizing the charge assignment function W in Eq and substituting the charge density with the atom representation hj, we can derive the continuous filter convolution (CFconv) proposed in SchNet [17]. To elaborate further, get additional mesh representations as:
(m + a){ MLP
Σ JEM((
Ri.
W
cassigg
(18)
maaa
This Atom2Mesh module be regarded as the information flow from the short-range part the long-range part. Similarly, the Mesh2Atom module takes the same input and geometric graph but outputs additional atom representations (a <- m)', which could be viewed as the back-interpolation operation. It allows for the information flow in the inverse direction, from the long-range part to the short-range part. The long-range Mesh2Mesh module together with the Atoh2Mesh and Mesh2Atom modules corresponds to the particle-ment mesh part (long-range term) in the P'M.
Ultimately, shown in Fig merge the information updated by each part itself with the normalized information received from the other part, and we also incorporate a residual connection to obtain the final output
h1+1= h' + +h' +LN((a m)')
(19)
ml = m + m + LN((m + a)))
(20)
3.4 Decoder Block
As we are interested in the prediction of molecule-level energies and atom-level forces, an additional decoder is applied to the final atom representations 1/4 and mesh representations m² to get the atom-wise energies out and mesh-wise energies out We follow previous work to assume additive property of energy to sum all atom-wise energies as the short part of the molecule energy Eshort
Esshor
Σ
Σ MLP((N(()4))
(21)
JEU
jEU
We also sum all mesh-wise energies as the long part of the molecule energy flong
blong
Σ
m
Σ MLP(LN(())))
(22)
JEV
JEV
The final potential energy is calculated as: E = Blong + Eshort, Furthermore, although direct prediction of forces is possible, we instead use the negative gradient of the energy as the prediction of
forces: F = V E. The final training objective is a weighted loss between energy and force:
AF
= >
E
Σ
R++
(23)
3N

---

Experiment
4.1
Experimental Setup
In this section, we conduct comprehensive validations of our Neural P'M framework using diverse datasets and configurations. First, we intuitively demonstrate the necessity of incorporating long- range interactions through toy dataset Ag used in Allegro [16]. Subsequently, we integrate various geometric GNNs [17
23 with our Neural P framework on two prevalent datasets 3 to demonstrate versatility and effectiveness. All results are evaluated using mcan absolute error (MAE) test sets, and the bascline results sourced directly from the corresponding papers. Unless stated otherwise, almost all hyperparameters align with the baseline GNNs. For more comprehensive overview of hyperparameter settings and implementation details, please refer to the Appendix D and
9 18
g
OE62 [20] and MD22
4.2
Toy Dataset: Ag
The Ag dataset comprises 1,159 struc- tures sampled from 1, 11K AIMD simulation [16]. These structures were generated from bulk face- centered-cubic lattice with vacancy, encompassing 71 atoms subject to pe- riodic boundary conditions. For con- sistency with Allegro, we randomly split them into 950 structures for train-
Energy MIEE mmeeaatom)))
0.750
0.574
0,328
ing, 50 structures for validation and the remaining structures for testing. As shown in Fig. 3 compared to the strictly local Allegro model, ViS- Net, which has only one layer, offers slightly improved force prediction, yet the energy prediction significantly de- teriorates. This may be caused by the fact that the model can only perform message passing once, with a lack of long-range interactions. Long-range interactions can be complemented in theory by raising the cutoff from 4.0 to 12.0 Å, but this does not work in practice, because it could potentially lead to information over squashing problems, as mentioned in LSRM 13 When ViSNet with single layer is integrated into our framework, long-range interactions can be effectively captured, significantly improving the accuracy of energy and force predictions compared to the vanilla ViSNet and Allegro. This toy experiment intuitively demonstrates the critical need incorporate long-range interactions and emphasizes the significance of well-crafted methodology in incorporating them.
Allegro
NPM (Cutoff=4.0
IO 4.0
(Cutom-
Force NAE (mew()
19.1
25
Allgro
NPMM
[Cutolf=4,0 Lm11
(Cutoff 12.0 L-11
(CuRoft
Figure 3: Mean absolute errors (MAEs) for energy and force predictions on Ag dataset are compared among Allegro, ViSNet, and our proposed framework.
4.3
MD22
The MD22 dataset consists of MD trajectory datasets, which present challenges due to their larger system sizes, ranging from 42 to 370 atoms. The number of structures in each molecule dataset ranges from 5,032 to 85,109. We calculate the diameter of each molecule, defined as the average of the maximum distance between any two atoms within a molecule. The smallest diameter observed
approximately 10.75 Å, while the largest molecule measures about 32.39 Å. We train a separate model for each molecule and randomly split the dataset according sGDML B.
Table I demonstrates the results of the ViSNet model incorporating with our Neural PdM framework (ViSNet-NP'M for short) on MD22. ViSNet-NP'M achieves the state-of-the-ard (SoTA) performance on both energy and force predictions across the four largest molecules and also achieves the lowest mean absolute error (MAE) for energy or force predictions in the remaining three smaller molecules.
7

---

Table 1: Mean absolute errors (MAE) of energy (kcal/mol) and forces (kcal/mol/A) for seven large molecules on MD22 compared with state-of-the-ar models. The best one in each category is highlighted in bold.
Molecale
Diamever (A)
SGDML
SO3KRATES
Allegro
Equiffommr
MACE
ViSNet
Baseline
Ewald
LSRM
Neural P'M
NH-e
10.75
energy
0.3902
0.337
0U0620
04775
0.034
0.7968
0.1068
0.0804
0J0876
0.0922
0.0814
0.0902
DIAA
14.58
eneryy
lorces
0.7474
1.3117
0.379
0.242
0,1155
0.0732
0.1788
0.0506
0.1317
0.0646
0.1326
0.0668
0.0933
0.0664
0.8833
0.0998
0,0712
0.0019
Stachyose
13.87
energy
4.0497
0.442
0.2485
1404
0.1244
0.1283
0.1089
0.1055
0,8856
forces
0.6744
0.435
0.0971
0.0635
0.0876
0.0869
0.0976
0.0767
0.0940
AT-AT
17.03
0,7235
06611
0.178
0,1428
0.1309
0.1083
0. 0.1688
0.1487
0.0885
0.0772
0.0714
AT-AE-CG-CG
energy
forces
0.7028
0.332
0.1280
0.1252
0.3153
0.1563
0.1571
0.1115
0.1063
Buckyball coscher
15.89
energy
foroes
1.1962
0.6820
0.381
0.237
0,5258
0.0887
0.3978
0.1114
0.4812
0.0853
0.4121
0.1335
0.3575
0,0989
0.4220
0.1026
0.3543
0.08466
Doubll-walled nanotube
32,59
energy
forues
4.0122
0.5231
0.993
0.727
0,3428
2.2097
0.2747
1.1945
0.2767
16655
0.3959
1.0899
0.7909
0,2875
0.3391
1,820
0.2561
0.7751
When compared to the vanilla ViSNet, ViSNet-NP*M showed an average improvement of 34.6% and 21.2% in energy and force prediction, respectively. Notably, our framework exhibits more substantial improvement when compared ViSNet-LSRM and ViSNet-Ewald, both of which utilize ViSNet as the short-range model. As shown in Appendix Table another statc-of-the-art model, Equiformer, when integrated with our Neural M'M framework, also demonstrates significant enhancements the short-range model itself. These impressive results highlight our framework's 's ability to effectively improve the learning of potential long-range interactions large molecules.
It's worth noting that for the two supramolecules that cannot be fragmented by LSRM, our Neural P'M achieves a significant performance improvement in energy prediction, with 57.48% increase for the double-walled nanotube and a 16.07% increase for the buckyball catcher. This suggests that our Neural P'M is general solution for various molecules, which is not limited by traditional fragmentation methods like BRICS.
4.4
OE62
We further take our analysis by incorporating four prevailing geometric GNNs including SchNet [17), PaiNN [18], DimeNet++ [9]. and GemNet-T [ on the OE62 dataset [20] to confirm the framework's versatility. The OE62 dataset consists of about 62,000 large organic molecules, each with the energy calculated by Density Functional Theory (DFT). The structures within the OF62 dataset are non- periodic yet can span large spatial dimensions, exceeding 20 Å. The dataset is strictly split into train, validation, and test according to Ewald MP [12.. The same dataset preprocessing process as Ewald MP is also applied.
The numerical results presented in Table 2 and Appendix Table 6 indicates that the Neural P'M framework, which combines four models, delivers more performance gains than Ewald MP and LSRM when using the same hyperparameters. Additionally, our framework exhibits faster computation time than Ewald MP, likely due to the efficiency of FFT implementation by Pytorch. An unexpected observation is the speed performance of DimeNet++. Given that DimeNet++ does not inherently facilitate message passing between atom embeddings, Ewald MP compensates by integrating long- range interactions in each output block. In contrast, approach exchanges short-range and long- range representations each layer, which might account for our marginally slower speeds compared to Ewald MP. We also provide detailed profiling results for the number of model parameters, GPU memory usage, and other relevant metrics in Appendix G For more details on the implementation on the four models, please refer to the Appendix
4.5 Ablation Study
4.5.1 Architecture
We first investigate the impact of the Atom2Mesh and Mesh2Atom tom modules. We remove the Atom2Mesh module from the original model to avoid the information flow from short-range blocks

---

Table 2: Energy MAEs and computation times per input structure for the OE62 dataset compared with Ewald MP and other baseline methods. The data was sourced directly from [2].
OE62-val
OE62-test
Forward Pass
Forward & Backward Pass
Model
Variant
Rel.
MAE meV
Rel.
Rumtime ms/struct.
Rel.
Runtime ms/struct.
Rel.
meV
%
% f
%
%
SchNet
Baseline
133.5
131.3
0.13
0.28
Emheddings
144.7
136.7
-4.1
0.14
15.2
0.33
17.8
Cutoff
257.4
-92.8
254.8
-94.1
0.14
13.6
0.31
11.6
SchNett-L
86.6
35.1
89.2
32.1
0.32
156.0
0.75
171.7
Fwald
79.2
40.7
81.1
38,2
0.70
461.6
1.03
271.4
Neural P'M
70.2
47.4
69.1
47.4
0.37
184.6
0.57
103.6
PaiNN
Baseline
63.3
1.52
3.16
Embeddings
63.5
63.1
1.4
Cutoff
65.1
-6.0
64.4
-2.2
1.84
20.9
3.91
23.6
SchNct-LR
58.3
5.1
58.2
7.7
1.84
20.7
4.21
33.1
Ewald
57.9
5.7
59.7
5.7
2.29
50.5
4.57
44,4
Neural P'M
54.1
11.9
52.9
16.4
2.17
42.8
4.19
32.6
DimeNet++
Baseline
51.2
53.8
Embeddings
50.4
53.4
Cutoff
48.3
5.7
48.1
10.6
2.68
34.7
6.10
43.4
SchNet-LR
51.4
-0.5
$4.4
-1.1
2.37
19.0
4.73
11.2
Fwald
46.5
9.2
48.1
10.6
2.70
35.5
5.93
39.5
Neural PDM
40.9
20.1
41.5
22.9
3.11
56.3
5.62
31.9
GemNet-T
Baseline
51.5
53.1
Fmheddings
52.7
-2.3
53.9
1.5
0.4
Cutoff
47.8
7.2
47.7
10.2
4.02
31.2
8.88
27.7
SchNet-LR
51.2
0.6
52.8
0.5
3.32
8.3
7.73
11.1
Ewald
47.4
8.0
47.5
10.5
4.05
32.0
8.86
27.4
Neural P'M
47.2
8.3
47.4
10.7
3.93
28.0
7.71
10.8
to long-range blocks and vice versa. Table[3]emoonstrattt that both modules contribute synergistically to the model's overall performance. The results illustrate the necessity of enabling information exchange between the long-range and short-range blocks.
4.5.2 Hyperparameters
Compared to the vanilla model, our framework introduces only two new hyperparameters: the assignment cutoff distance between mesh points and atoms, denoted as passign and the number of mesh points in each dimension, represented as Na, Ny. N2.
We find that the selection of the number of mesh points is crucial for the model's final performance. As illustrated in Appendix Fig. b the mean absolute error (MAE) in energy increases with the number of mesh points, while the forward computation time also extends. This decline in performance may attributed to instances where each atom is assigned to multiple mesh points simultaneously. As such occurrences become more frequent, the model may struggle effectively learn the appropriate assign- ment rules. In practice, we typically set the cutoff' distance to either 4.0 or 5.0 Å, ensuring that the product of the num- ber of mesh points and the cutoff distance is approximately equivalent to the cell size in each dimension.
Table 3: Energy MAE of SchNet-NP*M variants on the OE62 dataset. The best one is highlighted in bold.
Architecture Variants
Energy MAE
Original Without Mesh2Atom
69.10
76.14
Without Atom2Mesh
74.48
Without Both
72.07
Additionally, provide further ablation studies on the impact of the assignment cutoff distance (without the k-NN graph) to examine the effects of multiple assignments. As shown in Appendix Fig. c). all experiments exhibit slight decrease in performance due to multiple assignments. However, appropriately chosen cutoff (4 or 5 À) still yields relatively optimal results. Notably, the results do not worsen further as the assignment cutoff increases. We hypothesize that this may be because larger assignment cutoff creates broader neighborhood environment, facilitating the

---

learning of the assignment function with a fixed number of meshes, thereby alleviating the challenges associated with multiple assignments.
Related Work
Geometric Graph Neural Networks Geometric graph neural networks preserve equivariance toward the rigid transformation in space, which can be categorized according to their emphasis on specific types of structural features and their respective methods of integration. SchNet [77 stands out as the pioneering approach to applying continuous filter convolution on molecular distances. Subsequently, DimeNet++ [9] and GemNet [8] explicitly incorporate angles and dihedrals using Fourier-Besse functions. To address the computational complexity associated with angles extractions, PaiNN [8] and ViSNet (33 adopt the density trick and reduce the complexity to linear time. Additionally, many works are based on high-order geometric tensors [2 16 (2), which ensure rigorous theoretical guarantees of equivariance through the use of Clebsch- Gordan product. Despite these advancements, all these existing methods are constrained to the local atomic environment, and are unable approximate the long-range interactions. Hence, there is an urgent need for comprchensive framework to address this challenge.
Long-range Interaction Modeling Incorporating long-range interactions into a short-range model is challenging. Early studies attempted to compensate these long-range effects by integrating physical equations with either hand-crafted terms 19 or predicted charges (2I). While, recent works have shifted towards creating carefully designed models that can directly learn long-range interactions from data. The LSRM framework [13] for instance, captures long-range interactions in real space by using specific algorithms to fragment molecules into discrete groups and models their interactions hierarchically. Other methods 12
24 15 handle long-range components in reciprocal space, employing concepts like Ewald summation 4 Our approach differs from these works by introducing the discretized meshes and facilitating the exchange of information between long-range and short-range components.
6 Conclusion
In this paper, we introduce a novel framework, termed Neural P3M, designed to enhance the long- range interaction modeling for various geometric GNNs. In addition, Neural PaM stands by not being confined to any specific fragmentation approach, making it adaptable to various molecular systems. Neural PdM achieves significant performance improvement on prevalent benchmarks by capturing short-range and long-range interactions at both atom and mesh scales, and enabling the exchange of information between them.
The limitation of our study is that it does not thoroughly investigate the impact of the number of meshes, nor does it explore potentially more effective methods for modeling long-range interactions beyond FFT. Nonetheless, our paper offers the community fresh perspective molecular geometry modeling. Our proposed Neural P framework is extensive of existing geometric GNNs for energy and force prediction of molecules. The prediction molecular energies and forces has diverse applications in downstream tasks including molecular dynamics simulation and molecular property prediction. As our framework better captures the long-range interaction within the molecule, it can potentially accelerate the pharmaceutical discovery and understanding of diverse molecules that have positive impacts on treating diseases. On the other hand, we are also aware of the potential negative impact if the model misused, as our understanding of different molecules is still very limited. We will work closely with both the machine learning and the science community ensure the proper usage of our model for the good of society.
Limitation and Societal Impacts:
7 Acknowledgments and Disclosure of Funding
We thank the reviewers for their valuable comments. This work was supported by NSFC under grant No. 62088101 102.
10

---

References
[1] I. Batatia, D. P. Kovacs, G. Simm, C. Ortner, and G. Csányi Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35: 1423-1122222 1436, 2022.
[2] S. Batzner, A. Musaclian, L. Sun, M. Geiger, J. P. Mailoa, M. Korblluth, N. Molinari, T. E. Smidt, and B. Kozinsky. E (3)-cqivarriant graph neural networks for data-cfficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
[3] S. Chmiela, V. Possilev-Galiinee O. T. Unke, A. Kabylda, H. E. Sauceda, A. Tkatchenko, and K.-R. Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 92))eadd0877, 2023.
[4] S. W. de Leeuw, J. W. Perram, and E. R. Smith. Simulation electrostatic systems in periodic boundary conditions. i. lattice sums and dielectric constants. Proceedings of the Royal Society of London. A. Matheaatcall and Physical Sciences, 373(1752) 1980.
[5] J. Degen, C. Beegshhid-GGerlade A. Zaliani, and M. Rarey. On the art of compiling and using drug-like' chemical fragment spaces. ChemMedChem, 3(10):1503, 2008.
[6] M. Deserno and C. Holm. How mesh up ewald sums. i. theoretical and numerical comparison of various particle mesh routines. The Journal of chemical physics, 109(18);7678-769,, 1998.
[7] P. P. Ewald. Die berechnung optischer und elektrostatischer gitterpotentiale. Annalen der physik, 369(3):253-277, 1921.
[8] J. Gasteiger, F. Becker, and S. Günnemann. Gemnet: Universal directional graph neural networks for molecules. Advances in Neural Information Processing Systems, 34:6790-6802, 2021.
[9] Gasteiger, S. Giri, J. T. Margraf, and S. Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrium molecules. arXiv preprint arXiv:2011.14113 2020.
[10] J. Gasteiger, J. GroB, and S. Günnnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.00 033122,, 2020.
[11] R. W. Hockney and J. W. Eastwood. Computer simulation using particles. ere Press, 2021.
[12] A. Kosmala, J. Gasteiger, N. Gao, and S. Günnemann. Ewald-based long-range message passing for molecular graphs. In International Conference on Machine Learning, pages 17544-17563, PMLR, 2023.
[13] Y. Li. Y. Wang, L. Huang. H. Yang. X. Wei, J. Zhang. Wang, Z. Wang, B. Shao, and T.-Y. Liu. Long-short-range message- message-passing: A physics-informec framework to capture non-local interaction for scalable molecular dynamics simulation. arXiv preprint arXiv:2304.1 13542....2, 2023.
[14] Y.-L. Liao and T. Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint 2tti::22006119990 2022.
[15] Y. Lin, K. Yan, Y. Luo, Y. Liu, X. Qian, and S. Ji. Efficient approximations of complete interatomic potentials for crystal property prediction. In International Conference on Machine Learning, pages 21260-21287. PMLR, 2023.
[16] A. Musaelian, S. Batzner, A. Johansson, L. Sun, C. J. Owen, M. Kornbluth, and B. Kozin- sky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023.
[17) K. Schütt, P.-J. Kindermans, H. E. Sauceda Felix, S. Chmiela, A. Tkatchenko, and K.-R. Müller. Schnet: continuous-filet convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
[18] K. Schütr, O. Unke, and M. Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. International Conference on Machine Learning, pages 9377-9388. PMLR, 2021.

---

[19] C. G. Staacke, H. H. Heenen, C. Scheurer, G. Csányi, K. Reuter, and J. T. Margraf. On the role of long-range electrostatics in machine-learred interatomic potentials for complex battery materials. ACS Applied Energy Materials, 4(1 1:12562-12569) 2021.
[20] A. Stuke, C. Kunkel, D. Golze, M. Todorović, J. T. Margraf, K. Reuter, P. Rinke, and H. Ober- hofer. Atomic structures and orbital energies of 61,489 crystal-forming organic molecules. Scientific data, 7(1):58, 2020.
[21] O. Unke, S. Chmiela, M. Gastegger, K. T. Schütt, H. E. Sauceda, and K.-R. Müller. Spookynet: Learning force fields with electronic degrees of freedom and nonlocal effects. Nature communi- cations, 12(1):7273, 2021.
[22] Y. Wang, S. Li, т. Wang, B. Shao, N. Zheng, and T.-Y, Liu. Geometric transformer with interatomic positional encoding. Advances in Neural Information Processing Systems, 36, 2024.
[23] Y. Wang. T. Wang, S. Li, X. He, M. Li, Z. Wang. N. Zheng, B. Shao, and T.-Y. Liu. Enhancing geometric representations for molecules with equivariant vector-scalar interactive message passing. Nature Communications, 15(1):313, 2024.
[24] H. Yu, L. Hong, S. Chen, X. Gong, and H. Xiang. Capturing long- range interaction with reciprocal space neural network. arXiv preprint arXiv:2211.16684, 2022.
12

---

Supplemental Material
A
Notations
Table 4: Glossary of notations
Notation
Description
The index of atoms meshes
7
The index blocks
x,
The coordinates of particles (atoms)
Tp, x
The coordinates of the meshes
ri
The displacement vector between the particle and
c m
The cell vectors
The frequency vectors
erf
The function
p())
The charge density of the particle
5(-)
The delta function
The point charges
pv())
The charge density of the mesh point
+0)
The pair-wise electrostatic potential
The potential generated by all particles
(1))
The potential generated by all particles excluding the particle The potential generated by the particle i. The Fourier transformed function g, P
(.(-)
9. 7:7%
The convolution operation
N
The number of particles in unit cell The charge assignment function influence function
W
Vgnia V
The volume of the discrete grid and cell.
E, E
The ground truth and prediction of potential energy
F
The ground truth and prediction of atomic forces The atom types
z
U,
The of atoms and meshes
No, Ng. N:
The number of discretizations along cach dimension x, y/
Shbrtt passinn
The cutoff distance of radius graphs
Eshomt Eussign
The edge set radius graphs
M().(())
The neighboring nodes of the target atom (mesh) node. The atom representations
h
m
The mesh representations
f
The edge representations
F,F-!
The Fourier transformer and inverse Fourier transform on the discretized mesh (FFT and IFFT)
I|:1|2
The 2-norm of a vector
a(-)
The activation function (SiLU)
GNN(.)
The short-range graph neural network (learnable)
Wlong
The weights in long-range block in real space (learnable)
G
The weights in long-range block in Fourier space (learnahle)
W
Wac
The weights representation assignment (learnable)
MLP(.)
The multi layer perception (learnable)
LN(-)
The layer normalization
λ AF
The weights in the loss between energy and forces
13

---

B Detailed Derivation of Eq
Let's start our derivation by replacing the the square of the p(m)'s modulus the product of itself with its conjugate:
1
E
(24)
2V
m=30
m) (m)
(25)
2V
m20
(m)
q e
(26)
2V
m=0
We can then confidently interchange the summation symbols and put the normalization factor 4 within the summations as:
Es
9
Σ
(27)
j=1
m3
N
Σ
9
1
Σ m
((m))
(28)
2
V
Using convolution theory, which states that the convolution of two functions is the pointwise product of their Fourier transforms, it becomes clear that the expression in parentheses represents the inverse Fourier transform. Consequently, we can rewrite the expression as follows:
1 _ p 2
E E
1gg p(r,)
(29)
We refer 7 as the smeared Coulomb Green function (influence function), and altering it when assigning charges with different charge assignment function W.
14

---

C Detailed Implementation of Cell Construction
Cell construction trivial for periodic systems like erystals, as a canonical cell can always be assigned. We now describe the cell construction for non-periodic systems. Given a set of atom coordinates {x: first derive a canonical coordinate frame U the eigenvectors of the covariance matrix:
UAU = (X - 1 -
(30)
where is the diagonal matrix of the eigenvalues of the covariance matrix, X E Rax3 is coordinate matrix, and = Ex=1 x//n. For any rotation matrix R and X' = XR, it is easy to see that U RU new eigenvector matrix for the new covariance matrix. Therefore, we use canonical coordinates as = (X DUUT which is invariant under global translation and rotation. After the transformation, the principle components of the coordinates now align with the coordinate frame. We can define the cell vectors to follow the directions of the coordinate with the cell length defined by the maximum coordinate span with additional padding d on both sides:
Cy
T;
min I; + 2d 15i5"
18i5
Cy
max Y1 1SiSn
min D. + 2d ey
(31)
15
C,
min
2, + 2d
e
where ã, J. 2 are coordinate components of the transformed molecules. In practice, we used d = 0.5A. The mesh coordinates are obtained via Eq and the final atom coordinates are obtained by moving the molecule inside the cell as:
min 32 d. min D: d min 2n =d)a U Isisn
Y= x
(32)
15t5n
15$5
There are rare cases when the molecule exhibits high symmetry. However, as only consider different atom types and treat the same type atoms as indistigguishable the final molecule and mesh are also indistinguishable and unique in this sense.
15

---

D Detailed Implementation for Integrating Various GNNs into Neural P'M
In this section, we first provide the pseudocode for the Neural P'M block facilitate understanding of our framework, followed by a detailed explanation of the implementation. We emphasize the distinct integration strategies required the varying inputs and outputs of short-range geometric GNNs. For detailed insights into the specific implementations within geometric GNNs, we direct readers to the original paper.
D.1 Pseudocode for Neural P'M Block
The pscudocode for the Neural P'M block is presented in Algorithm I as general framework for iteratively and interdependently updating the atom features h and mesh features m. The GNN in the algorithm can incorporate most geometric GNN frameworks that use node features, the atom graph Eshor, and edge features fshort input to update the node features. The FNO serves the long-range block, updating mesh features according to Eq 17 The representation assignment then calculates the relevant features based the assignment graph assign and edge features fasign, Finally, the overall representation is updated using information from both the short-range and long-range blocks.
Algorithm Neural P block
1: Input: Atom feature h', mesh feature m', atom graph Eshort, assignment graph gussign and edge features fsheet fassign
2: h GNN(W) Eshort fshort)
Atom2Atom (Short-range)
3: m' + FNO(m²)
>> Mesh2Mesh (Long-range)
4: (a + (mm)) + MLP EjEA(h m. Wi passign at-mJij 11 passign
Mesh2Atom (Repr. Assignment)
5: (m a)}: HMPP
MLP (EEM(C) Ri. W 11
Atom2Mesh (Repr. Assignment)
h/+1 11' hh' LN((( - m))
Mesh2Atom (Update)
7: ml41 m² + m LN((In + a))
Atom2Mesh (Update)
8: return h1*1, m
D.2 SchNet
SchNet [17] utilized continuous graph convolutional kernels generated from edge features of radial basis functions (RBFs) capture the geometric information of interatomic distances. In each Neural P'M Block, the atom representations hí and mesh representations m' are initially subjected to layer normalization before being processed by SchNet Block and FNO Block, respectively.
SchNet (lN()(
(33)
m FNO(LN(mm))
(34)
Following this, the representation assignment block updates these representations separately.
(m + a) Atom2Mesh(h"....)
(35)
(a + m) = Mesh2Atom (m'
(36)
The exchanged representations then normalized and combined with their corresponding updated representations via an addition operation. Finally, we employ residual concatenation to obtain the final representation:
h1+1 'h' LN((a m)') m m + m + LN((m + a)a)
(37)
(38)
D.3 PaiNN
PaiNN [8] is an equivariant graph neural network based on scalar-vector interactions. Each hidden state is described by tuple of scalar representations /e' and vector representations vee! and updated
as follows:
h/ vēe' PaiNN Block(LN((h), vec'...)
(39)
16

---

We only use scalar representations to exchange information with mesh representations, vector repre- sentations can also get long range information when interacting with scalars. The implementations of other parts are consistent with SchNet.
D.4 DimeNet++
DimeNet++ ( is improved version of the original DimeNet [10] architecture. In addition to distance, it further leverages the geometric information of any angles formed by three nodes and applies 2D spherical Bessel functions to embed the angles. Thus, the hidden state f' of DimeNet++ is at the edge level. To exchange information between atoms and meshes, we need to aggregate the
edge-level representations to the node-level representations as follows:
DimeNet Block(LN(f')....
(40)
14 = Σ F/j W
RBF ([|x) x)112)
(41)
JEN(()
The subsequent implementations consistent with SchNet, while in order obtain the final edge- level representations, we combine the atom representations on both sides of the edge, and finally update as follows:
(Aedge f m)(5 = 55 (W Concar Concati(a m)/ (a t m)(l)
(42)
f(+1 =f + f + LN(((edge + m)))
(43)
m m' + m + LN((M + a)²)
(44)
D.5 GemNet
GemNet g further extends DimeNet to incorporate geometric information of dihedral angles formed by four atoms and applies high-order Bessel functions to embed the dihedral angles. However, since the computational complexity of quadruplets is too high, GemNet-T used this paper still uses triplets, which can be viewed as more complex DimeNet. GemNet updates both atom-level and edge-level representations as follows:
i = GemNet (loc((),( L,(()))
(45)
We use node-level representations to exchange information with mesh representations, and subsequent implementations are consistent with SchNet. The final representations are updated as follows:
1
(46)
g++ f' LN((Gedge +- m))) m 11 = m' + i + LN((N+ +a)a)
(47)
(48)
should be noted that updates are made solely the edge-level representations to prevent information redundancy. Our observations indicate that edge-level representations are predominantly parts of GemNet, hence, we focused our updates there. Additionally, we remove the scaling factor from our implementation.
D.6 ViSNet
ViSNet [23] is upgraded version of PaiNN, also utilizing scalar-vector interactions that can describe angles, dihedral angles, and improper angles in linear time complexity. When training ViSNet on the MD22 dataset, we find that ViSNet suffers from unstable training when learning rate is relatively large, so we slightly modified the implementation. Unlike the first 4 models, instead of exchanging information using the representations after updating, we use the input representations directly after layer normalization:
R vēe', f ViSNet (LN(h,, vec' LN(Ff))
(49)
m' = FNO(LN(m²))
(50)
(m - a) Atom2Mes
(51)
17

---

(a t m)' = Meshh.AtomLLN(mm..
(52)
The final representations are modified as follows:
h 1++1 = h² + it' + (a t m)
(53)
m +1 = mm' m (m +- a)
(54)
ا = f + fr
(55)
This modification is similar to altering from post-normalization to pre-normalization in the standard Transformer.
18

---

E
Additional Results
E.1 Performance of Integrating Equiformer into Neural PDM
We further evaluate the performance of the Equiformer model integrated with our Neural P'M on the two largest molecules MD22, as shown in Table The results demonstrate our framework's 's con- sistent improvement over mainstream state-of-the-ar models, further highlighting the compatibility of our approach.
Table 5: Mean absolute errors (MAE) of energy (kcal/mol) and forces (kcal/mol/A) for the two largest molecules on MD22 compared with Equiformer baseline 14]. The best one in each category highlighted in bold.
Molecule
Diameter (A)
Equiformer
Baseline
Neural PDM
Buckyball eatcher
15.89
energy forces
0.3978
0.3038
0.11
0.1018
Double walled nanotube
32.39
energy forces
1.1945
0.6208
0.2747
0.2399
E.2
Comparison of Performance OE62 between LSRM and Neural PåM
Due to the diversity of OE62, the fragmentation algorithm used by LSRM is not suitable for all molecules in this dataset. Nevertheless, for the sake of thorough comparison, applied filtering to select a subset of molecules and used this slightly modified dataset to evaluate LSRM's performance OE62. The results in T indicate that while LSRM outperforms the baseline, its performance remains below that of our Neural M
Table 6: Energy MAEs the OE62 dataset compared against the LSRM and baseline models. The
best one in each category is highlighted bold.
OE62-val val
OE62-tee test
Model
Variant
MAE meV
Rel.
MAE meV
Rel.
% 1
% f
SchNet
Baseline
133.5
131.3
LSRM
72.9
45.4
72.6
44.7
Neural P
70.2
47.4
69.1
47.4
PaiNN
Baseline
61.4
63.3
LSRM
56.6
7.8
56.4
10.9
Neural P
54.1
11.9
52.9
16.4
DimeNet++
Baseline
51.2
53.8
LSRM
47.9
6.4
50.4
6.3
Neural PdM
40.9
20.1
41.5
22.9
GemNet-T
Baseline
51.5
53.1
L.RM
50.8
1.4
51.5
3.0
Neural P
47.2
8.3
47.4
10.7
19

---

4.0 Á Hyperparameters of Neural P
E1 Common Hyperparameters
Ag Dataset We use a compact ViSNet which has only single layer with 128 hidden dimensions and maximum spherical harmonic order mmxx 1. For training, we employ the AdamW optimizer with a batch size of 4. The learning rate is dynamically adjusted using the ReduceLROnPlateen scheduler with decay factor of 0.8, triggered after a patience interval of 30 epochs without improve- ment. The initial learning rate is set to 0.0018, 18, is preceded by warm-up phase of 1000 steps. In our loss function, energy and force are weighted at ratio of 0.1 0.9, respectively, balance their importance during the training process. We employ an early stopping mechanism that terminates training if the validation metric does not improve after 600 epochs. Experiments are conducted on NVIDIA 16G V100 GPU.
MD22 Dataset We employ a ViSNet that consists of 6 layers with 128 hidden dimensions, and maximum spherical harmonic order of mmxx = 1 to enable fair comparison with LSRM model. We adjust the batch size for each molecule to achieve approximately 1000 steps per epoch (a batch size of 6 for Ac-Ala3-NHMa
8 for DHA and so on). The initial learning rate is carefully tuned within the range of 0.001 to 0.0018 to optimize performance. Additionally, the weights of energy and force in the loss function is customized for different molecules, with supramolecules using weight of 0.005 for energy and 0.995 for force, while other molecules using a ratio of 0.05 0.95. Other settings remain the same the Ag dataset. Experiments conducted NVIDIA 16G V100 GPU.
Regarding the four models trained on the OE62 dataset, providing a detailed hyperparameters each challenging due to their uniqueness. However, ensure fair comparison, we set the hyperparameters in line with Ewald MP exactly. The only difference that after eliminating the scaling factor from the GemNet implementation, range of 0.0001 to 0.0005. Experiments are conducted on a NVIDIA 80G A A100 GPU.
OE62 Dataset
tuned the initial learning rate within the
F22 Hyperparameters about Mesh Construction
In this subsection, we detail the hyperparameters employed during the mesh construction process. The empirical principles guiding their selection are discussed in Section 4.5 here we focus on the specific hyperparameters in practice.
Table 7: Hyperparameters employed during the mesh construction process different molecules
Dataset
Molecule
Expand
Short-range cutoff (pthmn)
Assignment cutoff (pussiz)
N₂
Ag
4.0
Ac-M
Å
Á
Á
Stachyose
1.0 Á
4.0 Á
Á
MD22
AT-AT
Á
5.0
5.0 Á
AT-AT-CG-CG
1.0 Á
5.0 A
5,0 Á
Buckyball catcher
1.0 Á
4,0 A
5.0 Á
4
Double-wald walled nanotube
1.0
4,0 Á
5.0 A
7
OE62
1.0
6.0 Á
A
20

---

G Profiling Results of Neural P'M on OE62
We present the number of parameters and memory usage (with standard settings and batch size of of the largest molecule in OE62) as well as the maximum batch size that can be accommodated on single A100 GPU the Table X The bulk of the memory usage is still attributed to the short-range modules-for for instance, 167 19 MB versus 19945 MB in GemNet. As anticipated, the integration of the mesh concept and additional modules means that Neural PSM has higher parameter count and slightly greater memory usage than Ewald MP. Nevertheless, this modest increase in resource demand offset by the significant performance improvements offered by Neural PoM, along with the computational acceleration brought by FFT.
Table 8: Profiling results the OE62 dataset compared with Ewald MP and other baseline methods. The best one in each category highlighted in bold.
Model
Variant
# of Parameters (M)
GPU Memory (MB)
Max. Batch Size
SchNet
Baseline
2.8
400
Embeddeds
14.4
344
Cutoff
2.8
1671
392
SchNet-I.R
5.3
4835
128
Ewald
12.2
2675
240
Neural P'M
19.1
2283
280
PaiNN
Baseline
12.5
8135
80
Embeddeds
15.7
9073
72
Cutoff
12.5
20480
32
SchNet-LL
15.1
11289
56
Fwald
15.7
9901
64
Neural P
28.7
11195
56
DimeNet++
Baseline
2.8
12013
48
Embeddeds
5.4
13865
40
Cutoff
2,8
48128
8
SchNet-LR
3.7
13813
40
Ewald
4.7
13725
40
Neural P'M
6.4
17191
32
GemNettT
Baseline
14.1
16719
32
Embeddeds
16.1
17643
32
Cutoff
14.1
33792
16
SchNet-L.R
15.0
19131
32
Ewald
15.8
18819
32
Neural PdM
16.8
21

---

H
Ablation Study
We describe our ablation experiments in detail here. We chose the simplest SchNet model and evaluate on the OE62 dataset.
H.1 Empirical Analysis of the Effectiveness of Atom2Mesh & Mesh2Atom Modules
We set the cutoff distance distance between atoms and mesh points y to a constant value of 4.0 and fix the number of discretizations to 3. Subsequently, we incrementally removed either Atom2Mesh, Mesh2Atom, both from the original architecture prevent information exchange between short-range and long-range blocks, thereby assessing the impact of these modules. The results are presented in Table
2.5
2.23
Forward time (sec)
2.0
1.5
0.97
1.0
0.55
0.37
0.5
0.0
b
78.6
Energy MAE (meV)
80
77.5
74.8
75
69.1
70
65
3
5
7
10
Number of meshes
C
76
Energy MAE (mev)
74.7
74
73.3
73.4
73.3
73.3
72.7
72.0
71.9
72
6
7
8
9
10
Assignment cutoff (w/o k-NN Graph)
Figure 4: Relationships between the number of meshes and forward time (a) and energy MAE (b),
as well as the relationship between assignment cutoff without k-NN graph and energy MAE (c).
22

---

H.2 Empirical Analysis of the Number of Mesh Points
We ensure that the cutoff between the atoms and the mesh points aasign is constant (4.0 A) and that the number of discretizations is the same in all three directions, i.e., N₂ N2 N₂. The results of the forward time performance and performance with the number of mesh points are shown Fig.
4 a) and (b).
H.3
Empirical Analysis of the Assignment Cutoff Distance
We conduct ablation studies to examine the impact of the assignment cutoff distance. For the performance results reported in T] setting the maximum number of neighbors to which generally minimizes multiple assignments. To assess the effect of multiple assignments, this ablation experiment uses only the radius graph, varying the assignment cutoff distance from 3 to 10. The performance results based on different assignment cutoff distances shown in Fig. 4cc..
a combination of a radius graph and k-NN graph,
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We have accurately stated the paper's contributions and scope in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have discussed the limitations of our proposed framework in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]

---

Justification:
We have provided mathemtical backgrounds on Ewald summation in Section
and Appendix B
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have disclosed all the needed information for reproducibility in Appendix D F We will open-source for reproducibility once our paper gets published.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
25

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All datasets used in this paper are publicly and freely accessible. We have included sufficient instructions the datasets and our experimental settings in Section We will open-source for reproducibility once our paper gets published.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
mublic//uuiess/Ddeetbmi.sss oos.ssaobsssiooooooooooooo... for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/mipss.cc/pullc/giiddee htt///CooSSSeiiaaeeeeeee ssisssssicyyyyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permirted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have included all details of the model architecture, data processing, and hyperparameter settings in Appendix and for reproducing and understanding our results.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification: We followed previous work report performance on single seed. We also fixed the seed for reproducibility instead of averaging across multiple seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
26

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We have indicated the needed computational resources in Appendix F
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https::/neurrpp: http://hwwrrppp..pppppp..ccceeea
Answer: [Yes]
Justification: Ethics.
The research conducted in this paper conforms with the NeurIPS Code of
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We have discussed the potential societal impacts of our proposed framework
in Section O
Guidelines:
The answer NA means that there is no societal impact of the work performed.
27

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper does not release any high-risk data or models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best. faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have adequately and properly cited and credited the datasets and models used in this papaer.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
28

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
29

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "yRaehqX2iY", "reviewer": "Reviewer_M1Fm", "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.", "questions": "1. What do the \"Embeddings\", \"Cutoff\", and \"SchNet-LR\" mean in Table 2?\n\n2. What is the computational complexity of the proposed method?", "limitations": "The authors have addressed the limitations and potential societal impacts.", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "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. \n\nThe improvements in benchmarks are promising, which demonstrates the power of the proposed model.\n\nThe 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]. \n\n2. 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.\n\n[1] Ewald-based long-range message passing for molecular graphs. ICML 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your recognition of our work’s superior performance and clear organization. We will address your questions and concerns as follows.\n\n**Weakness 1, Novelty & Contribution of Neural P$^3$M**\n\nWhile 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.\n\n**Weakness 2, Additional Experimental Results**\n\nThanks 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).\n\nFor the memory consumption, please see Table 3 and its discussion in the Global Rebuttal for more details. \n\n**Question 1, Concepts in Table 2**\n\nThe 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.\n\n[1] Ewald-based long-range message passing for molecular graphs. ICML 2023\n\n**Question 2, Computational Complexity**\n\nThe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KboIwy9EkW", "reviewer": "Reviewer_7zWV", "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. \nThe 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.\nThe 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.\nNeural 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.", "questions": "1. Does Neural P3M have some restrictions of geometric GNN models? Or it can be combined with most GNNs?\n2. 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?\n3. 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.", "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.\n2. Profiler results and the number of parameters in each variant are not presented.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n2. The framework is built upon well-established principles such as Ewald summation and P3M methods. This theoretical grounding lends credibility to the approach.\n3. 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.\n4. 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.\n5. 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.\n2. Additionally, you can present the P3M Blocks by some pseudocode. \n3. 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.\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**Weakness 1, Mathematical Details**\n\nThank 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.\n\n**Weakness 2, Pseudocode**\n\nThank 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.\n\n**Weakness 3, Computational Complexity**\n\nThank 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.\n\n**Weakness 4, Typos in Table 4**\n\nThanks for pointing these out. We have fixed them in the revised manuscript.\n\n**Question 1, Restriction on GNN**\n\nWe 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.\n\n**Question 2, Computational Bottleneck**\n\nWe 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].\n\n| Base Model | Short Range Forward Time | Long Range Forward Time |\n| --- | --- | --- |\n| SchNet | 0.4797ms | 0.9814ms |\n| PaiNN | 0.8986ms | 0.9827ms |\n| DimeNet++ | 0.9761ms | 0.9702ms |\n| GemNet-T | 2.9914ms | 0.9817ms |\n\n[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.\n\n[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.\n\n[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.\n\n**Question 3, Freezing Common Layers**\n\nFollowing 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.\n\n|  | Test EMAE |\n| --- | --- |\n| Freeze | 83.9 |\n| Finetuned | **68.6** |\n| From scratch | 69.1 |\n\n**Limitations 1, Distributed Training**\n\nOur 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.\n\n|  | Baseline | Cutoff | SchNet-LR | Ewald | Neural P$^3$M |\n| --- | --- | --- | --- | --- | --- |\n| EMAE | 895 | 869 | 984 | 830 | **693** |\n| FMAE | 61.1 | 60.3 | 65.3 | 56.7 | **55.6** |\n\n**Limitations 2, Profiler results**\n\nSee Table 3 and its discussion in the Global Rebuttal for more details.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EWo4zoQ81C", "reviewer": "Reviewer_iSFT", "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.", "questions": "- What’s the reason for choosing VisNet to add Neural P3M in Table 1 instead of Equiformer/MACE which has better baselines?\n- 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?\n- Is it possible to add results with LSRM in Table 2 results as well? Was there a reason for not including that?\n- 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)\n- 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.", "limitations": "Yes, this work does mention its limitations and potential negative impact.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The need for long-range GNN methods is crucial for various molecular representation tasks \n    - This work tackles this pertinent issue utilizing a novel combination of atom and mesh representation\n- The underlying method is well-described\n- The distinction of this approach to other recent and relevant approaches is well-described\n- The paper is well-written and easy to walk through\n- Strong results \n    - Demonstrates improvement over vanilla VisNet on Ag dataset\n    - SOTA results on most targets and molecules of MD22 and gives significant improvements on larger molecules\n    - 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.\n- The anonymous link to the code is not in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank your recognition of our work’s novelty, contribution and clear organization. We will address your questions and concerns as follows.\n\n**Weakness 1, Evaluation & Improvement**\n\nWe resolve the issues point-to-point in the **Questions** Section.\n\n**Weakness 2, Anonymized Code**\n\nWe 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.\n\n**Question 1, ViSNet as the Baseline**\n\nWe 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.\n\n| Molecule |  | Baseline | Neural P$^3$M |\n| --- | --- | --- | --- |\n| Buckyball catcher | Energy | 0.3978 | **0.3038** |\n|  | Forces | 0.1114 | **0.1018** |\n| Double-walled nanotube | Energy |  1.1945 | **0.6208** |\n|  | Forces | 0.2747 | **0.2399** |\n\n**Question 2, Choice of Mesh & Cutoff**\n\nAs 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.\n\n**Question 3, LSRM on OE62**\n\nWe 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.\n\n**Question 4, GPU Memory Usage**\n\nSee Table 3 and its discussion in the Global Rebuttal for more details.\n\n**Question 5, Ablation Study of Atom2Mesh & Mesh2Atom**\n\nWe 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.\n\n| Architecture | Test EMAE |\n| --- | --- |\n| Without Both | 72.07 |\n| Without Mesh2Atom | 76.14 |\n| Without Atom2Mesh | 74.48 |\n| Neural P$^3$M | **69.10** |", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "K9CYJupoJ1", "reviewer": "Reviewer_g8KM", "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.", "questions": "1. Discuss the different choices for the distance cut-off.\n2. 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?\n3. Provide ablation studies on the proposed component and the discussion on the effectiveness of the FFT.", "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.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n- 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.\n- 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.\n- 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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for your recognition of our work’s contribution and clear organization. We will address your questions and concerns as follows.\n\n**Weakness 1, Ablation Study of Atom2Mesh & Mesh2Atom and Efficiency of FFT**\n\nWe 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).\n\n| Architecture | Test EMAE |\n| --- | --- |\n| Without Both | 72.07 |\n| Without Mesh2Atom | 76.14 |\n| Without Atom2Mesh | 74.48 |\n| Neural P$^3$M | **69.10** |\n\nRegarding 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.\n\n**Weakness 2, More Explanation of the Parameterizing Long-Range Interaction**\n\nSorry 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.\n\n**Weakness 3, Ablation Study of Assignment Cutoff**\n\nWe 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.\n\n|  Assignment cutoff (Without k-NN Graph) | Test EMAE |\n| --- | --- |\n| 3 | 74.7 |\n| 4 | 72.0 |\n| 5 | 71.9 |\n| 6 | 72.7 |\n| 7 | 73.3 |\n| 8 | 73.4 |\n| 9 | 73.3 |\n| 10 | 73.3 |\n\n**Questions 1**\n\nSee Weakness 3 for details.\n\n**Questions 2**\n\nSee Weakness 2 for details.\n\n**Questions 3**\n\nSee Weakness 1 for details.\n\n**Limitation 1, Extending to DNA/Protein**\n\nThanks 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to express our sincere gratitude to all the reviewers for dedicating their time to read our manuscript and for offering their valuable suggestions. We appreciate the recognition our manuscript has received from the reviewers. We have also addressed each concern raised on a point-by-point response through a total of 28 experiments and a comprehensive profiling. Here we show some common results for all reviewers.

**Results 1, Performance of LSRM on OE62 Dataset**

Due to the LSRM's limited generalizability and its dependency on the fragmentation algorithm, we have excluded certain molecules from OE62 dataset (validation:285, test:293, train:2387). We present the performance results (in the format of *Valid EMAE/Test EMAE*) of the LSRM on this slightly modified dataset.

|  | Baseline | LSRM | Neural P$^3$M |
| --- | --- | --- | --- |
| SchNet | 133.5/131.3 | 72.9/72.6 | **70.2/69.1** |
| PaiNN | 61.4/63.3 | 56.6/56.4 | **54.1/52.9** |
| DimeNet++ | 51.2/53.8 | 47.9/50.4 | **40.9/41.5** |
| GemNet-T | 51.5/53.1 | 50.8/51.5 | **47.2/47.5** |

**Results 2, Performance of Ewald MP on MD22 Dataset**

We present the performance results (*MAE*) of Ewald MP combined with ViSNet on MD22 dataset.

| Molecule |  | Baseline | Ewald | Neural P$^3$M |
| --- | --- | --- | --- | --- |
| Ac-Ala3-NHMe | Energy | 0.0796 | 0.0775 | **0.0719** |
|  | Forces | 0.0972 | 0.0814 | **0.0788** |
| DHA | Energy | 0.1526 | 0.0932 | **0.0712** |
|  | Forces | 0.0668 | **0.0664** | 0.0679 |
| Stachyose | Energy | 0.1283 | 0.1089 | **0.0856** |
|  | Forces | **0.0869** | 0.0976 | 0.0940 |
| AT-AT | Energy | 0.1688 | 0.1487 | **0.0714** |
|  | Forces | 0.1070 | 0.0885 | **0.0740** |
| AT-AT-CG-CG | Energy | 0.1995 | 0.1571 | **0.1124** |
|  | Forces | 0.1563 | 0.1115 | **0.0993** |
| Buckyball catcher | Energy | 0.4421 | 0.3575 | **0.3543** |
|  | Forces | 0.1335 | 0.0989 | **0.0846** |
| Double-walled nanotube | Energy | 1.0339 | 0.7909 | **0.7751** |
|  | Forces | 0.3959 | 0.2875 | **0.2561** |

**Results 3, Profiling results of Neural P$^3$M**

We present the number of parameters and memory usage (with standard settings and a batch size of 8 of the largest molecule in OE62) as well as the maximum batch size that can be accommodated on a single A100 GPU in the following table. The bulk of the memory usage is still attributed to the short-range modules—for instance, 16719 MB versus 19945 MB in GemNet. As anticipated, the integration of the mesh concept and additional modules means that Neural P$^3$M has a higher parameter count and slightly greater memory usage than Ewald MP. Nevertheless, this modest increase in resource demand is offset by the significant performance improvements offered by Neural P$^3$M, along with the computational acceleration brought by FFT.

|  |  | # of Parameters | GPU Memory (MB) | Max Batch Size |
| --- | --- | --- | --- | --- |
| SchNet | Baseline | 2.8M | 1623 | 400 |
|  | Embeddings | 14.4M | 1865 | 344 |
|  | Cutoff | 2.8M | 1671 | 392 |
|  | SchNet-LR | 5.3M | 4835 | 128 |
|  | Ewald | 12.2M | 2675 | 240 |
|  | Neural P$^3$M  | 19.1M | 2283 | 280 |
| PaiNN | Baseline | 12.5M | 8135 | 80 |
|  | Embeddings | 15.7M | 9073 | 72 |
|  | Cutoff | 12.5M | 20480 | 32 |
|  | SchNet-LR | 15.1M | 11289 | 56 |
|  | Ewald | 15.7M | 9901 | 64 |
|  | Neural P$^3$M  | 28.7M | 11195 | 56 |
| DimeNet++ | Baseline | 2.8M | 12013 | 48 |
|  | Embeddings | 5.4M | 13865 | 40 |
|  | Cutoff | 2.8M | 48128 | 8 |
|  | SchNet-LR | 3.7M | 13813 | 40 |
|  | Ewald | 4.7M | 13725 | 40 |
|  | Neural P$^3$M | 6.4M | 17191 | 32 |
| GemNet-T | Baseline | 14.1M | 16719 | 32 |
|  | Embeddings | 16.1M | 17643 | 32 |
|  | Cutoff | 14.1M | 33792 | 16 |
|  | SchNet-LR | 15.0M | 19131 | 32 |
|  | Ewald | 15.8M | 18819 | 32 |
|  | Neural P$^3$M | 16.8M | 19945 | 32 |

Meta-review: This paper consider the incorporation of mesh points into geometric GNNs for materials science. The reviewers largely agree that this is a valuable contribution that is well-supported, and I believe concerns have been adequately addressed. I accordingly recommend acceptance.
Decision: Accept (poster)