So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems
Authors: Thorben Frank
Oliver Thorsten Unke
Klaus Robert Muller
Venue: NeurIPS 2022 Accept (2022)
Keywords: equivariance
self-attention
message passing
molecules
force fields
non-local

Abstract: The application of machine learning methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab-initio methods. However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which are often neglected due to the difficulty of modeling them efficiently. This work proposes a modified attention mechanism adapted to the underlying physics, which allows to recover the relevant non-local effects. Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher-order geometric information for each atom in a molecule, enabling a non-local formulation of attention in the SPHC space. Our proposed model So3krates - a self-attention based message passing neural network - uncouples geometric information from atomic features, making them independently amenable to attention mechanisms. Thereby we construct spherical filters, which extend the concept of continuous filters in Euclidean space to SPHC space and serve as foundation for a spherical self-attention mechanism. We show that in contrast to other published methods, So3krates is able to describe non-local quantum mechanical effects over arbitrary length scales. Further, we find evidence that the inclusion of higher-order geometric correlations increases data efficiency and improves generalization. So3krates matches or exceeds state-of-the-art performance on popular benchmarks, notably, requiring a significantly lower number of parameters (0.25 - 0.4x) while at the same time giving a substantial speedup (6 - 14x for training and 2 - 11x for inference) compared to other models.

So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems
J. Thorben Frank
Oliver T. Unke 1233
Klaus-Robert Müller
Machine Learning Group, TU Berlin, 10587 Berlin, Germany
BIFOLD, Berlin Institute for the Foundations of Learning and Data, Germany
Google Research, Brain team, Berlin
Department of Artificial Intelligence, Korea University, Seoul 136-713, Korea Max Planck Institut für Informatik, 66123 Saarbrücken, Germany
Abstract
The application of machine learning methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab-initio methods. However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which often neglected due to the difficulty of modeling them efficiently. This work proposes a modified attention mechanism adapted to the underlying physies, which allows recover the relevant non-local effects. Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher-order geometric information for each atom in molecule, enabling non-local formulation of attention in the SPHC space. Our proposed model SO3KRATES a self-attention based message passing neural network uncouples geometric information from atomic features, making them independently amenable to attention mechanisms. Thereby we construct spherical filters, which extend the concept of continuous filters in Euclidean space to SPHC space and serve as foundation for spherical self-attention mechanism. We show that in contrast other published methods, $03KRATE, is able describe local quantum mechanical effects over arbitrary length scales. Further, we find evidence that the inclusion of higher-order geometric correlations increases data efficiency and improves generalization. S03KRATES matches or exceeds state- of-the-art performance on popular benchmarks, notably, requiring significantly lower number of parameters (0.25-0.4x) while at the same time giving a substantial speedup (6-14x for training and 2-1 1x for inference) compared to other models.
1 Introduction
Atomistic simulations use long time-scale molecular dynamics (MD) trajectories to predict macro- scopic properties that arise from interactions on the microscopic scale [T-]]. Their predictive reliability is determined by the accuracy of the underlying force field (FF). which needs to be queried every time step. This quickly becomes computational bottleneck if the forces are determined from first principles, which may be required for accurate results. To that end, machine learning
field
" froggiinn.rnnngggggmiii....mmmm
I Vkaus-robertmmullleeeu--betineee
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

(n) Design Principle
(b) Soakrates architecture
Current approaches
Proposed approach
shared represenation for invariant fratures and equivariant information
information exchange attention
represcutations with
Festure
SPHC initialization
initialization
invarlant
invariant
self- attention
sele-anttntoon
spherical
lapers
SoDkrates Layer
equivariant
equivation:
So3krates Layer
reduce complexity
(c) Model Sixe and Accuracy
(d)
Training and Inference Time
parannless
SM
model siate nisan Brror
training
musm crror
0.4
inference
2M
time
IM
0.2
NewtonNet PAINN SoSkrates DimoNet GemNett Spooky NequiP
0.0
0.0
Solkrates
DinneNet GemNetQ" NoqulP
(e) Spherical Harmonic Coordinates (SPHCs)
(f) Localization in SPHC'
(g) Nou-local effects cumulene
Layer
Layer
Layer
fromt view
far apazt
Atom
c000000
view
SPHCC-epaco
Local Model
ab-inilio
Atom
enesgy (6V)
So3krates
close
Atom
135 dthedral angle @ (*)
180
Figure 1: (a) Design principle our proposed peessageppassingg scheme, where invariant features and equivariant information represented by spherical harmonic coordinates (SPHCs) X separate and exchange information via (spherical) self-attention. (b) Overview of the Sо3KRATES architecture. (c) Comparison of prediction accuracy and model size of different architectures. (d) Comparison of training and inference times of different architectures. (e) Illustration of learned SPHCs different layers of the So3KRATES architecture. (f) Low- oow-dimeniional projection of atomic SPHCs, showing that atoms far apart in Euclidean space can mapped close together SPHC space. (g) In contrast to local models, So3KRATER is able to learn the long-range correlations between hydrogen rotors in cumulene and reproduces the ab-initio ground truth faithfully.
FFs (MLFFs) offer computationally more efficient, yet accurate empirical alternative expensive ab-initio methods 2 24
In recent years, Geometric Deep Learning has become popular design paradigm, which exploits relevant symmetry groups of the underlying learning problem by incorporating a geometric prior 25 26.. This effectively restricts the learnable functions of the model to subspace with a meaningful inductive bias. Prominent examples for such models are e.g. convolutional neural the group of translations, graph neural networks 28 which are invariant W. r. node permutation.
12
networks (CNNs) [27]. which equivariant
(GNNs)
For molecular property prediction, it has been shown that equivariance group (3(3) greatly improves data efficiency and accuracy the learned FFs equivariance, architectures either rely on feature expansions in terms of spherical harmonics (SH) [ or explicitly include (dihedral) angles 29 s2
the 3D rotation
29 32
To achieve
While the latter scales quadratically (cubically) in
the number of neighboring atoms and has been shown to be geometrically incomplete
451 the
calculation of spherical harmonics scales only linear in the number of neighboring atoms, which
makes them fast and accurate alternative 30 B2 36
As our own timings were measured on different GPU, we decreased the reported times according greedup-factorg reported in B7 For full details, appendix A.6
Reference times were taken from 34

---

However, current higher-order order geometric representations based on SHs usually result in expensive transformations. since an individual feature channel per SH degree (and order) required. As result, going to higher degrees is computationally expensive and comes at the price of increasing complexity, resulting in state-of-the-arr (SOTA) models with millions of parameters 30 order to be applicable to large molecular structures, models are required to be both efficient and accurate on all length scales.
R2 34 However, in
Non-local electronic effects have been outlined as one of the major challenges for a new generation of MLFFs They result in non-local, higher-order geometric relations between atoms. Most current architectures implicitly assume locality of interactions (expressed through local neighborhood), which prohibits an efficient description of all relevant atomic interactions at larger scales. Simply increasing the cutoff radius used to determine local neighborhoods is not an adequate solution, since only shifts the problem to larger length scales (30]
In this work, we propose spherical harmonic coordinates (SPHCs), which encode higher-order geometric information for each node in a molecular graph (Fig. Te). This is in stark contrast current approaches, which consider molecules as three-dimensional point clouds with learned features and fixed atomic coordinates: we propose make the SPHCs themselves learned quantity. Through localization in the space of SPHCs (Fig. D]] models are able to efficiently describe electronic effects that are non-local in dine-dimmesiooal Euclidean space (Fig. Ig.
We then present SO3KRATES (Fig. [D), a self-attention based message-passinng neural network (MPNN), which decouples atomic features and SPHCs and updates them individually (Fig. [а). This resembles ideas from equivariant graph neural networks [25]. but allows to go to arbitrarily high geometric orders. The separation of higher-order geometric and feature information allows to overcome the parametric and computational complexity usually encountered in models with higher- order geometrie representations, since we only require a single feature channel (instead of one per SH degree and order). Thus, S03KRATES resembles some early architectures like SCINET I0 or in parametric simplicity. We further show that SO3KRATES outperforms the popular 38 kernel model by a large margin in the low-data regime, domain which has so far been considered to dominated by kernel machines (21]. Numerical evidence suggests that the data efficiency of So3KRATES is directly related to the maximal degree of geometric information encoded in the SPHCs. We then apply SO3KRATES our model achieves SOTA results, despite is light- structure and having only 0.25 0.44xx the number of parameters of competitive architectures (Fig. [.), while achieving speedups of 6-1 4x and 2-11x for training and inference, respectively (Fig. Id).
PHYSNET
1
SGDML
the well-established MD17 benchmark and show that
Although we focus on quantum chemistry applications in this work, the developed methods also applicable to other fields where long-ranged correlations in three-dimensional data are relevant. For example, models based on SPHCs may also be applicable to tasks like 3D shape classification or computer vision.
2 Preliminaries and Related Work
In the following, we review the most important concepts our method is based upon and relate it to prior work.
carry over many of the benefits of convolutions to unstructured domains and have thus become one the most promising approaches for the description of molecular properties. Their general working principle relies the repeated iteration of message passing (MP) steps, which can be phrased as follows [4] 25
Message Passing Neural Networks MPNNs 14
m:j = =mi(f Lj, r())
m; - my
(2)
JEN(()
f u(f,, m.).
Here, mij is the message between atoms and j computed with the message function mn-), m, is the aggregation of all messages in the neighborhood, (i)) of atom i, and u(-) an update function returning updated features f. based on the current features f; and message m, The neighborhood
3

---

(i) consists of all atoms which lie within a given cutoff radius around the atomic position тæ, which ensures linear scaling in the number of atoms n. While earlier variants parametrized messages only in terms of inter-atomic distances 13 19. more recent approaches also take higher-order geometric
information into account 25
Molecules as Point Clouds A molecule can be considered as a point cloud of n atoms PaD (R, F), where R (T.)...Tx denotes the set of atomic positions r; E R2 and F (fi.... f,) is the set of rotationally invariant atomic descriptors, or features, f; E R We write the distance vector pointing from the position of atom i the position of atom j as ray ry Ts, the distance rúj = | and the normalized distance vector as faj = Tij/Taj. Given the point cloud, a density over Euclidean space assigning a vector value to each point can be constructed as
p(r)
Σ 8(1r: - (((l) f,
(4)
where 5 is the Dirac delta function. It can be shown that applying convolutional filter on p(r) resembles the update steps used in MPNNS (T).
Equivariance Given a set of transformations that act on a vector space A as Sy : + to which
associate an abstract group G, function f f: à + B is said to be equivariant w. Gf
f(S, = - (x),
(5)
where To : 1B + B is an equivalent transformation on the output space Thus, in order to say that I is equivariant, it must hold that under transformation of the input, the output transforms "in the same way". While equivariance has been a popular concept in signal processing for decades (cf. c.g.
42 or wavelet neural networks 43 recent years have seen efforts to design group
equivariant NNs and kernel methods, since respecting relevant symmetries builds an important 277 which are equivariant w. t. translation, GNNs 14 28 which are invariant (T, T) w. r. permutation, or architectures which are equivariant t. the (3) group 2 Kll
inductive bias 12
45 Examples are CNNs
56 A6 In this work, we consider the SO(3) group of rotations, such that A is the Euclidean space 3 where the corresponding group actions are given by rotation
33
matrices Rg E 1x333
Spherical Harmonies The spherical harmonics are special functions defined on the surface the sphere S = {i E R²; : = }} and form an orthonormal basis for the irreducible representations (irreps) of S3)).. In the context of tensor field networks [33], they have been introduced as elementary building blocks for 50(3)-equivarian
neural networks. The spherical harmonics are commonly denoted as y 2 r> R. where the degree I determines all possible values of the order
m E (-I...... -1}. They transform under rotation as
r1 (Re f) Dmmm (Ro) Yi (f),
(6)
where D D' (Ro) are the entries of the Wigner-D matrix D'(Ro) E ((2++))((++1)
47
Based on
the spherical harmonics, we define vector-valued function Y S + 22++1 for each degree
I. with entries Y.M for all valid orders m of given degree l. Since ( F) = D' Y ( (F)
(cf. eq. 6 Y (1) is equivariant w. SO(3).
Tensor Product Contractions The irreps (h) and (12) can be coupled by computing their tensor product (11) @ (12),, which can equivalently be expressed as a direct sum 33 48
(1 x22)
hitt +
y()) (()) =
(1
1==
where the entry of order m3 for the coupled irreps (1s) is given by
4
ada m3
Σ
dadaa
yd yls
(8)
mi mi
my
=
mg

---

and Chi-aala mm. mag. are the so-called Trebsch-Gordon coefficients. In the following, we will denote the tensor product of degrees I and l2 followed by "contraction" l3 (meaning the irreps of degree 13 in the direct sum representation of their tensor product) as (Y (41) 0t r((a)). which is mapping of the form R (222t+)))))) + 2 since my E fl.,.....3))
3 Methods
In the following, we describe the main methodological contributions of this work. We introduce the concept of an adapted point cloud PAD (R. X F).., which incorporates the set of spherical harmonics coordinates (SPHCs) X' (xi,.... Xn) (see below) in addition to features F and Euclidean coordinates R. However, contrary to R. SPHCs X' are refined during the message passing updates. Having SPHCs as part of the molecular point cloud extends the idea current MPNNs, which learn message functions on R, only. Instead, we learn a message function m (cf. eq. i on both, the (fixed) atomic coordinates R as well as on the SPHCs X. This adapted paasage-passing scheme allows to learn non-local geometric corrections. Based these design principles, propose the SO3KRATES architecture.
Initialization
Feature vectors are initialized from the atomic numbers 2; N (denoting which chemical element atom belongs to) by embedding map
f. femb (z)),
(9)
where femb N + RF. We define SPHCs as the concatenation of degrees C == {Imin Lmax
x
(min)
Umax
IR (lmax Lmin +1)?
(10)
@RaImia
ER 2imia
such that their transformation under rotation be expressed in terms of concatenated Wigner-D matrices (see appendix A. The short-hand (9) E 2R2++ refers to the subset of SPHCs with degree They are initialized as
(rew( - (Pig) (EN(())
Σ
(11)
C
where C, EsEN(() Pram (riz). Pr.m over the neighborhood. (i) of atom
+ R is the cosine cutoff function and the sum runs
Message Passing Update Two branches of attention- weighted MP steps are defined for the feature vectors f and SPHCs x (see Fig. I After initialization (eqs. and (I1b,, the features are updated as
=
Proa (riz) Ozj fzy (EN(()
Σ
(12)
where Ony RR self-attention it is possible to define an MP update for the SPHCs as
4 50
coefficients (see below). In analogy the feature vectors,
x! +
Σ JEN(()
Dram (ra)). a. (ff)))
(13)
where individual attention coefficients iy E R for each degree of the SPHCs are computed using multi-head attention 449 However, with this definition, both MP updates are limited local neighborhoods N (i). To be able to model non-local effects, we introduce the SPHC distance matrix X E R with entries Xij =lix: Xill2. i.e. distances between two atoms and in SPHC space for all possible pair-wise combinations of n atoms. To have uniform scales, we further apply the softmax along each row of X to generate rescaled matrix X = softmax (X) with entries Xzj- A polynomial cutoff function Φx... [29 is then applied to X to define spherical neighborhoods. Nx (i) (see A.2 which may include atoms that are far away Fuclidean space (see Fig. IT. The spherical
cutoff distance is chosen Xeut 1/n to ensure that spherical neighborhoods remain small, even
5

---

when going to larger molecules. We then incorporate non-local geometric corrections into the MP update of the SPHCs
Xi +
Prew (a)).. (fa)+ JEN
2xom (x (x)))... . (fO)) (14)
JEN(()
X- but nom-local in 3t We will show in the first part of the experiments, how geometric corrections from SPHC space allow for modelling non-local quantum effects, inaccessible to current architectures. In the second part, we use SO3KRATES model without geometric corrections, which makes it a traditional MPNN in the sense of only localizing in R We find this architecture to be highly parameter, data and time efficient while capable of reaching SOTA results.
local R
Spherical Filter and Self-Attention as
The self-attention coefficients in eqs.
are calculated
0i (w.r (wiy ff)//FF
(15)
where Wij E IR is the output of a filter generating function and "O" denotes the element- product. The filter maps the Euclidean distance Tky and per-degree SPHC distances
(1) = llx; x; ll2
between the current SPHCs of atoms and j into the feature space IF (as a short-hand, we write the vector containing all per-degree SPHC distances as Lxi lec)). is built as the linear combination of two filter-genrrating functions
Wij (Tas) + ds (VE))
(16)
radial filter
spherical filter
which separately act on the Euclidean and SPHC distances. We call d IR - IR.. the radial filter function and o. R + IR the spherical filter function (an ablation study for és can be found in appendix A.4 Since per-atom features f,, interatomic distances rij. and per-degree distances Xij are invariant under rotations (proof in appendix I.II. so are the self-attention coefficients aij.
(I)
While we choose to pass the per-degree norms directly into the filter generating function d., future work might explore the possibilities of alternative metrics (instead of the L2 norm) or an expansion in terms of basis functions as it is common practice for inter-atomic distances (see appendix A.3 eq. 30
Atomwise Interaction according to
After each MP update, features and SPHCs are coupled with each other
f. = + (fe, (xi (2) leEc, (EC))EC).
(17)
x; + + f.. 1x 10) NEEe, 1x (xE") Xi + t02 (Ix IEE)) (1)
X, X;
(18)
where Φ1 F+2 1 R 02 F+2 R. and 03
(1) R R. the inputs to 01,2,3,
degree-wise scalars x (1) = xlx U
Il2 are used to preserve equivariance. The coupling step additionally
includes cross-degree coupled SPHCs x! for each degree 1. Following
48 they are constructed as
Imax
kit (x:
(11) 31 X:
(1a)
(19)
11 Ii=inn
1g=1++1
where ,2,1 E are learnable coefficients for all valid combinations of 41 la given I and the term in brackets is the contraction of degrees l and la into degree I (eq.
S03KRaTES architecture Using the design paradigm above, we build the transformer network So3KRATES, which consists of a self-attention block on F and X' (eqs. (12) and (3)), respectively. as well as an interaction block (eqs. 17 and (8)) per layer. After initialization of the features and the SPHCs according to eqs. and layers. Atomic energy contributions E, are predicted from the features of the final layer using a two-layered output block. The individual contributions are summed to the total energy prediction F F= Σ. F,, See Fig. D for an overview. More details on the implementation, training details and network hyperparameters are given in appendix A.3 and
they are updated iteratively by passing through my

---

(a) Overcome the effective cutoff induced by message passing
eutoff
eff. cutoff
cutoff eff. cutoff
cutoff
eff. cutoff
pooos
oooo 0000
1.2
Nequll' (1as=1)
1.2
Nequll' (uss-1 NequlP (lxx
1)
ab-initto
1.2
Noqull' (pse1) 1]
ob-initto
energy
NequIP Maas 33
NequlP (1xx=3))
0.6
0.6
0.6
0-
ed-initio
0
0
90
135 dibedral angle ())
180
90
133
180
90
135 dihudral angle (*)
180
dibedral angle (*)
rol. onergy (eV)
1.2
SoSkrates (Lame-) 1)
1.2
SoSkrates (Lac 1)
1.2
Suxkerater (Laee 1)
with corrections
with corrections no corrections
with corrections no corrections
0.6
no corrections
0.6
0.6
ab-initio
ob-iniiio
on-iniiio
135 dihedral angle @ (*)
180
135 dihedral anple @ (*)
180
dibedral angle e (*)
(b) Results for different models applied largest (right)
curmulene structure
1.2
SchNet (Ten = GÁ)
1.2
GGMMI (glohal)
1.2
SpookyNet (- non-local block) ab-initio
rel. energy
SchNet. (Tan 12A)
ab-instio
0.6
ob-snitio
0.6
0.6
135 dihedral angle @ (*)
180
135 dihedral angle @ (")
180
135 dihedral angle @ (*)
180
Figure 2: (a) Energy predictions for different cumulene structures with NEQUIP (maa 1
and max = 3) and SO3KRATES (lmax = 1) models with and without geometric corrections
(eqs. and (4)). Cutoff radius and number of layers are kept constant at reut = 2.5 À and n. 4, respectively. (b) We additionally test other SOTA models on the largest cumulene structure (9 carbons). We find, that even for a very large cutoff, SCHNET fails to describe the dihedral angle profile. The same holds true for an inherently global sGDML model and even for a SPOOKY model that explicitly includes non-local corrections.
4 Experiments
In the first subsection, we show how non-local quantum effects can be incorporated by using non- local corrections from the space of SPHCs. In the second part of the experiments, we remove the non-local part which yields traditional, Ro-local MPNN which reaches SOTA results on established benchmarks while requiring much less computational time and parameters than competitive models. A scaling analysis as well as an accuracy comparison for both model variants can be found in appendix A.8 and A.5
Non-Local Geometric Interactions For efficiency reasons, MPNNs only consider interactions between atoms in local neighborhoods, i.e. within a cutoff radius Teut- Thus, information can only be propagated over a distance of Teut within a single MP step. Although multiple MP updates increase the effective cutoff distance, because information can "hop" between different neighborhoods as long as they share at least one atom, each MP step is accompanied by an undesirable loss of information, which limits the accuracy that be obtained. Consequently, MPNNs local effects on length-scalle scales that exceed the effective cutoff distance. To illustrate this problem, we consider the challenging open task [2] of learning the potential energy of cumulene molecules with different sizes (see Fig. 2 Here, the relative orientation of the hydrogen rotors at the far ends of the molecule strongly influences its energy due non-local electronic effects [21]. In order to be able to successfully learn the energy profile with a local model, the effective cutoff has to be large enough to allow information propagate from hydrogen rotor to the other.
unable to describe non-
As a representative example for MPNNs, we consider the recently proposed NEQUIP model 32 which achieves SOTA performance several benchmarks. We find that even when the effective cutoff radius is large enough in principle, an MPNN with ni 4, Tent. = 2.5 A, and Imax < 1 fails to learn the correct energy profile. This is due to the fact that the relevant geometric information "cancels out" (similar to addition of vectors oriented in opposite directions) within each neighborhood,
7

---

(8)
(b) Low-dimensional representation of atomic positions in X-space
Before training
After training
Layer
Layer
Layer
Layer
Layer
Layer
Layer
Layer
136"
Far
185"
close
Gor mporr
6 =
9
closs ?
e 181"
== 180
Figure 3: (a) Cumulene structure in Euclidean space. (b) cow-dimensional projections of the SPHCs (muax 1) SO3KRATES before and after training. Carbon atoms are grey and hydrogen atoms on different ends of the molecule are green and red, respectively. During training, the model maps hydrogen atoms opposite ends close together SPHC space (which not the case initialization).
underlining the limited expressiveness of mean-field interactions in MPNNs. Only by including higher-order geometric correlations, e.g. going to 1max 3, the correct energy profile can be recovered (at the cost of computational efficiency). When going to even larger cumulene structures. however, the effective cutoff becomes too small and it necessary to increase the number MP layers solve the task (again, at the cost of lower computational efficiency), which is illustrated appendix AA.1 Neither increasing the maximum degree of interactions maxx nor the number of layers ng. is satisfactory workaround: Instead of offering general solution to describe non non-local interactions, both options decrease computational efficiency, while only shifting the problem to larger length-scales higher-order geometric correlations.
We further apply three additional models to the cumulene structure with nine carbon atoms. To that end, we use an invariant SCINET model with varying cutoff distances Á and 12 A), an inherently global but invariant SGDML model and the SPOOKYNET architecture which explicitly includes global effects using a non-local block. We find that none of the three is capable of describing the rotor energy profile of cumulene.
In contrast, our proposed S03KRATES architecture able reproduce the energy profile for cumulene molecules of all sizes independent of the effective cutoff radius. Crucially, even with Imax = 1, the predicted energy matches the ah-initio reference faithfully. We find that geometric corrections in the MP update of the SPHCs (cf. eq. (14) are responsible for the increased capability of describing higher-order geometric correlations,
SO3KRATES model with a naive MP update (cf. eq. (13) fails to solve this task with max = 1 (see Fig. 2.. We further confirm that the model picks up on the physically relevant interaction between the hydrogen rotors by analysing the attention values after training (see Fig. appendix A.7.. To illustrate how SO3KRATES is able to describe non-local effects, we show a low-dimensional projection of the atomic SPHCs before and after training for the largest of the cumulene molecules (Fig- 5 After training, the SPHCs for hydrogen atoms at opposite ends of the molecule are embedded close together SPHC space, allowing SO3KRATES to efficiently model the non-local geometric dependence between the hydrogen rotors.
Generalization to structures, larger than those in the training data are usually associated with the re-usability of the learned, local representations. For that reason, it is unclear if this property still holds when non-local corrections are used. As we show in appendix A.5 a SO3KRATES model with non-local corrections still generalizes well to completely unknown and larger structures.
Benchmarks, Data Efficiency and Generalization As pointed out in 3 and equivariant features not only increase performance, but also improve data efficiency. The latter is particularly important, as ab-initio methods for reference data generation can become exceedingly expensive when high accuracy is required. Here, we use a subset of the recently introduced QM7-X data set 5II, which we call QM7-X250. It contains 250 different molecular structures, each with data points for training. 10 data points for validation and 11 11-3748 data points for testing (for details, see appendix A.9). The small number training/validation samples per molecule makes it particularly suited for evaluating model behavior in the low data regime. In the following, train (1) one model

---

(a)
Accuracy on QM7-X250
(b) Accuracy for Unknown Structures
125
220
av. farce mos
av. foro man
[meV A-
(meN
0
fax
2
3
Sodkarates (individual)
Sodkrates (joint)
GDDML
Figure 4: (a) Average force MAE as a function of (max on 250 structures from QM7-X [11] The blue solid line corresponds to Sо3KRATES models trained on each of the 250 structures (à 80 training points), individually. The jointly trained SO3KRATES model (red solid line) is a single model trained on all 250 x 80 training points. (b) Accuracy of force predictions for unknown structures as function of Lmax- Generalization is only investigated for the jointly trained model. Due to the way the molecular descriptor is designed, SGDML 38 models can only be trained on individual structures.
per structure in QM7-X250 and (2) one model for structures in QM7-X250 (250 80 20k training points), which we refer to as individual and joint models, respectively.
We start by investigating the performance eгror strongly decreases with higher max (Fig. 4.. As kernel methods are known to perform well in the low-data regime [2I] we compare our results SGDML [38] kernel models, which only use distances as a molecular descriptor (corresponding to mmax 0). For lmax O, we find SGDML gives competitive results, whereas for Imax 1, Sо3KRATES starts to outperform SGDML. As soon as mmax 2 2, however, the prediction accuracy of SоžKRATES is greater than that of SGDML by large margin. Thus, increasing the order of geometric information in the SPHCs leads to strong improvements in the low-data regime. For jointly trained models, we find that SO3KRATES outperforms SGDMI. even for Imax = O, with continuous improvement for increasing 4max- In appendix A. we report energy and force errors across degrees and further experimental details.
function of the maximal degree lmsx and find that the
The generalization capability of SO3KRATES is tested, by applying a jointly trained model to 25 completely unknown molecules from the QM7-X data set (see, Fig. Ab. details in appendix A. Again, we find that force MAEs decrease with increasing max. For reference, we compare performs on par,
S03KRATES to individually trained SGDML models and find that S03KR or even slightly better. for /max 2 Going beyond Imax is found to only marginally improve generalization. In addition, we report results for a model trained on the full QM7-X data set in appendix A.5 following 330.
For completeness, we also apply SO3KRATES to the popular MD17 benchmark (see Table D. We find, that SO3KRATES outperforms networks that have the same parameter complexity by large margin (PAINN and NEWTONNET). Notably, it requires significantly less parameters than other SH based architectures (NEQUIP and SPOOKYNET), while performing only slightly worse or even on par with them. Furthermore, Sо3KRATES outperforms DIMENET, its closest competitor in timing (cf. Fig. d), consistently by a large margin. Compared to current SH based approaches, GEMNETQ needs less parameters (still ~ 2.5x more than SO3KRATES) to achieve competitive results. However, requires the explicit calculation of dihedral angles which scales cubically in the number of neighboring atoms. Due to its linear scaling (see A.8) and lightweight structure, SO3KRATES can significantly reduce the time for training and inference (see Fig. and A.6).
5 Discussion and Conclusion
Due to the locality assumption used in most MPNNs, they are unable to model non-local electronic effects, which result in global geometric dependencies between different parts of a molecule. The length-scales of such interactions often greatly exceed the cutoff radius used in the MP step, and even though stacking multiple MP layers increases the effective cutoff, ultimately, MPNNs are not capable of efficiently modeling geometric dependencies on arbitrary length scales.
In this work, we contribute conceptually by proposing an efficient and scalable solution to this problem. We suggest a set of refinable, equivariant coordinates for point clouds in Fuclidean space, called spherical harmonic coordinates (SPHCs). Non-local geometric effects can then be efficiently

---

Table 1: MAE for energy (in keal mol-1) and forces (in kcal - of current state-of-the-ar machine learning models on the MD17 benchmark for 1 training points. For each model, the numher of paramacters, well as the scaling in the number of neighbors m shown.
NEQUIP 1S4 3M (mc)
SPOOKYNET 3M (M)))
30
GUMNErO E4
DIMENET 19M1 (O(ee))
PAINN M
NEWTONNET SS 200k ((n))
SoŠKRATES 700k ((n))
2:2M (M(w))
1 (1)
Aspirin
cnergy forces
0.13
0.151
0.204
0.159
0.168
0.139
0.19
0.258
0.217
0.499
0.371
0.545
0.236
Echanol
0.088
0,064
0,230
Malondialdehyde
onergy fores
0.08
0.079
0.104
0.091
0.096
0.077
0.15
0.167
0.139
0.383
0.319
0.522
0.147
Naphthalene
energy forces
0.117
0.04
0.089
0.051
0.215
0.083
0.084
0.074
Salicyclic Acid
0.114
0.134
0.114
0.115
0.106
0.374
Toluene
energy
0.09
0.091
0,102
0.097
0.094
0.095
forees
0.05
0.087
0.060
0.216
0.102
0.00s
0,073
energy
forces
0.08
0.119
0.104
0.301
0.140
0.149
0.111
modeled by including geometric corrections, which are localized in the space of SPHCs, but non-local in Euclidean space. Further, we show that introducing spherical filter functions acting on the SPHCs increases geometric resolution and predictive accuracy.
We then propose the So3KRATES architecture, self-attention based MPNN, which decouples atomic features from higher-order geometric information. This allows to drastically decrease the parametric complexity while still achieving SOTA prediction accuracy. We show evidence that increasing the geometric order of SPHCs greatly improves model performance in the low-data regime, as well as generalization to unknown molecules.
limitation of the current implementation of So3KRATES that spherical neighborhoods NNx in eq. [4] are computed from all pairwise distances in SPHC space. An alternative implementation could use a space partitioning scheme to find neighborhoods more efficiently. In a broader context, our work falls into the category of approaches that help to reduce the vast computational complexity of molecular and material simulations. This can accelerate novel drug and material designs, which holds the promise of tackling societal challenges, such as climate change and sustainable energy supply [53] Of course, our method could also be used for nefarious applications, e.g. design of chemical warfare, but this is true for all quantum chemistry methods.
Future research will focus on applications of SоžKRATES to materials and bio-molecules, which are typical examples of chemical systems where the accurate description of non-local effects is necessary to produce novel insights. Efficient treatment non-local effects in point cloud data goes beyond the domain of quantum chemistry. One way of representing non-local dependencies are non-local neural networks [54]. In comparison to the presented approach they compute a relation in feature rather than in Euclidean space, making it incapable of capturing direct geometric relations in Euclidean space. However, this might be necessary if the relative orientation of objects far apart from each other plays role for identifying different objects.
6 Acknowledgements
All authors acknowledge support by the Federal Ministry of Education and Research (BMBF) for BIFOLD (01IS18037A) KRM was partly supported by the Institute of Information Communi- cations Technology Planning & Evaluation (ITP) grants funded by the Korea govermment(MSIT)
(No. 2019-0-00079, Artificial Intelligence Graduate School Program, Korea University and No. 2022-0-00984, Development of Artificial Intelligence Technology for Personalized Plug-and- Explanation and Verification of Explanation). and was partly supported by the German Ministry for Education and Research (BMBF) under Grants 011514013A-E, AIMM, 01GQ1115, 15, O1GQ0850, 011$18025A and 011S 8037A: A; the German Research Foundation (DFG). We thank Stefan Chmiela, Mihail Bogojeski and Nicklas Schmitz for helpful discussions and feedback the manuscript.
10

---

References
[1] Mark E Tuckerman. Ab initio molecular dynamics: basic concepts, current trends and novel applications. J. Phys. Condens. Matter, 14(50):R1297, 2002.
[2] Frank Noé, Alexandre Tkatchenko, Klaus-Robert Müller, and Cecilia Clementi. Machine learning for molecular simulation. Annu. Rev. Phys. Chem., 71:361-390, 2020.
[3] John Keith, Valentin Vassilev-Galindo, Bingqing Cheng, Stefan Chmiela, Michael Gastegger, Klaus-Robert Müller, and Alexandre Tkatchenko. Combining machine learning and com- putational chemistry for predictive insights into chemical systems. Chemical Reviews, 121 (16):9816-9872 2021. URL httpp::/pubs.... acs. cc00107
11.1111111....... chemrev
[4] Jörg Behler and Michele Parrinello. Generalized neural-network representation of high- dimensional potential-energy surfaces. Phys. Rev. Lett., 98(14):146401 2007.
[5] 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. Phys. Rev. Lett., 104(13): 136403, 2010.
[6] Jörg Behler. Atom-centered symmetry functions for constructing high-dimensional neural network potentials. J. Chem. Phys., 134(7):074106, 2011.
[7] Albert P Bartók, Risi Kondor, and Gábor Csányi. On representing chemical environments. Phys. Rev. B, 87(18))184115, 2013.
[8] Zhenwei Li, James R Kermode, and Alessandro De Vita. Molecular dynamics with on-the-fly machine learning of quantum-meehanical forces. Phys. Rev. Lett., 114(9):094405, 2015.
[9] Stefan Chmicla, 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. Sci. Adv., 3(5): 1603015, 2017.
[10] Kristof T Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus-Robert Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nat. Commun., 8: 13890, 2017.
[11] Michael Gastegger, Jörg Behler, and Philipp Marquetand. Machine learning molecular dynamics for the simulation of infrared spectra. Chem. Sci., 8(10))6224-6935, 2017.
[12] Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-learned force fields. Nat. Commun., 9(1): 3887, 2018. doi: 44733///44447008-06666669999
[13] Kristof T Schütt, Huziel E Sauceda, Pieter-Jan Kindermans, Alexandre Tkatchenko, and Klaus- Robert Müller. SchNet a deep learning architecture for molecules and materials. Chem. Phys., 148(24):241722, 2018.
[14] Justin Gilmer, Samuel S Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning. pages 1263-1272. Pmlr, 2017.
[15] Justin S Smith, Olexandr Isayev, and Adrian E Roitherg. ANI-1: an extensible neural network potential with DFT accuracy force field computational cost. Chem. Sci., 8(4):3192-3203, 2017.
[16] Nicholas Lubbers, Justin S Smith. and Kipton Barros. Hierarchical modeling of molecular energies using a deep neural network. J. Chem. Phys., 148(24):241155, 2018.
[17] Martin Stöhr, Leonardo Medrano Sandonas, and Alexandre Tkatchenko. Accurate many-body repulsive potentials for density-functione tight-binding from deep tensor neural networks. Phys. Chem. Lett., 11:6835-6844 2020.

---

[18] Felix A Faber, Anders S Christensen, Bing Huang, and O Anatole von Lilienfeld. Alchemical and structural distribution based representation for universal quantum machine learning. J. Chem. Phys., 148(24):241117,22 2018.
[19] Oliver T Unke and Markus Meuwly, PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges. J. Chem. Theory Comput., 15(6):3678-3693, 2019.
[20] Anders S Christensen, Lars A Bratholm, Felix A Faber, and O Anatole von Lilienfeld. FCHL revisited: faster and more accurate quantum machine learning. J. Chem. Phys., 152(4):044107 2020.
[21] Oliver T Unke, Stefan Chmiela, Huziel E Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schüit,, Alexandre Tkatchenko, and Klaus-Robert Müller. Machine learning force fields. Chemi- cal Reviews, (21(16);
2021.
[22] Yaolong Zhang, Ce Hu, and Bin Jiang. Embedded Atom Neural Network Potentials: efficient and accurate machine learning with a physically inspired representation. J. Phys. Chem. Lett., (11):4962-49677 2019.
[23] Silvan Käser, Oliver Unke, and Markus Meuwly. Reactive dynamics and spectroscopy of hydrogen transfer from neural network-based reactive potential energy surfaces. New J. Phys., 22:55002, 2020.
[24] Anatole von Lilienfeld, Klaus-Robert Müller, and Alexandre Tkatchenko. Exploring chemical compound space with quantum- bbsed machine learning. Nat. Rev. Chem., 4(7):347-358, 2020.
[25] Victor Gareia Satorras, Emiel Hoogeboom, and Max Welling. E (n) equivariant graph neural networks. In International Conference Machine Learning, pages 9323-9332. PMLR, 2021.
[26] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13472 2021.
[27] Yann LeCun, Yoshua Bengio, et al. Convolutional networks for images, speech, and time series. The handbook of brain theory and neural networks, 3361(10):1955, 1995.
[28] Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. arXiv preprint arXiv: /60.:02900000 2016.
[29] Johannes Klicpera, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs. arXiv preprint arXiv:2003.03123 2020.
[30] Oliver Unke, Stefan Chmiela, Michael Gastegger, Kristof T Schütt,, Huziel Sauceda, and Klaus-Robert Müller. SpookyNet: Learning force fields with electronic degrees of freedom and nonlocal effects. Nat. Commun., 12:7273, 2021.
[31] Kristof ' Schütt, Oliver T Unke, and Michael Gastegger. Equivariant message passing for the prediction tensorial properties and molecular spectra. arXiv preprint 1xXi::2102.03150, 2021.
[32] Simon Batzner, Tess Smidt, Lixin Sun, Jonathan Mailoa, Mordechai Kornbluth, Nicola Molinari, and Boris Kozinsky. SE (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. arXiv preprint 1rXiv:2101.03/64, 2021.
[33] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: rotation-an and cunslation-equivaaan neural networks for 3D point clouds. arXiv preprint arXiv:1802.082199 2018
[34] Johannes Klicpera, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. arXiv preprint 10X.:22006008033.. 2021.
[35] Sergey N Pozdnyakov, Michael J Willatt, Albert P Bartók, Christoph Ortner, Gábor Csányi, and Michele Ceriotti. Incompleteness of atomic structure representations. Physical Review Letters, 125(16):166001,
2020. 2020.
12

---

[36] Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling. Se (3) transsforrsss 3d ooto-rranslation equivariant attention networks. arXiv preprint 20Xv:2000611503,,, 2020.
[37] Aneja Alisha. Performance compariso between nvidia's geforce gtx 1080 and tesla p100 for 2017.
deep learning. https:://itthub.com./
[38] Stefan Chmiela, Huziel Sauceda, Igor Poltavsky, Klaus-Robert Müller, and Alexandre Tkatchenko. sgdml: Constructing accurate and data efficient molecular force fields using machine learning. Computer Physics Communications, 240:38 400:38455,22 2019.
[39] Brandon Anderson, Truong-Son Hy, and Risi Kondor. Cormorant: covariant molecular neural networks. arXiv preprint arXiv: 19066990001111 2019.
[40] Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations, 2021.
[41] Matan Atzmon, Haggai Maron, and Yaron Lipman. Point convolutional neural networks by extension operators. arXiv preprint 1rXivv1803.. 1003, 2018.
[42] J-F Cardoso and Beate H Laheld. Equivariant adaptive source separation. IEEE Transactions on signal processing, 44(12):001733030 1996.
[43] Antonios K Alexandridis and Achilleas D Zapranis. Wavelet neural networks: with applications in financial engineering, chaos, and classification. John Wiley & Sons, 2014.
[44] Taco S Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pages 2990-2999. PMLR, 2016.
[45] Geoffrey E Hinton, Alex Krizhevsky, and Sida D Wang. Transforming auto-encoders. In International conference on artificial neural networks, pages 44-51. Springer, 2011.
[46] Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International Conference on Machine Learning, pages 5361-5370. PMLR, 2020.
[47] Eugene Paul Wigner. Gruppentheorie und ihre anwendung auf die quantenmechanik der atomspektren. 1931.
[48] Oliver Unke, Mihail Bogojeski. Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus- Robert Müller. -quivvaana prediction of molecular wavefunctions and electronic densities. Advances in Neural Information Processing Systems, 34, 2021.
[49] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. arXiv preprint arXiv: 1706.03666207 2017.
[50] Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. arXiv preprint arXiv:1710.1090: 2017.
[51] Johannes Hoja, Leonardo Medrano Sandonas, Brian Ernst, Alvaro Brrquez-Mayygoitte Robert A DiStasio Jr, and Alexandre Tkatchenko. Qm7-x, comprehensive dataset of quantum- mechanical properties spanning the chemical space of small organic molecules. Scientific data, 8(1):1-11, 2021.
[52] Mojtaba Haghighatlari, Jie Li, Xingyi Guan, Oufan Zhang, Akshaya Das, Christopher J Stein, Farnaz
Meili Liu, Martin Head-Gordon Luke Bertels, et al. Newtonnet: A newtonian message passing network for deep learning of interatomic potentials and forces. arXiv preprint arXiv:2108.02913 2021.
[53] Anubhav Jain, Shyue Ping Ong, Geoffroy Hautier, Wei Chen, William Davidson Richards, Stephen Dacek, Shreyas Cholia, Dan Gunter, David Skinner, Gerbrand Ceder, al. Commentary: The materials project: materials genome approach accelerating materials innovation. API. materials, 1(1):01 1002, 2013.
13

---

[54] Xiaolong Wang, Ross Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In Proceedings of the IEEE conference on computer vision and pattern recognition. pages 7794-7803, 2018
[55] Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: neural network library and ecosystem for JAX, 2020. URL
http://github. wtt.:a...coooeee www.coo/oooooo coon/google/ffa.
[56] Matteo Hessel, David Budden, Fabio Viola, Mihaela Rosca, Eren Sezener, and Tom Hennigan. Optax: composable gradient transformation and optimisation, in jax! 2020, URL http:
[57] Charles R. Harris, K. Jarrod Millman, Stéfan van der Walt, Ralf Gommers, Pauli Vir- tanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten van Kerkwijk, Matthew Brett, Allan Hal- dane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre
Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis Oliphant. Array programming with NumPy. Nature, 585(7825) September 2020. doi: 11038/44188--00000 2649-2. URL. TRI.https::: httph: :ttipp::/doi..o http::://oooooor//////o 10. 00..0/////1186- 020 000-2-699--22--2-222--- 2649-2
[58] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman- Milne,, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://githie com/ google/jax
[59] Ameen Ali, Thomas Schnake, Oliver Eberle, Grégoire Montavon, Klaus-Robert Müller, and Lior Wolf. Xai for transformers: better explanations through conservative propagation. arXiv preprint http:2202.07344 2022.
[60] Tsz Wai Ko, Jonas A Finkler, Stefan Goedecker, and Jörg Behler. A fourth- suurth--eeeration high- dimensional neural network potential with accurate electrostatics including non-local charge transfer. Nature communications, 12(1):1-11, 2021.
[61] Diederik P Kingma and Jimmy Ba. Adam: a method for stochastic optimization. arXiv preprint 1xXiv:1412.6988 2014.
14


Reviews:
[{"review_id": "y-fq56zhb0A", "reviewer": "Reviewer_toeX", "summary": "The authors consider the problem of capturing non-local effects in molecular modelling. To address this, they construct a model based on:\n\n- Spherical Harmonic Coordinates (SPHC): A wavelet transform with different orders of spherical harmonic is used to represent the molecule geometry in a SPHC space. An equivariance relation between rotations in the euclidean input and Wigner-D transformations in SPHC gives the model the right symmetries, and the intention is that non-local interactions in Euclidean space may become local in SPHC.\n- An adapted message passing scheme which passes messages in local neighbourhoods both in Euclidean and SPHC space while preserving equivariance.\n\nThe authors provide demonstration of their method on an educational example of cumulene twisiting as well as larger benchmarking on QM7-X250 and MD17.", "questions": "The proposed SPHC solution is well suited to the specific angle-dependent challenge of cumulene, however there are no other examples of non-local interaction such as dipole-dipole or charge-dipole interaction highlighted in the paper. Is there another case study where SPHC is the correct representation to capture the non-local effects and a 6A cutoff is so clearly deficient? How does So3krates do at modeling water clusters or the Au2-MgO system considered in the SpookyNet paper or other examples originating from Ko et al?", "limitations": "Locality is often built-in to neural-network based force fields to both (1) limit computational cost and (2) enable generalisation from smaller sub-systems to larger simulations. Do we lose these when we use a non-local model? To what extent can So3krates be scaled to something like the 100 million atom system recently run on 128 GPUs using Allegro [1]. Perhaps such an experiment is hard to run, but some evidence that the speed benefit of So3krates on MD17 translates to a large scale simulation would help the case for non-local model development.\n\n[1] https://arxiv.org/pdf/2204.05249.pdf", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall the authors have come up with a principled approach for including long ranged geometric features into a message passing architecture. I think the method is interesting and sufficiently novel for publication, and the results that are presented are favourable but perhaps slightly incomplete (see significance section below).\n\nOriginality\n\nThe idea to perform some sort of convolution with spherical harmonic wavelets to provide additional geometric input is not new, and nor is the discussion of the equivariance of this operation. However, the interpretation of these SPHC features as living in a coordinate space where interesting neighbourhoods can be learned for message passing that correspond to very distant communication in Euclidean space is an interesting new insight that I had not thought about before.\n\nQuality\n\nThe main takeaway for me was that the proposed model is (1) solving the specific case of cumulene highlighted in literature before (2) achieving similar performance on MD17 as a number of other methods, but with a small number of parameters and higher train/test speed. I have comments on (1) in the Questions section below and (2) in the significance section below. However I think overall the quality of results meets the publication bar because the lightweight nature of the model seems a promising direction.\n\nClarity\n\nThe paper is easy to follow.\n\nSignificant\n\nI feel like the significance of the paper suffers slightly from a couple of missing baselines:\n\n(1) I’m not an expert on ML forcefields, but I think SpookyNet makes more of a deliberate effort to model non-local behaviour than NequIP, so I was not sure why NequIP was chosen as the baseline in fig 2. A more considered description of how models like SpookyNet, PhysNet, sGDML etc capture non-locality and what sort of non-locality SPHC is particularly adept (or not!) at capturing might be more educational than the somewhat obvious (and already published) statement that a NequIP model cannot capture effects larger than it can see.\n\n(2) I’m not sure I know how to evaluate the differences between the numbers on MD17: Although the statement that the errors are low and the parameter counts are low are true, all the models in table 1 have ~10^6 parameters in order of magnitude and all have broadly similar performance on MD17. Do these differences matter in practice? The timings of So3krates in fig 1d seem the most convincing difference to existing methods, but some of the baseline models are missing from this plot and DimeNet (the closest competitor in timing) is missing from the rest of the paper.\n\n(3) SpookyNet has ~10-20 meV/A on the full QM7-X set, but this is not mentioned so I was not sure why So3krates was compared with sGDML rather than something like SpookyNet. Also I think I missed why So3krates is only tested on a subset of QM7-X.\n\nI also flag a slight concern that this paper mainly addresses the cumulene problem which is not of broad interest to NeurIPS, but I personally found this interesting so I did not factor this into my score.\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "nqDy1qCbv-r", "reviewer": "Reviewer_jF7P", "summary": "The paper proposes to use separate representations of spherical harmonic coordinates \njointly with representations of atomic features in learning models of molecular force fields.\nThis allows message passing to local neighborhoods in the spherical harmonic coordinates\nwhich translate to non-local neighborhoods in Euclidean space. As a result the method\nis able to model long-range interaction effects that are prohibitively expensive to model with\nexisting, hop-limited message-passing graph neural networks.  It is also helpful to the\nlow-data cases since the fewer number of parameters in the network requires less training\ndata to estimate.", "questions": "The idea of using harmonic functions like wavelets in deep learning has been\npursued before [for example, see \"Wavelet Neural Networks: \nWith Applications in Financial Engineering, Chaos, and Classification\", \nAntonios K. Alexandridis, Achilleas D. Zapranis, John Wiley & Sons, 2014].\nHow does the current proposal compare to those uses?\n", "limitations": "The authors discuss the limits of the current applications to materials\ndiscovery, while projecting a wider scope of application to other spatial modeling needs\nlike in computer vision.  However, in those wider areas, there have been more active\nconsiderations of other representations to capture spatio-temporal dependences.\nSome review of where the proposal approach could distinguish itself can add to the\nargument.\n", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The proposal is sound and convincing, and echoes the long established methods\nin signal processing.  Experimental demonstration and comparison to prior art are convincing.\n", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "QE8bFZBTnFW", "reviewer": "Reviewer_SNvt", "summary": "The authors propose a new ML force field based on higher order interactions and equivariant self-attention. The model represents atomic features using spherical harmonic coordinates that are updated during the forward pass through the network, and spherical neighborhoods are constructed using these representations. This helps the model learn relations between distant atoms. As a result, the proposed method is computationally efficient, requires few parameters and achieves comparable performance to recent methods like Nequip.", "questions": "1. How does the performance vary with the number of nodes? Spherical neighborhoods computation and self-attention could be very expensive for large systems like bio-molecules and materials.\n2. MD17 is a small, relatively homogenous dataset. Do the authors have a sense for how the model would perform on a large and diverse dataset like OC20 (opencatalystproject.org)?\n", "limitations": "Seems adequate.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n* SPHC representation helps the model capture global dependencies in a natural way.\n* The entire model is equivariant, which is a desirable property to have for its data efficiency.\n* The proposed method requires very few parameters to achieve good performance.\n* It is also very fast for inference, which is an important property to have for many applications like MD simulations where millions of inference steps are needed.\n\nWeaknesses\n* The paper does not present results from multiple, diverse datasets (unlike many other papers in this area). So, it is unclear how generalizable the results are.\n* The results in table 1 show that Nequip obtains better results (though at higher training & inference cost). Therefore, this method is only suitable when computational efficiency is more important than accuracy.", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept