GemNet: Universal Directional Graph Neural Networks for Molecules
Authors: Johannes Gasteiger
Florian Becker
Stephan Günnemann
Venue: NeurIPS 2021 Poster (2021)
Keywords: GNN
graph neural network
molecules
equivariance
point clouds
weight initialization
normalizer-free
message passing

Abstract: Effectively predicting molecular interactions has the potential to accelerate molecular dynamics by multiple orders of magnitude and thus revolutionize chemical simulations. Graph neural networks (GNNs) have recently shown great successes for this task, overtaking classical methods based on fixed molecular kernels. However, they still appear very limited from a theoretical perspective, since regular GNNs cannot distinguish certain types of graphs. In this work we close this gap between theory and practice. We show that GNNs with directed edge embeddings and two-hop message passing are indeed universal approximators for predictions that are invariant to translation, and equivariant to permutation and rotation. We then leverage these insights and multiple structural improvements to propose the geometric message passing neural network (GemNet). We demonstrate the benefits of the proposed changes in multiple ablation studies. GemNet outperforms previous models on the COLL, MD17, and OC20 datasets by 34%, 41%, and 20%, respectively, and performs especially well on the most challenging molecules. Our implementation is available online.

GemNet: Universal Directional Graph Neural Networks for Molecules
Johannes Gasteiger, Florian Becker, Stephan Günnemann
Technical University of Munich, Germany {j gasteiger, beckerf, faakkerf,gennenninee wgw.gnueennnneenee... tum. de
Abstract
Effectively predicting molecular interactions has the potential to accelerate molec- ular dynamics by multiple orders of magnitude and thus revolutionize chemical simulations. Graph neural networks (GNNs) have recently shown great successes for this task, overtaking classical methods based on fixed molecular kernels. How- ever, they still appear very limited from theoretical perspective, since regular GNNs cannot distinguish certain types of graphs. In this work we close this gap between theory and practice. We show that GNNs with directed edge embeddings and two-hop message passing are indeed universal approximators for predictions that are invariant to translation, and equivariant to permutation and rotation. We then leverage these insights and multiple structural improvements to propose the geometric message passing neural network (GemNet). We demonstrate the benefits of the proposed changes in multiple ablation studies. GemNet outperforms pre- vious models on the COLL, MD17, and OC20 datasets by 34 %, 41 %, and %, respectively, and performs especially well on the most challenging molecules. Our implementation is available online.
1
Introduction
Graph neural networks (GNNs) have shown great promise for predicting the energy and other quantum mechanical properties of molecules. They predict these properties orders of magnitudes faster than methods from quantum chemistry at comparable accuracy. GNNs can thus enable the accurate simulation of systems that are orders of magnitude larger. However, they still exhibit severe theoretical and practical limitations. Regular GNNs are only as powerful as the I-Weisfeiler Lchman test of isomorphism and thus cannot distinguish between certain molecules 45 require a large number of training samples to achieve good accuracy.
60
Moreover, they
In this work we first resolve the questionable expressiveness of GNNs by proving sufficient conditions for universality the case invariance to translations and rotations and equivariance to permutations; and then extending this result to rotationally equivariant predictions. Simply using the full geometric information (e.g. all pairwise atomic distances) in a layer does not ensure universal approximation. For example, if our model uses rotationally invariant layer we lose the relative information between components. Such a model thus cannot distinguish between features that are rotated differently. This issue is commonly known the "Picasso problem": " An image model with rotationally invariant layers cannot detect whether a person's 's eyes are rotated correctly. Instead, we need a model that preserves relative rotational information and is only invariant global rotations. To prove universality the rotationally invariant case we extend a recent universality result based on point cloud models that representations of the rotation group SO(3) [8]. We prove that spherical representations are actually sufficient; full SO(3) representations are not necessary. We then generalize this to rotationally equivariant predictions by leveraging recent result extending invariant equivariant
https://www daml. in. tum. Geneentet
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

predictions |57. We then discretize spherical representations by selecting points on the sphere based on the directions to neighboring atoms. We can connect this model to GNNs by interpreting these directions as directed edge embeddings. For example, the embedding direction of atom would be defined by atom e, resulting in the edge embedding Cex- Updating the spherical representation of atom based on atom b then corresponds to two-hop message passing between the edges Cca and edb via Cba. with atoms c and defining the embedding directions. This message passing formalism naturally allows us to obtain the molecule's s full geometrical information (distances, angles, and dihedral angles), and the direct correspondence proves the model's universality.
We call this edge-basett based two-hop message passing scheme geometric message passing, and propose multiple structural enhancements to improve the practical performance of this formalism. Based on these changes we develop the highly accurate and sample-efficien geometric message passing neural network (GemNet). We furthermore show that stabilizing the variance of GemNet's activations with predetermined scaling factors yields significant improvements over regular normalization layers.
We investigate the proposed improvements in a range of ablation studies, and show that each of them significantly reduces the model error. These changes introduce little to no computational overhead over two-hop message passing. Altogether, GemNet outperforms previous models for force predictions on COLL by 34 %, on MD17 by 41 %, and on OC20 by 20 % on average. We observe the largest improvements for the most challenging molecules, which exhibit dynamic, non-planar geometries. In summary, our contributions are:
Showing the universality of spherical representations and two-hop message passing with directed edge embeddings for rotationally equivariant predictions.
Geometric message passing: Symmetric message passing enhanced by geometric information.
Incorporating all proposed improvements in the Geometric Message Passing Neural Network (GemNet), which significantly outperforms previous methods for molecular dynamics prediction.
2
Related work
Machine learning potentials. Research on predicting a molecule's 's energy and forces (so-called machine learning potentials) started with hand-fitted analytical functions and then gradually moved towards fully learned models. Arguably, classical force fields are their very first instances. They use analytical functions with coefficients that were hand-tuned based experimental data. A popular example for these is the Merck Molecular Force Field (MMFF94) 31I. The next wave of methods used kernel ridge regression based on fixed, hand-crafted molecular representations 3 g 23 Finally, modern research mostly focusses on fully end-to-end learnable models based on GNNs 30 [3]] These models can also be combined with molecular features from quantum mechanical caleulations to improve performance [48]. We consider this combination as orthogonal research.
Directional GNNs. We can also achieve equivariance and invariance rotations without relying on group representations. Directional GNNs achieve this by representing directional information explicitly [54] or in the form of angles 29 and dihedral angles 25 33]. Our work is focused this class of models, proving their universality and proposing an improved variant, GemNet.
Expressiveness of GNNs. A large part of GNN research has been focused on their (limited) ex- pressiveness. Morris et al. 45). Xu al. 660 first proved that they are only as expressive as the Weisfeiler-Lehman test of isomorphism and Garg et al. (27 showed the limitations of basic direc- tional message passing. Kondor et al. [37] Maron al. (3 Morris et al. [45] 46) then investigated higher-order representations to vircumvent this issue. Finally, Azizian & Lelarge
(2), Maron et al.
42 showed that so-called folklore GNNs are the most expressive GNNs for a given tensor order.
Equivariant neural networks. Equivariance and invariance have recently emerged as one of the foundational principles of modern neural networks [33 [5]. This is especially relevant for models in physics, for which we often know the symmetries a priori. Equivariant models for the SO(3) group were first investigated
the context of spherical convolutions by Cohen et al. [14, Esteves al. 21,, Kondor et al. (36). These methods leverage group representations to achieve full equivariance. They were then transferred to the context of 3D point clouds and molecules by Anderson et al.
1. Thomas et al. [56). Weiler et al. (58), and further developed by Batzner et al. 4. Finzi et al.
24 Fuchs et al. 226] Importantly, Yarotsky (6T) proved the universality of 2D convolutional
networks, and Bogatskiy et al. [5] extended this result to general groups. Maron et al. [4] proved

---

universality for models invariant to S, and equivariant to an additional symmetry. Dym & Maron [ combined these results to prove universality for the joined group of translations, rotations, and permutations. Apart from reflections this is the exact group relevant for general molecules.
3
Universality of spherical representations
GNNs for molecules typically incorporate directional information in one of two ways: Via (0(3) representations [ 56] or by using directions real space [29 (4 Directions in real space are associated with the three-dimensionaa Sa sphere, while the S3)3) group is double covered by the four- dimensional S sphere. Directional representations thus use one degree of freedom less than SO(3) representations, making them significantly cheaper. And, as we will prove in this section, directional representations actually provide the same expressivity as SO(3) representations for predictions R We achieve this by showing that the SO(3)-base) tensor field network (TFN) [56] variant used by Dym & Maron [18] is equivalent to a similar model based on spherical representations, in the case of rotationally invariant predictions. We then generalize recent result by Villar et al. [57], which lets us extend our theorem to the rotationally equivariant case. Afterwards, we relate this universality to directional GNNs by interpreting them discretization of spherical representations.
Preliminaries. We consider a point cloud with n points (atoms), each associated with a position and set of rotationally invariant features (e.g. atom types), defined as X E R and Hin IR In this section we define model classes by sets of functions F. As a first step, we are interested in proving that the set F defining our model is equal to the full set of functions G that are invariant to the group of translations T and rotations SO(3), and equivariant to the group of permutations S.. We denote the codomain of functions in G WT, where WT is some representation of SO(3). We denote a vector's norm by z = || ||2. its direction by & = x a/z, and the relative position by Pba 2b Ta. Proofs are deferred to the appendix. Note that this section is not intended as an introduction to the (3(3) group. For a concise introduction in the context of machine learning see e.g. Weiler al. (58] Section 3] Kondor al. (36]
Tensor field network. In order to show the equivalence of the TFN to spherical representations, we first need to define this model. Following Dym & Maron [18, we split the model into two parts: Embedding functions Freu that lifts the input into an equivariant representation, and pooling functions F pool that aggregate the results of multiple embedding functions on each point and computes the model output. The overall model is then defined as the set of functions
K
FUHN = = {f 1 H.) = K(D)D
(k) pool
(X, Hmm)), plad E FULN (D), (H) E F (D)},
Ueat
(1)
where D E N denotes the function's maximum polynomial degree, K(D) E N is chosen such that Theorem I is fulfilled (Dym & Maron [18] only prove the existence of this function), and f" denotes elementwise application of on all points. We then define the set FTPN pool as all rotationally equivariant linear functions on the SO(3) group, i.c all SO(3) convolutions (38). Note that these are more expressive than the self-interaction layers used originally (56). The embedding functions
JTFN (D) = { f(2D) oo f()) ((r E F consist of an auxiliary function T22 X, H) = H Teat
and a series of tensor product functions (called convolution by Dym & Maron (18) FIEN {
f(x, H) = (X, FFN X, H))}- The intermediate representations are H E Wh where Wrent is representation of SO(3) indexed by the degree / and the order m. For H we have I= = 0. The main update is defined by
TPN(L.) amo
H) = (H
U.) amo
Σ Σ
Σ beN
F(ll)) T (I'- Eg H. bime?
(2)
fyimy L.m,
where 0 a (learned) scalar and N₂, are the neighbors of point a. The Blebsch-Gordan coefficients Cllysma) ((p.m))((l.mm)) arise from decomposing the tensor product of two input SO(3) representations (the of output representations. Their exact values are not
filter and input representations) into relevant for this discussion. We index the output with degree l, and order m., the learned filter with
ly and my and the input with 4 and m,, m (x) R0O (z) Yim (x) is rotationally equivariant
3

---

filter, with (learned) radial part R, which is any polynomial of degree S D, and the real spherical harmonics Ym with degree I and order m. The spherical harmonics are the basis for the Fourier transformation of functions on the sphere, analogously to sine waves for functions on R. We can prove universality for TENs by using the universality of polynomial regression and showing that TFNs can fit any polynomial (see Dym & Maron [T8] for details), resulting in:
Theorem (Dym Maron (8).. Consider the set of functions G mapping Raxattxx + WT that equivariant rotations and permutations and invariant translations. For all n E E.
1. For D Ng, every polynomial p of degree
FIFN D(D)D
2. Every continuous function f can be approximated uniformly on compact sets by functions in UDEND F-ENN KDD.DD
Spherical networks. Instead of intermediate SO(3) representations we now switch to spherical
representations, which are functions on the sphere H S" + We define the set of functions jSphere
R(D).D analogously to FIEN K(D),D' However, for juphere (D) we use
Feat
Husphere X (f) (H) +
Faphere Dbb La. r) Hp))),
(3)
BEN.
with the filter function F'sphere (x, 8) Σi, Hdhh (x (x)
Ry (#) (f)), using the real part R the complex spherical harmonics y- The set of pooling functions for invariant predictions is
spheee pool
(f ((H) = -ppool { dF))
(4)
with the learnable parameter pant- We obtain the universality theorem by showing the equivalence between this model and TFN for rotationally invariant functions. The proof based on the connection between spherical harmonics and the CLebsch-Gordan coefficients [51, 3.7.72] (sce|App. Ap.
Theorem 2. Consider the set of functions g' mapping IR 133x++xx - WT that are equivariant to permutations and invariant translations and rotations. For all n E N.
1. For D No. every polynomial G' of degree D is in p-here K(D),D 2. Every continuous function G' can be approximated uniformly compact sets by functions in UDEN. prsphere (D),DD
Next, extend Theorem 2 rotationally equivariant functions. We do this by generalizing recent result by Villar et al. [57] to obtain (spe B):
Theorem 3. Let h: pdxn+hxn
dd be any function that is equivariant to permutations and rotations and invariant to translations. For all E [I, let the set of relative vectors {æea E
[L, n not span a (d - 1)-dimensional space. Then there are n functionn f(e); pdxx+txan - RR" such that
hag x H = Σ [X) H He
(5)
c=1
cafa
where fle) is equivariant to permutations, but invariant rotations and translations.
This theorem lets us extend rotationally invariant model to an equivariant one, while providing universality guarantees. Together, Theorem and Theorem (with d 3) thus show that we can approximate any rotationally equivariant function using only representations on the S sphere. We thus do not need SO(3) representations, spin- wiighweghted spherical harmonics (22), triplet embeddings, or complex- -valued functions. This result puts theory back line with practice, where the best results currently achieved without relying on these more expensive representations [54]
4 From spherical representations to directional message passing
Directional representations. To use spherical representations in a model we first need find tractable description. Instead of using spherical harmonics, we propose to sample the representations

---

in specific directions f.. If we look at recent models, we see that they implicitly use the directions to each atom's 's neighbors for this purpose, i.e. they embed the edges in the molecule's 's graph. These directions define an equivariant mesh that circumvents the aliasing effects that would arise from fixed grids [36]. Schütt et al. [54] flexibly define the directional mesh in each layer by aggregating directions, while Gasteiger et al. [29] and others use a fixed mesh for each atom. We can refine this mesh of directions e.g. by using more neighbors or by interpolating between directions. The approximation error of this directional mesh is related to the spherical harmonic expansion via the mesh norm and the separating distance between directions [33 34 Note that depending on the discretization scheme the resulting mesh might not provide a universal approximation guarantee.
Eq. (3) only defines the relationship for a fixed direction, while models commonly use different directional meshes for the input and output. To incorporate this we add a convolution with learned filter F2, which can only improve the model's 's expressiveness. Since the input and output are spherical functions, the used filter F2 has to be zonal, i.e. it can only depend on one angle. This can be
expressed as 17
H (X, (fx))) (H.(Fo)+ [So(3) BEN. (Ho +
Σ
Esphere (acbere(Xiinn (Ihaa R()
Σ
Ho P.) (r)) dR
VERs
Σ B. RR
Σ
Esphero (Tba, Hn
(6)
where Rs denotes the directional mesh of atom b with mesh directions denoted by f.s and i. specifies the output direction. The integral vanishes due to the Dirac delta 8.
General filters. To see the relationship to GNNs we furthermore need to generalize the filter Fophere (Dba, f.). This filter only depends on the angle 0r since it is rotationally invariant:
Lemma 1. Fsphere Rx, Rf) = Fshhere (z, f) for any rotation matrix R.
We can therefore substitute Fsphere with general learnable filter Fi that is parametrized by this relative angle. Since Fsphere arises as a special case we do not lose expressivity. We thus obtain
H (X (X (H.
Σ E
Fi LF, F2 F2 (L
bEN MRRS
We have now arrived at message passing scheme that has universal approximation guarantees and is only based on relative directional information. To see the connection to GNNs we interpret these discretized spherical representations as edge embeddings pointing towards t. and f.. Eq. (7) then corresponds to two-hop message passing between the edge embeddings of c. and A. via the edge Tba. Interestingly, the central learnable part of of. (7) is the product of the filters Fi (xba. Lř, and F2 (t) with the input representation, which is strikingly similar to the Hadamard product used in modern GNNs
28 53
except that these only one-hop message passing.
5 Geometric message passing
Geometric representation. We now develop a specific two-hop message passing scheme based [Eq (7) We use embeddings based on interatomic directions, and embed all atom pairs with distance Fea < Cemb- i, and are thus instantiated the inter- atomic directions &c,a and a We denote directional embeddings as mcc = H. (Tea). Message passing is thus based on quadruplets of atoms - two atoms are interacting (a and b) and two atoms define the directions (c and d). We denote the angle between directions by Pabd = L To improve empirical performance we addition- ally use the dihedral angle 6eabd ai 1 Tba and substitute Lioi; = Lt. with Peab- Fig. I illustrates the three angles Peab, Pabds and Orube we use for updating the embedding m based on m.n.- To ensure that all angles are well-defined we ex- clude overlapping atom quadruplets, i.e. of d. We represent the relative directional information using spherical Fourier-
d
mea
mab
Ocabel
poals
Patet
a.b
(rotate)
Figure 1: Angles used in geo- metric message passing. The di- hedral angle Beaba becomes vis- ible when rotating the molecule so that atoms a and b lie on top of each other (right).
5

---

bases with polynomial radial envelopes to ensure smoothly differentiable predictions, as proposed by Gasteiger al. [29]. We split the basis into three parts to incorporate all available geometric information. Before the envelope, these are:
2
sin
nm dbbb)
ĒRBF, (2db)
V
Feub
(8)
Cemb
Itdb
ZIn
ēCBF,In In (Ibas (aaaa) = V
jil
Iba (io
(9)
3 2
Rip. 1 (ztn)
Gint
2
21u
CSBF,mm (Fea: Peub, leabá)
ji(
Ica) Yta) (Peat, Beabd),
(10)
L (2 Comb
with the interaction cutoff Cints the spherical Bessel functions ji. and the n-th root of the l-order Bessel function 3in. Note that Gasteiger et al. [29] only used the first two parts ERBF and ECBF. These representations are then transformed using two linear layers to obtain the filter F. In order to maintain smoothly differentiable cutoff we cannot use bias in this transformation. Altogether, the core geometric message passing scheme is
Σ bEA"HMtch deNe ((a.c)
SBFI eesreessFFeee Peob, leabs)) "eatt))" (W((WBBF2 eCBF Pabd))
maa
(11)
@ @WRBFZ WRBFICRBF @ mai))
where W denotes a weight matrix, W denotes a weight tensor. The first weight matrix of each repre- sentation part has small output dimension. This causes bottleneck that improves generalization.
Symmetric message passing. Whenever we have directional embedding Mca. also have the opposing embedding mac since both are based on the same cutoff Cemb- Whether we associate the embedding Mcaa mac with atom is arbitrary. A more principled approach is to jointly interpret both embeddings as a representation of the atom pair a and c. this view, an update to M should also influence mac- This would normally require executing the above message passing scheme twice, once for updating m,a. based on m.,. and once for updating mac based on m.d.. We propose to eircumvent this double execution by calculating the update (Eq. (1I) only once and then using it for both m.n and mac- To preserve the distinction between the two directions and ensure that mea * mae- we transform the two updates using two separate learnable weight matrices. One single message passing update thus carries information for both embeddings, which is then dissected by the two weight matrices. In practice, this only requires a simple re-indexing operation that maps the edge ca to ac
Efficient bilinear layer. The whole message passing scheme, i.e. basis transformation, neighbor aggregation, and bilinear layer, only use linear functions. We can therefore freely optimize the order of summation without changing the result, as proposed by Wu al. (59 (see Appp. D for details). Doing so can provide a faster and more memory-efficient model, reducing memory usage by 50 % even for Hadamard products. Moreover, since everything is based on efficient matrix products, this allows us to use the bilinear layer at practically no additional cost compared to Hadamard product. Note that this requires using padded matrices instead of the usual gather-scatter operations to prevent excessively large intermediate results.
6 GemNet: Geometric message passing neural network
GemNet. The geometric message passing neural network (GemNet) is significantly refined architecture based on DimeNet++ (28) Hippocratic license 2.1]. GemNet predicts the molccular energy E' and forces F 3g based the atomic positions X e 3 and the atomic numbers # E N The architecture illustrated in Fig. 2 A comprehensive version with low-level layers and hyperparameters is described in App. F GemNet was developed on the COLL dataset, but generalizes to other datasets such as MD17 without architectural changes. Every change we propose either improves model performance biases since we found them to be irrelevant or even detrimental to accuracy. We show the impact of the most relevant changes via ablation studies in Sec. 7
reduces model complexity. For example, GemNet uses no
Interactions. GemNet incorporates three forms of interactions. The first is geometric message passing, as described in Sec. 5) The second is one-hop form of geometrie message passing. This

---

Model:
Interaction:
Message passing:
Q/T-MP:
nohd
Riot
Jix Poob oz
Fook
CRBF
Max
1) Ag 1)
Mas
RRS
daylo
(1-1)
RBF
CBF
e SBF pe
EDWF
MP
(WO)
WD
+28
o(W
esur
(co) Cuun
CBF
(colad) SBF
Residual
T-MP
Coli) CHI
CCBF
Atom emb
Wn WD
Q-MP
$2/64
Embedding
hn
esBF
The
Icabd) SBB
Interaction
o (WD)
self infinneraciion
WO
PPW
Interaction
Σ
Residual
Residual
32/66
Interaction
Interaction
OWWO
(W,
Atom emb
128
128
m
WC
20
00
mass
Figure 2: The GemNet architecture (comprehensive version App. F). denotes the layer's input, concatenation, and a non-linearity. Directional embeddings Mea are updated using three forms of interaction: Two-hop geometric message passing (Q-MP), one-hop geometric message passing (T-MP), and atom seff-interactions. Differences between Q-MP and T-MP are denoted by colors and dashed lines.
interaction uses a single cutoff c Cemb and passes messages only between directional embeddings pointing towards the same atom, similarly DimeNet [29]. This provides both angle-based based pair interactions and atom self-interactions, thanks to the symmetric message passing scheme described in Sec. 5 The third interaction is a pure atom self-interane interacttaat based atom embeddings. We first aggregate the directional embeddings of one atom to obtain an atom embedding. We then use this atom embedding update all directional embeddings. We found all three interaction forms to be beneficial, and show this in our ablation studies.
Stabilizing activation variance. The variance of activations in model is usually stabilized using normalization methods, which has various positive effects on training [6 41 [2]. However, they also have multiple undesirable side effects, especially in the context of molecular regression. Batch normalization introduces correlations between separate molecules and atoms. Layer normalization forces all activation scales to be constant, while atomic interactions actually cover large range of scales directly bonding atoms have substantially stronger interaction than atoms at a long range. To circumvent these issues, stabilize GemNet's variance by introducing constant scaling factors, as suggested by Brock et al. 6 We found that the activation variance is primarily impacted by four components: Skip connections, non- non-linearities message aggregation, and Hadamard/bilineet layers. The two summands in a skip connection = x + f(x) have no covariance at initialization due to random weight matrices. We can thus remove its impact by scaling the output by 1/v2. We remove the non-linearity's impact by scaling its output with gain of y 1/0.6 for SiLU, similarly to (35). Note that we do not center SiLU's output but instead choose slightly lower 7 to account for mean shift. Additionally, we standardize the weight matrices have exactly zero mean and 1/fan-in variance. The sum aggregation and Hadamard/bilinnar layers have more complex impact on the variance, which we cannot determine priori (see App. for details). We therefore estimate the variance after these layers based on random batches of data. We then rescale their output accordingly to obtain roughly the variance of the layer input at initialization. These simple empirical scaling
factors are sufficient to keep the activation variance roughly constant (see Fig. 3) We found that
other measures such as adaptive gradient clipping 77, scaled weight standardization 6 or weighting the residual block with zero at initialization [16 not beneficial for model accuracy.
7

---

Table I: MAE on COLL, in meV/Á and meV. GemNet is 34 % more ac- curate for forces. The higher energy error due to its lower loss weight.
Table 2: Force MAE for MDI7@CCSD in V/À. GemNet
outperforms previous methods by 44 % on average.
sGDML NequlP GemNet-Q GemNet-T
Aspirin
33.0
14.7
10.4
10.3
Forces Energy
Benzene
1.7
0.8
0.7
0.7
SchNet
172
198
Ethanol
15.2
9.4
3.1
3.1
DimeNet ++ +
40
47
Malonaldchyde
16.0
16.0
6.0
5.9
GemNet-Q
26.4
53
Toluene
9.1
4.4
2.5
2.7
GemNet-T
31.6
60
GemNet-dQ
38.1
60
GemNet-dT 43.1
55
GemNet-Q and GemNet-T. Geometric message passing is comparatively expensive since it is based on quadruplets of atoms. Its runtime thus scales with (nnhinhems) where kint the number of interacting neighbors, and kemb is the number of embedded directions. For this reason we investigate two message passing models in our experiments one with two-hop geometric message passing (GemNet-Q) and one using only the two cheaper forms of interaction (GemNet-T), Their complexities are O Onhhimkamb) and O(nkemb): respectively. Note that GemNet-T is thus a direct ablation of the two-hop message passing scheme implied by our theoretical results.
Direct force predictions. GemNet predicts forces by calculating F. - OE/Ox, via backpropa- gation. This form of calculation guarantees conservative force field, which is important for the stability of simulations. However, by using [qq. (5) we can also directly predict forces and other vector quantities. This essentially means predicting a magnitude for each directional embedding and then summing up over the vectors defined by this magnitude and the embedding's associated direction, similarly to Park et al. [47. We denote this variant as GemNer-dQ and GemNer-dT. Interestingly. GemNet thus able to generate rotationally equivariant predictions despite only using invariant representations. Direct predictions substantially accelerate the model, especially for training. For most datasets, the resulting accuracy on par with most previous models, but significantly worse than GemNet's accuracy via backpropagation. However, this is not true for OC20, where we found GemNet-dT to converge faster and perform on par with GemNet-T.
Limitations. GemNet is focused on one specific, important task: Predictions for molecular simula- tions. We do not make any statements regarding its performance beyond this scope. The GemNet architecture might seem more complex than some previous models, due to its larger variety of interac- tions and blocks. However, its number of parameters and training or inference time is actually on par with previous models. Two-hop message passing introduces significant computational overhead. We mitigate this effect with a down-projection layer and additionally introduce the ablated GemNet-T model. This model performs surprisingly well on MD17, but not on COLL. This suggests that one-hop message passing is expressive enough for some practical use cases, but two-hop message passing gives an advantage for the more challenging task of fitting multiple molecules at once.
Societal impacts. Accelerating molecular simulations can have positive effects in a wide range of applications in physics and chemistry. At the same time, however, this can be used for malicious purposes such as developing chemical agents or weapons. To the best of knowledge, this work does not promote these use cases more than regular chemistry research does. To somewhat mitigate negative effects we will publish our source code under the Hippocratic license [19].
7 Experiments
Experimental setup. We evaluate our model on four molecular dynamics datasets. COLL 28 CC-BY 4.0] consists of configurations taken from molecular collisions of different small organic molecules. MD17 99] consists of configurations of multiple separate, thermalized molecules, consid- ering only one molecule at a time. MD17@ CCSD [I0] uses the same setup, but calculates the forces using the more accurate and expensive CCSD or CCSD(T) method. The open catalyst (OC20) dataset g CC-BY Y 4.0] consists of energy relaxation trajectories of solid catalysts with adsorbate molecules. This dataset split into three tasks: (1) Structure to energy and forces (S2EF), which the same task used by the COLL and MD17 datasets, (2) initial structure to relaxed structure (IS2RS), where
8

---

Table 3: Force MAE for MD17 in V/À. GemNet outperforms all previous methods by wide margin, on average by 41 %.
Kernel methods
GNNs
GemNet
sGDML FCHL19 DimeNet SphereNet NequlP PaiNN GemNet-Q GemNet-T
Aspirin 29.5
20.7
21.6
18.6
15.1
14.7
9.4
9.5
Benzene
8.1
7.7
8.1
6.3
6.3
Benzene 10
2.6
2.3
1.5
1.4
Ethanol
14.3
5.9
10.0
9.0
9.0
9.7
3.8
3.7
Malonaldehyde
17.8
10.6
16.6
14.7
14.6
14.9
6.9
6.7
Naphthalene
4.8
6.5
9.3
7.7
4.2
3.3
2.2
2.4
Salicylic acid
12.1
9.6
16.2
15.6
10.3
8.5
5.4
5.5
Toluene
6.1
8.8
9.4
6.7
4.4
4.1
2.6
2.6
Uracil
10.4
4.6
13.1
11.6
7.5
6.0
4.5
4.2
Table 4: Results for the three tasks of the open catalyst dataset (OC20), averaged across its four test sets. GemNet outperforms all previous methods in all measures, average by 20 %.
*DimeNet" -large uses separate models for energy and force prediction for IS2RE.
S2EF
IS2RS
IS2RE
Energy MAE Force MAE Force cos meV. 4
AFbT
ADwT
Energy MAE
meV/A +
+
% 1
%
meV
ForceNet-large 336.3
31.2
0.520
12.7
49.6
DimeNet - -large*
31.3
0.544
21.8
51.7
559.1
SpinConv GemNet-dT
29.7
0.539
16.7
53.6
434.3
292.4
24.2
0.616
27.6
58.7
399.7
an energy optimization carried based the model's predictions and measure how close the final structure is to the true relaxed structure (average distance within threshold, ADwT) and whether the final forces are close zero (average forces below threshold. AFbT), and (3) initial structure to relaxed energy (IS2RE), where we predict the energy of the relaxed structure, based on an energy optimization starting at the initial structure. All presented OC20 models trained on the S2EF data. Following the setup of Batzner et al. [ we use 1000 training and validation configurations for MD17, and 950 training and 50 validation configurations for MDI7@ CCSD. We focus on force predictions and use high force loss weight since they determine the accuracy of molecular simulations.
measure the mean absolute error (MAE), averaged over all samples, atoms, and components. We compare with the results reported by several state-of-the-arr models: sGDML [10], FCHL.19 [12, SchNet [53], DimeNet
(29). DimeNer++ 28), SphereNet [39], NequlP
44. PaiNN (54) ForceNet (32, and SpinConv (55) For further details see App. G
Results. Tables 4 show that GemNet-T and GemNet-Q consistently perform best all molecular dynamics datasets investigated and by large margin. This is true both in comparison to previous GNNs and for kernel methods - despite the latter typically being more sample efficient. The improvements are largest for chain-like molecules, such as ethanol and malonaldehyde, These molecules are the most challenging since they exhibit wide range of movement. GemNet even performs better than some previous models that were trained with 50x more training samples. For example, performs better than SchNet with 000 training samples on out of eight MD17 molecules (see Table 9b. Interestingly. the two-hop message passing scheme implied by our theoretical results (GemNet-Q)) yields significant improvements on COLL, but performs approximately on par with the ablated GemNet-T on MD17. To investigate this disagreement combined dataset all MD17 molecules. Table 13 shows that GemNet-Q again performs better than GemNet-T in this setting. These results suggest that regular MD17 is too simple to show the benefits of two-hop message passing. It seems to be particularly important in more difficult settings that cover large variety of configurations and molecules.
trained GemNet
Computational aspects. GemNet-Q is roughly two times slower than GemNet-T (see Table 14). Thanks to the efficient aggregation, GemNet with bilinear layers is as fast with regular Hadamard products. Efficient aggregation also reduces the memory usage for regular Hadamard products by around 50% & (from 4. G.I to 2.2GB for a batch of 32 Toluene molecules). Note that GemNet has not

---

been optimized for runtime and can likely be accelerated substantially. GemNet-Q uses 2.2M and GemNet-T 1.9M parameters. which is comparable to previous models such DimeNet + which uses 1.9M parameters. See I for further details.
Direct force prediction. Directly predicting the forces accelerates training by four times average and inference by 1.6 times on average in our experiments (see [Tble 14), while reducing memory consumption roughly factor of two. While using direct predictions instead of backpropagation increases the MAE by 44 % on COLL and by 48 % on MD17 (see Tables I and 77. they actually perform better on the S2EF task on OC20. This likely due to OC20 being orders of magnitude larger than COLL and MD17. Whether to use direct predictions thus depends on the dataset and the application's 's computational requirements.
Ablation studies. We investigate the pro- posed architectural improvements COLL in Table 5) The proposed symmetric message passing scheme yields significant accuracy im- provements, as does using a bilinear layers in- stead of Hadamard product. We also see that removing any of the three interaction forms described in [Sec. 6jnnreases the error, show- ing that this combination is indeed beneficial. The proposed scaling factors also yield de- cent improvements, while regular layer nor- malization actually increases the error. Two- hop message passing yields the largest single improvement. Taable 10 shows that our archi- tectural improvements yield similar benefits for DimeNer++, Overall, the error improve- ments are quite evenly distributed. This suggests that GemNet's improved performance is not due to one single change, but rather due to the full range of improvements proposed in this work.
Table 5: Ablation studies COLL. Force MAE in meV/Á after 500 000 training steps. All proposed components yield significant improvements.
Model
Forces
without symmetrie message passing
28.5
Hadamard product instead of bilinear layer
29.3
without atom embedding updates without one-hop message passing
28.3
31.3
without two-hop message passing
32.4
without scaling factors use layer norm instead (without centering)
29.1
33.3
with bias GemNet-Q
27.2
27.0
8 Conclusion
In this work we proved the universality for GNNs using directional embeddings. We proposed geometric message passing based on these insights, and improved this scheme with symmetric message passing and efficient bilinear layers. We incorporated these improvements in the GemNet architecture, which substantially improves the eтror on various molecular dynamics datasets. We showed that all of the proposed enhancements yield sigmificant performance improvements. Most of our proposed improvements are of independent interest for other molecular GNNs.
Acknowledgments and Disclosure of Funding
We would like to thank Soledad Villar for help with proving Lemma as well as Nicholas Gao and Aleksandar Bojchevski for their invaluable feedback and comments.
This research was supported by the Deutsche Forschungsgemeinschattt (DFG) through the Emmy Noether grant GU 1409/2-1 1 and the TUM International Graduate School of Science and Engineering (IGSSE), GSC 81.
References
[1] Brandon M. Anderson, Truong-Son Hy, and Risi Kondor. Cormorant: Covariant Molecular Neural Networks. In NeurlPS, 2019.
[2] Waiss Azizian and Mare Lelarge. Expressive Power of Invariant and Equivariant Graph Neural Networks. In ICLR, 2020.
[3] Albert P. Bartók, Mike C. Payne, Risi Kondor, and Gábor Csányi. Gaussian Approximation Potentials: The Accuracy of Quantum Mechanics, without the Electrons. Physical Review Letters. 104(13):136403, 2010.
10

---

[4] Simon Batzner. Tess E. Smidt. Lixin Sun. Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, and Boris Kozinsky. SEq3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials, arXiv, 2101.03164, 2021.
[5] Alexander Bogatskiy, Brandon Anderson, Jan Offermann, Marwah Roussi, David Miller, and Risi Kondor. Lorentz Group Equivariant Neural Network for Partiele Physies. In ICMI., 2020.
[6] Andrew Brock, Soham De, and Samuel L. Smith. Characterizing signal propagation to close the perfor- mance gap in unnormalized ResNets. In ICLR, 2021.
[7] Andrew Brock, Soham De, Samuel L. Smith, and Karen Simonyan. High- High-Porformance Large-Scale Image Recognition Without Normalization. arXiv, 2102.06171, 2021.
[8] Lowik Chanussot, Abhishek Das, Siddharth Goyal, Thibaut Lavril, Muhammed Shuaibi, Morgane Riviere, Kevin Tran, Javier Heras-Dominges Caleb Ho. Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. Open Catalyst 2020 (OC20) Dataset and Community Challenges. ACS Catalysis, 11(10):6059-600 6072, 2021.
[9] Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt,, and Klaus- Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, 1633x16 2017.
[10] Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-learred force fields. Nature Communications, 9(1):1-10. 2018.
[11] Anders Christensen and Anatole von Lilienfeld. On the role of gradients for machine learning of molecular energies and forces. Machine Learning: Science and Technology, 1(4):045018, 2020.
[12] Anders S. Christensen, Lars A. Bratholm, Felix A. Faber, and O. Anatole von Lilienfeld. FCHL revisited: Faster and more accurate quantum machine learning. The Journal of Chemical Physics, 152(4):044177, 2020.
[13] Taco Cohen and Max Welling. Group Equivariant Convolutional Networks. In ICML, 2016.
[14] Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling, Spherical CNNs. In ICLR, 2018.
[15] Taco Cohen, Mario Geiger, and Maurice Weiler. A General Theory of Equivariant CNNs on Homogeneous Spaces. In NeurlPS, 2019.
[16] Soham De and Sam Smith. Batch Normalizatinn Biases Residual Blocks Towards the Identity Function in Deep Networks. In NeurlPS, 2020.
[17] J. R. Driscoll and D. M. Healy. Computing Fourier Transforms and Convolutions on 2-Sphere. Advances in Applied Mathematics, 15(2):002-250, 1994.
[18] Nadav Dym and Haggai Maron. the Universality of Rotation Equivariant Point Cloud Networks. In ICLR, 2021.
[19] Coraline Ada Ehmke. The Hippocratic License 2.1: An Ethical License for Open Source., 2020. URL https://fixrstddonnarr.
[20] Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Siemoid-weighted linear units for neural network function approximation reinforcement learning. Neural Nenworks, 107:3-11, 2018.
[21] Carlos Esteves, Christine -lanchela,,,,tte, Ameesh Makadia, and Kostas Danilliiis. Learning SO(3) Equivariant Representations with Spherical CNNs. In ECCV, 2018.
[22] Carlos Esteves, Ameesh Makadia, and Kostas Daniilidis. SeighWeighted Spherical CNNs. In NeurlPS. 2020.
[23] Felix A. Faber, Anders S. Christensen, Bing Huang, and O. Anatole von Lilienfeld. Alchemical and structural distribution based representation for universal quantum machine learning. The Journal of Chemical Physics, 148(24):241717, 2018.
[24] Mare Finzi, Samuel Stanton. Pavel Izmailov, and Andrew Gordon Wilson. Generalizing Convolational Neural Networks for Equivariance Lie Groups Arbitrary Continuous Data. In ICML, 2020.
[25] Daniel Flam- Shepherd, Tony C. Wu, Pascal Friederich, and Alan Aspuru-Guzia Neural Message Passing High Order Paths. Machine Learning: Science and Technology, 2021.
11

---

[26] Fabian Fuchs. Daniel Worrall, Volker Fischer, and Max Welling. Trass)-ransformer::: 3D Roto-Transsation Equivariant Attention Networks. In NeurIPS, 2020.
[27] Vikas Garg, Stefanie Jegelka, and Tommi Jaakkola. Generalization and Representational Limits of Graph Neural Networks. In ICML., 2020.
[28] Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and Uneertainty- Aware Directional Message Passing for Non-Equilibrium Molecules. In Machine Learning for Molecules Workshop, NeurIPS, 2020.
[29] Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional Message Passing for Molecular Graphs. In ICLR, 2020.
[30] Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural Message Passing for Quantum Chemistry. In ICML, 2017.
[31] Thomas A. Halgren. Merck molecular force field. I. Basis, form, scope, parameterization, and performance of MMFF94, Journal of Computational Chemistry, 17(5-6):490-519 1996.
[32] Weihua Hu, Muhammed Shuaibi, Abhishek Das, Siddharth Goyal, Anuroop Sriram, Jure Leskovee, Devi Parikh, and C. Lawrence Zitnick. ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations. arXiv, 2103,01436, 2021.
[33] Kurt Jetter, Joachim Stöckler, and Joseph Ward. Error estimates for scattered data interpolation on spheres. Mathematics of Computation, 68(22)):73--777, 1999.
[34] Jens Keiner, Stefan Kunis, and Daniel Potts. Efficient Reconstruction of Functions on the Sphere from Scattered Data. Journal of Fourier Analysis and Applications, 13(4):435-458, 2007.
[35] Günter Klambauer. Thomas Unterthiner, Andreas Mayr, and Sepp Hochreiter. Self- celf-normalizing neural networks. In NeurlPS, 2017.
[36] Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network. In NeurlPS, 2018
[37] Risi Kondor, Truong Son Hy. Horace Pan, Brandon M. Anderson, and Shubhendu Trivedi. Covariant Compositional Networks For Learning Graphs. In International Workshop Mining and Learning with Graphs (MLG), KDD, 2019.
[38] Peter Kostelec and Daniel N. Rockmore. FFTs the Rotation Group Journal Fourier Analysis and Applications, 14(2):145-179, 2008.
[39] Yi Liu, Limei Wang, Meng Liu, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical Message Passing for 3D Graph Networks. arXiv, 2102.05013, 2021.
[40] Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. ICLR, 2018.
[41] Ping Luo, Xinjiang Wang. Wenqi Shao, and Zhanglin Peng. Towards Understanding Regularization in Batch Normalization. In ICLR, 2019.
[42] Haggai Maron, Heli Ben- Hem-Hamu, Hadar Serviansky, and Yaron Lipman. Provably Powerful Graph Networks. In NewrIPS, 2019,
[43] Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman. On the Universality of Invariant Networks. In ICML, 2019.
[44] Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya. On Learning Sets of Symmetrie Elements. In ICMI., 2020.
[45] Christopher Morris, Martin Ritzert, Matthias Fey, William L. Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and Leman Go Neural: Higher-Order Graph Neural Networks. In AAAI, 2019.
[46] Christopher Morris, Gaurav Rattan, and Petra Mutzel. Weisfeiler and Leman go sparse: Towards scalable higher-order graph embeddings. In NewrIPS, 2020.
[47] Cheol Woo Park, Mordechai Kornbluth, Jonathan Vandermause, Chris Wolverton, Boris Kozinsky, and Jonathan P. Mailoa. Accurate and scalable graph neural network force field and molecular dynamics with direct force architecture. npj Computational Materials, 7(1):199, 2021.
12

---

[48] Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and Thomas F. Miller. OrbNet: Deep learning for quantum chemistry using cunmetry-adapteed atomic-orbital features. The Journal of Chemical Physics. 153(12):12411, 2020.
[49] Zhuoran Qiao, Anders S. Christensen, Matthew Welbom, Frederick Manby, Anima Anandkumar, and Thomas Miller III. UNITE: Unitary N-body Tensor Equivariant Network with Applications to Quantum Chemistry. arXiv, 2100.11655,, 2021.
[50] Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the Convergence of Adam and Beyond. ICLR, 2018.
[51] Sakurai San Fu Tuan. Modern Quantum Mechanics. Revised, subsequent edition edition, 1993.
[52] Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normalization help optimization? In NeuriPS, 2018.
[53] Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. SchNet: A continuous-files convolutional neural network for modeling quantum interactions. In NeurlPS, 2017.
[54] Kristof T. Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In ICML, 2021.
[55] Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C. Lawrence Zitnick. Rotation Invariant Graph Neural Networks using Spin Convolutions. arXiv, 2106.09575, 2021.
[56] Nathaniel Thomas, Tess Smidt, Steven M. Kearnes, Lusamn Yang, Li, Kai Kohlhoff, and Patrick Riley. Tensor Field Networks: Rotation- and Translation-Equivrrian Neural Networks for 3D Point Clouds. arXiv, 1802.08219, 2018
[57] Soledad Villar. David W. Hogg. Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith. Scalars are universal: Equivariant machine learning, structured like classical physics. In Neur/PS, 2021.
[58] Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data. In NeurlPS, 2018.
[59] Wenxuan Wu, Zhongang Qi, and Li Fuxin. PointConv: Deep Convolutional Networks on 3D Point Clouds. In CVPR, 2019.
[60] Keyulu Xu, Weihua Hu, Jure Leskovee, and Stefanie Jegelka. How Powerful are Graph Neural Networks? In ICLR, 2019.
[61] Dmitry Yarotsky. Universal Approximations of Invariant Maps by Neural Networks. Constructive Approximation, 2021.
13


Reviews:
[{"review_id": "ywgskNKnopI", "reviewer": "Reviewer_Uuo9", "summary": "The paper proposes GemNet, an extension to DimeNet++ for predicting molecular properties. The additions include two-hop message passing, symmetric message-passing and efficient bilinear layers. GemNet achieves state of the art results on the MD17 and COLL benchmarks.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "The paper is well motivated and structured. Some aspects could be explained in more detail, in particular the \"symmetric message-passing\" scheme in Sec 5. The experiments and ablation studies demonstrate the performance and careful design choices of GemNet.\n\nThe sections on the universality of spherical  and directional message-passing seem a bit misleading to me. If I'm not mistaken, the proofs are only valid for general functions $F_\\text{sphere}$ / $F_\\text{TFN}$. However, in practice one needs to introduce a cutoff. In this case, the \"Picasso problem\" mentioned in the introduction becomes an issue again, if one only uses rotationally-invariant representations. Thus, the central issue is not just the universality of the architecture in theory, but the tradeoff between rotational symmetries of the representation (angular momentum of spherical harmonics) and locality. A discussion on this and how this might be connected to the two-hop scheme would be interesting.\n\nAnother minor issue is the discussion of inference time and model size:  the authors state that those are \"on par\" with previous models. While this is true for DimeNet, it is not the case for NequIP and PaiNN, which are much faster and smaller, while outperforming DimeNet. The timing of Toluene in Table 9 for NequIP (100ms, bs 4) also seems a bit high compared to the original paper (16ms, bs 1). It would be interesting to see also the accuracy and performance of a scaled-down GemNet with a comparable number of parameters.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Clarify practical implications of universality proof (size of cutoff etc., see above).", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "NtSBAB7JbP", "reviewer": "Reviewer_Nzov", "summary": "Many GNN models have recently been proposed for the problem of predicting molecular dynamics. These GNNs rely on rotationally invariant predictions since the underlying physics of the problem does not change with rotations. While such models have obtained good performance, there is limited theoretical understanding of these methods.\n\nThis paper makes two contributions to this task: first, it addresses the theoretical limitations by proving that GNNs with directed edge embeddings and two-hop message passing are universal approximators; second, it presents a new model called GemNet (an extension of Dimenet++ that uses two hop message passing), that obtains excellent results on COLL and MD17 datasets.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "## Originality\nThe theoretical contributions in the paper appear novel and interesting. The GemNet model builds upon the Dimenet and Dimenet++ models, but introduces some novel ideas, particularly the use of dihedral angles.\n\n## Quality\nThe paper appears technically sound. The theoretical claims made by the authors have been proven rigorously and the new model is shown to outperform previous methods experimentally on MD17 and COLL datasets. The new model is a natural extension to the theoretical claims, which provides a good backing for the theory.\n\nStabilizing activation variance and efficient interactions that the authors use in the GemNet paper are orthogonal to the other improvements and they can also be used in Dimenet++. It would be good to include ablations that incorporate these improvements into Dimenet++ so we can more clearly see the benefits of including dihedral angles.\n\n## Clarity\nThe paper was difficult to follow and suffers from mathiness. I would recommend the authors to try and simplify the presentation. \n\n## Significance\nThe methods presented in the paper are likely to be useful for many molecular chemistry problems.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors include a discussion of the limitations of their method.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "6P20eEr78Dz", "reviewer": "Reviewer_LzdH", "summary": "The authors propose a model Geometric Message Passing Neural Network (GemNet) which leverages directed edge embeddings and a new two-hop message passing scheme. The model is more expressive and performs convincingly better than the most related baseline (DimeNet) as well as other baselines across a large range of datasets.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "The model contributions seem theoretically sound (subject to my limited expertise on the topic) and in a way the model feels like a natural extension to DimeNet to include torsion angles. The paper is well written, with sufficient information for reproducibility, and the results compared to previous baselines (and specially compared to DimeNet) are compelling. While I did not go into the details, I appreciate the proof of universality of representations for rotationally invariant models, which is great way to qualify the expressivity of the model. Finally, the problem it tackles has very direct applications in science problems, and is of very high interest in state of the art research in rotationally invariant models.\n\nI think my main criticism is about the comparison between GemNet-Q and GemNet-T. Authors say “This suggests that one-hop message passing is expressive enough for some practical use cases, but two-hop message passing gives an advantage for the more challenging task of fitting multiple molecules at once”. As a potential user, it would be great to have a bit more insight (e.g. some rule of thumb) of when to use one or the other. Perhaps some investigation using synthetic datasets would have helped conveying this. For example considering that GemNet-Q, is technically more expressive than GemNet-T, in the cases where GemNet-Q performs worse, is this due to overfitting, or due to optimization problems?\n\n--------------\n\nPost rebuttal:\n\nThank you for your responses and additional experiments and clarification on the comparison between GemNet-Q and GemNet-T. I would like to maintain my original rating, which was already very positive.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See main review.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2"}, "extra_text": {}}, {"review_id": "3t9_yHU-73k", "reviewer": "Reviewer_YFgs", "summary": "* The paper proposes to use a type of networks that can be shown to be universal approximators for functions that are permutation equivariant, rotationally and translationally invariant.\n* This property is required to model functions that work on molecular structures. In essence, this provides a very strong prior on the functions we want to learn.\n* The experiments show that by enforcing this structure learning on molecular graphs can be improved (higher quality/label efficiency) compared to previous work. While, the experimental section is quite compact and heavily compressed nothing stood out as obviously wrong but I might be missing a detail.\n\nOverall I like the high level idea behind the paper and the results seem good. While, at a high level I am able to understand the paper, I am really struggling to understand the details and I do not feel like I would be able to reproduce the results or understand the nuances of the architecture. I am aware that the code is released, but this should not be a replacement for a clearly written paper. For this reason I suggest a score of marginally below the acceptance threshold. In the main review I have highlighted the issues and points were I got stuck. \n\n", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "main_review": "\n## Section 3\nThe tensor field network definition is a notational mess and I am struggling to properly understand the actual structure/meaning.\n* Eq. 1: K and D are not defined K(D) properly, $K,D\\inN$ is provided but not what they stand for.\n* Eq. 2: what does $l_o$ and $m_o$ stand for?\n* Eq. 2: indexing of the clebsch-gordan coefficients is not clear to me. (And a wikipedia search does not help either since it has a different notation). I digged deeper in the references in this paper but unfortunately the TFN paper refers to Griffiths (Introduction to quantum mechanics) which is 147 dollars on amazon at the day of the review and a bit to expensive/cumbersome to order and read for this review. Given that many people at Neurips might be in a similar situation I would recommend the authors to actually explain this properly. \n* Eq. 3. $F_{sphere}=\\sum_{l,m}$ what is l and m? Is R^l a part of the learned polynomial function.  Can you better describe how I should think about Y.  Looking at this is seems similar to a fourier basis but for functions defined on the sphere? It is not immediately obvious how to think about the function on the sphere. Eq. 3 depends on a direction $\\hat{r}$ and $x_b-x_a$, which is then reduced into a direction and a magnitude component. Providing an intuition for how to think about such a function would be helpful for the general ML researcher to better understand this paper.\n\n## Section 4\n* Eq. 5 what is $F_2$? What is $r_o$ and $r_i$.\n\n## Section 5:\n* There is a definition of $r_o$ and $r_i$. This would have been helpful in the previous section.\n* Figure 1, right hand side is not helpful. The flat arrow going in a circle does not really mean much to me. Why is a and a,b darker?.\n\n## Figure 2\n* Residual is not defined. are all the W's the same? What is the atom emb. What is W_arrowup, what is W_arrowdown. What is z. How is the first mca initialized? \n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Clarity", "review_text": "", "extra_scores": {"ethical_concerns": "Nothing obvious.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5-6 possibly more."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper was generally well received. Reviewers unanimously agreed that the paper was interesting and impactful, especially given the widespread interest in message passing neural networks to model atomistic systems. It seems clear that NeurIPS is an appropriate venue for this publication. 

The authors provided a number of clarifications and further experiments in response to the reviews that helped to address several concerns (in particular regarding baselines against existing models such as NequIP and ablation studies). The main issues that could not be addressed during the rebuttal period were the clarity of exposition which several reviewers commented on. I would encourage the authors to work on simplifying the paper as much as possible in the time leading up to the camera ready. 
Decision: Accept (Poster)