Spatial Attention Kinetic Networks with E(n)-Equivariance
Authors: Yuanqing Wang
John Chodera
Venue: ICLR 2023 poster (2023)

Abstract: Neural networks that are equivariant to rotations, translations, reflections, and permutations on $n$-dimensional geometric space have shown promise in physical modeling for tasks such as accurately but inexpensively modeling complex potential energy surfaces to guiding the sampling of complex dynamical systems or forecasting their time evolution.
Current state-of-the-art methods employ spherical harmonics to encode higher-order interactions among particles, which are computationally expensive.
In this paper, we propose a simple alternative functional form that uses neurally parametrized linear combinations of edge vectors to achieve equivariance while still universally approximating node environments.
Incorporating this insight, we design \emph{spatial attention kinetic networks} with E(n)-equivariance, or SAKE, which are competitive in many-body system modeling tasks while being significantly faster.

Published as a conference paper at ICLR 2023
SPATIAL ATTENTION KINETIC NETWORK WITH E(N) EQUIVARIANCE
Yuanqing Wang'and and John D. Chodera Computational and Systems Biology Program,
Sloan Kettering Institute
Memorial Sloan Kettering Cancer Center, New York, N.Y. 10065
yuanqing wuanging..wwngecnne.... www.@cooeeaaaaabbbbboooo
ABSTRACT
Neural networks that are equivariant to rotations, translations, reflections, and permutations ne-dimensional geometric space have shown promise in physical modeling-froot modeling potential energy surfaces to forecasting the time evo- lution of dynamical systems. Current state-of-the-an art methods employ spherical harmonics to encode higher-order interactions among particles, which are compu- tationally expensive. In this paper, we propose simple alternative functional form that uses neurally parametrized linear combinations of edge vectors to achieve equivariance while still universally approximating node environments. Incorporat- ing this insight, we design spatial attention kinetic networks with -nnuivavvre,, SAKE, which are competitive in many-body system modeling tasks while being significantly faster.
INTRODUCTION
Encoding the relevant symmetries of systems of interest into the inductive biases of deep learning architectures has been shown to be crucial in physical modeling. Graph neural networks (GNNs) (Kipf
and Welling 2016 Xu et al. 2018 Gilmer et al., 2017 Hamilton et al. 2017]] Battaglia et al. 2018).
for instance, preserve permutation equivariance by applying indexing- invariant pooling functions among nodes (particles) and edges (pair-wise interactions) and have emerged to become powerful workhorse in a wide range of modeling tasks for many-body system (Satorras et al., 2021).
When describing not only the topology of the system but also the geometry of the state, relevant symmetry groups for three-dimensional systems are SO(3) (rotational equivariance), SE(3) (rotational and translational equivariance), and E(3) (additionally with reflectional equivariance). A ubiquitous and naturally invariant first attempt to encode the geometry of such systems is to employ only radial information, i.e., interparticle distances. This alone has empirically shown utility in predicting quantum chemical potential energies, and considerable effort has been made in the fine-tuning of radial filters to achieve quantum chemical accuracy- kcal/mol, the empirical threshold to qualitatively reliably predict the behavior of quantum mechanical system-ant beyond (Schütt al. 2017).
Nonetheless, radial information alone is not sufficient to fully describe node environments the spatial distribution of neighbors around individual particles. The relative locations of particles around central node could drastically change despite maintaining distances to these neighbors unaltered. To describe node environments with completeness, one needs to address these remaining degrees of freedom. Current sfute-of-the-a approaches encode angular distributions by employing a truncated series of spherical harmonics to generate higher-order feature representations; while these models have been shown to be data efficient for learning properties of physical systems, these features are expensive to compute, with the expense growing rapidly with the order harmonics included (Thomas et al., 2018 Klicpera et al. 2021aa Fuchs et al., 2020 Batzner al. [021]] Anderson et al.) 2019). The prohibitive cost would prevent this otherwise performant class of models
AAlternative address: Ph.D. Program in Physiology, Biophysics, and System Biology, Weill Comell Medical College, Comell Univerisity, New York, N.Y. 10065

---

Published as a conference paper at ICLR 2023
from being employed in materials and drug design, where rapid simulations of large systems crucial to provide quantitative insights.
Here, design simple functional form, which call spatial attention, that uses the norm of set of neurally parametrized linear combinations of edge vectors to describe the node environment. Though simple in form, easy to engineer, and ultra-fast to compute, spatial attention is capable of universally approximating any functions defined on local node environment while preserving hun--invafaanceeepiivaaannne in arbitrary n-dimensional space.
After demonstrating the approximation universality and invariance of spatial attention, we incorporate it into novel neural network architecture that uses spatial attention parametrize fictitious velocity and positions equivariantly, which we call a spatial attention kinetic network with E(n)- (inn)--qvvaiineee or SAKE (pronounced saké (sah-keh), like the Japanese rice wine | To demonstrate the robustness and versatility of SAKE, we benchmark its performance on potential energy approximation and dynamical system forecasting and sampling tasks. For all popular benchmarks, compared to state- of-the-art models, SAKE achieves competitive performance on wide range of invariant (MD17: Table I QM9: Table 3 ISO17: Table 2 and equivariant (N-body charged particle Table 4 walking motion: Table 6) while requiring only fraction of their training and inference time.
2 BACKGROUND
In this section, we provide some theoretical background on physical modeling, equivariance, and graph neural networks to lay the groundwork for the exposition of spatial attention networks.
EQUIVARIANCE: PERMUTATIONAL, ROTATIONAL, TRANSLATIONAL, AND REFLECTIONAL
A function f : + y' is said to be equivariant to symmetry group G if
f(Tx(x)) Sy(f(x)
(1)
Vg E G and some equivalent transformations the two spaces respectively Tg X + X' and S, y +.
If on n--iimensional space X' = =0= =R, we have permutation P and (x(x) Px, Sy = Py satisfying Equation ! we say f is permutationally equivariant; Ty (x) = xR where Raxn a rotation matrix RRT = 1, and Sy))) = y we say f is rotationally equivariant: if T, = + Ax and (y) y + Ax, where _ we say is translationally equivariant; finally, if Ty (x) = Ref (x) and Sy) = Ry) and Refş is a reflection on n-dimensional space, we say f reflectionally equivariant.
2.2 GRAPH NEURAL NETWORKS
Modern GNNs, which exchanges and summarizes information among nodes and edges, are better analyzed through the spatial rather than spectral lens, according Wu et al. (2019) classification. Following the framework from Gilmer et al. 2001);;2. Xu et al. (018)) Battaglia et al. (2018). for node with neighbors u E (v),, in a graph G, with ( denoting the feature of node the k-th layer (or k-th round of message- message-passingg and A E RRC the initial node feature on the embedding space, the k-th cessage-passing step of GNN can be written three steps:
First, an edge update,
(k+ = o (h ( (M)))
(2)
where the feature embeddings h,, of two connected nodes u and v update their edge feature embedding Te followed by neighborhood aggregation,
(b+11) p
E (e))))
where edges incident to a node pool their embeddings to form aggregated neighbor embedding a.s and finally node update,
h
Implementation:
lab/ sake
2

---

Published as a conference paper at ICLR 2023
where NC- denotes the operation return the multiset of neighbors of node and d' and ø" are implemented as feedforward neural networks. Since the neighborhood aggregation functions pa are always chosen to be indexing- innariannnvaiiant functions, namely SUM or MEAN operator, Equation 3 and thereby the entire scheme, is permutationally invariant.
Problem Statement. We are interested in designing a class of parametrized functions fo X'x H X' x H that map from and to the joint spaces of (1-dimensional) geometry X E IR" and semantic embedding H E RC such that it is permutationally. rotationally, translationally, and reflectionally equivariant on A' and invariant on H. That is. for a given coördinate x E X, embedding h E H. and X' (rotation, translation, and reflection), we
any transformation mentioned in Section 2.TT: : X have:
Xf, h = fo(x,
(o((xx))
(5)
RELATED WORK: INVARIANT AND EQUIVARIANT GNNS
Invariant GNN. Although the notion of graph might not appear in their original publications, plethora of neural architectures could be viewed as de facto graph neural networks with invariance on geometric space (for survey, see (Han al. 2022 SchNet (Schütt al., 2017). for example, approximates the potential energy of molecules by applying only operate on the distances between atoms to ensure the model is -nvvirnvr,,,,,, effectively using a graph neural network architecture where nodes denote atoms and edges denote the distance separation between pairs of atoms within a cutoff radius. addition to interatomic distances, invariant models could also incorporate information about the angular distribution of atom neighbors by computing atomic features that incorporate angular information of eeighhor-atom-neiiee
convolutions that
triplets (Smith et al., 2017 Behler 2011 Klicpera et al. 2020a Waan et al. 2022 2023 Liu al. (022). Compared to SAKE, these models are not capable of equivariant modeling; even in an invariant setting, they are empirically less performant (See invariant tasks performance Section 6.11
Architectures achieving -E(uvvequvvariance by leveraging Bessel functions and spherical harmonics encode higher-order interactions (Thomas) et al. 2018: Klicpera et al.] [2021a] Fuchs et al.] [220)] Anderson et al., 2219]] Klicpera et al. 2021b]] Brandstetter et al. 2021| Liu al. 2022) shows outstanding performance and data efficiency, some of which is on par with SAKE (Table 4.. Villar et al.] (202T) discusses that these higher-order representation is not necessary to construct invariantly and equivariantly universal models. Meanwhile, they tend to be difficult to engineer and expensive to train and run. (See runtime benchmarks in
Equivariant GNN with spherical harmonics.
aforementioned tables.)
Equivariant GNN with dot product scalarization. Recent efforts (Schütt et al., [021 Thölke and Fabritis| 2022 Huang et al., 2022)) relying on dot product of edge tensors as equivariant nonlinearity have achieved competitive results on machine learning potential construction, among which none has been validated to perform well on both equivariant and invariant tasks. Experimentally. they are consistently outperformed by SAKE in invariant (Section 6.1
Table ID and equivariant (Section 6.2
Table () tasks.
Message passing between emebedding and geometric spaces. malized the link between graph neural networks and and equivariance and provided a generalizing framework encompassing iterative coometric-to-mmeeddiee
Satorras et al. (2021) has for-
and sotbedddin--t-geomeette type update, Like our proposed architecture, E(N) Equivariant Graph Neural Networks (EGNN), proposed in Satorras et al.]((22)))))) (2021), also uniquely describes the geometry of n--odyy system, although their argu- ment was based on a global scale given sufficient steps of message passing, while we can sufficiently describe the local geometric environment. This advantage is evidenced by extensive experiments (invariant: Table equivariant: Table and In ablation study (Section 663). we show that without spatial attention, EGNN is not competitive in potential energy modeling, even when all other tricks used in this paper were added. Similar to EGNN (Satorras et al.) 2021), our model is equivariant w.r.t. a general E(n) group and are not restricted to E(3). Also worth noting is that the Geometric Vector Perceptrons algorithm proposed in Jing et al. (2021) could be regarded as special case of our framework where the attention weights are learned globally, whereas we learn them in an amortized manner and thus can be transductively generalized across systems.
3

---

Published as a conference paper at ICLR 2023
4 THEORY: SPATIAL ATTENTION
Given a node v with embedding hy E = RC (where C' denotes the embedding dimension) and position X2 E X' R" (where denotes the geometry dimension) in graph 9, its neighbors u E (v) with connecting edges (cwe} with displacement vector čuv = X2 Xu and embedding hewv p===="" (lw, fx)) with some aggregation function pr-te, we define spatial attention ø X' x
H as
Nx
Ø (v)
+ 1=1
Σ oEA(())
(6)
where X; H - N₂ is a set of arbitrary attention weights-generating function that operates on the edge embedding, f f - X' is an equivariant function that operate on the edge vector, μ N - H. is an arbitrary function that takes the norms of N₂ linear combinations, and t denotes concatenation. We drop the explicit dependence of & on both geometric and embedding properties of and for simplicity. We name this functional form attention because it characterizes the alignment between edge vectors in various projections.
When implemented, À and take the form of feed-forward neural networks. In other words, for each node v, we first parametrize a NN- = ||)H H (is analogous to the number heads in multi-head attention) sets of attention weights based on the edge embeddings X₂ ). At the same time, we also equivariantly (on E(n) group) transform the edge vector into f (@uv) Next, we use those N,, sets of attention weights to linearly combine these vectors among edges and take the Euclidean norm of each of these linear combinations, resulting in N₂ norms. Lastly, we concatenate these norms and put them into feed-forward neural network y to compute the node embedding. It naturally follows that: Remark 1. Spatial attention is permutationallyy rotationally, transalationally, and reflectionally invariant on E(n).
since he.. is invariant w.r.t. indexing and the Euclidean norm function is invariant to rotation, translation, and reflection.
With all local degrees of freedom incorporated in spatial attention, it is perhaps intuitive to see that dSA (v) can uniquely define the local environment of nodes up to E(n) symmetry. We formalize this finding (Proof in Appendix Section I.I
Theorem 1. For a node v in a graph G with neighbors u E (v) connected to u by edges with positions xx, and Xu distinct embeddings heve, + heen,y V1 S i,j S [W(v)), and any E(n)-invariant continuous function g(Xx, Xu, (eew,)) spatial attention @SA can approximate with arbitrarily small errore e with some {}}, J, and μ..
It is worth noting here that we only consider the universal approximation power of mapping from the geometry space X' to scalar space. without considering that on the space of node (and edge) embedding; regardless of geometry. GNNs that operate on neighborhood node embeddings are generally not universal (Xu et al. 2018. Corso et al. 2020). This universal approximative power on F(n)-invariant functions can lead to universally approximative parametrization of F(()) ((m)--vqvvaarant functions, which we show in Remark 2 Practically, the inequality condition is not difficult to satisfy: we implement he as dependent upon edge length so even if the semantic embedding of edges are identical, the inequality hevn, # hean, holds as long the system is not strictly symemtrical at all times (namely the mirror symmetry presented in hydrogen molecule) despite distortions resulted from vibrations.
5 ARCHITECTURE: SPATIAL ATTENTION KINETIC NETWORK (SAKE)
we design a fast, efficient, and casy-to-implement architecture termed spatial attention kinetic network with Ein)-evquvariance, or SAKE, Here, we introduce the remaining components of SAKE before assembling them in modular way. The necessity of these designs are experimentally justified with ablation study in Section 6.3
Leveraging the simple functional form introduced in Section 4
Edge embedding. style continuous filter convolution and the simple concatenation of the scalar-valued distance in
To embed pairwise interactions, we combine the SchNet (Schütt et al. 2017)
4

---

Published as a conference paper at ICLR 2023
Algorithm Spatial Attention Kinetic Networks Layer
function SAKELAYER((15 (x(x)) 1,
> Input embedding, position, and velocity
for u E do for E (v) do hew + xz the , ha (eoe))
Edge update, Sec
Eq
end for
(1+11) hewr xxH
> Semantic attention and distance cutoff, Sec
5
Bq
10
for E (v) do
RSA,
(k+1)
(bear, Bun) (SA((
Spatial attention, See
Bq
end for av + Σ henr
Neighborhood aggregation, See
Eq
"cN(o)
((+11) t +0"+v((1 (v(()) )v +W. E A( wEN(u)
Vel. update, See
Bqq
12
Xi
(K)
(k) +v:e) (K)
Position update, Sec
5
5q
12
ho
(6+1) t 6'(ho ,00 ASA
(k) (8) (K)
& Node update, Sec 2.2
Eq
return (he (httt)) fxl to (v:
(k+1)) (k+1))
end for end function
Satorras et al. (2021) to achieve the balance between high radial resolution and large receptive field. The resulting edge embedding is thus
h (k)) oc (h!! @ h(*) + ||e(
@ RBF(
("
)))
where ④ denotes Hadamard product and f" is a (filter-generating) feed-forward network as in Schütt al. (2017).
Semantic attention and distance cutoff. To promote anisotropy (which Dwivedi et al. (2020) finds useful GNNs) in the pooling operation, apart from spatial attention introduced in Section 4 we compute the attention score on semantic and geometry space to weight interactions among particles based on embedding similarity and distances on n-dimensional space. The distance weights are calculated using the cutoff function proposed in Unke and Meuwly (2019):
{ cos( ~||6,,ا|| + 1), l|u,,vll S do;
X 16
dg
(8)
(0,
L do:
to filter out interactions outside a certain range (do). And semantic attention weights a are calculated similar Graph Attention Networks (Veličković et al. 2018),
exp(o exp(a(a"hee..))
H a
Σ explo (athee. hewr))
To produce models for the purpose of molecular simulation by ensuring continuous forces and gradients, one would need to choose as sigma a as at least C2-continuous activation functions, namely CeLU (Barron 2017). These weights are combined and normalzied:
XxH a M2
(10)
Σ EaxaH aH
trivial to expand multi-head attention with sets of (a......) .. and the resulting combined attention weights concatenated.
Since the edge embedding after mixed Fuclidean and semantic attention he.,, *a Xx already encodes the desired inductive bias that nodes farther away from each other would have less impact on each other's embedding, we directly use this representation and simply set X as a linear projection (with weights Wx) and f as identity in Equation 6
(hee...) = Wxhe.n f(C.v) = / l
(11)
We leave more elaborate choices of À and f functions for future study.
5

---

Published as a conference paper at ICLR 2023
Fictitious velocity integration. Similar to Satorras al. (2021) we keep track of a fictitious velocity variable V, for each node and linearly combine it (with weights Wx) update it and positions Xy in turns like Euler-discretized Hamiltonian integration.
v(++1) = =0">V(h
W.
Σ E X. MEA((e)
(12)
(k+1) X x + (R)
(13)
During velocity update, we scale the current velocity and, for the sake of parameter saving, reuse the same set of linear combinations edge vectors used in Equation 6
the additive term.
Remark 2. If hew- is distinct among edges, there exist sets of X (even when f I is the identity) such that Equation 12 can approximate any vector on the subspace spanned by edge vectors and is thus universal up equàvariance thereon.
Equivariance analysis. We inlay these modular components into the framework of graph neural
networks (described in Section 2.2) to produce the SAKE architecture, which we show in Algorithm 1. (A SAKEModel is defined by applying SAKELayer iteratively from k = 0 to & = K - 1, the depth of the network.) We have previously discussed that spatial attention is E(n)-invariant in Remark [ With the H(n) equivariance of velocity update and position update (Equation [2] being proved in Satorras et al. (202T) and the rest of the model only takes the norm of edges and is E(n) invariant, SAKE is H(n) equivarian.
Runtime analysis. The runtime complexity of spatial attention (Equation which is the bottleneck of this algorithm, is ((E)) i.c., linear in the number of graph edges [E), number of attention weights N,,, embedding dimension C, and geometric dimension D. When implemented, the FOR loops in Algorithm can be packed into multi-dimensional tensors that could benefit from GPU acceleration.
Relation to spherical harmonics-hased models. Under the framework of TFN TThoas et al. 2018). the embedding and position imput in Algorithm [h (1(()), {xx (x{²)} corresponds to the / 0 and
7 T= type tensors. The concatenation followed by neural network in Equation i loosely analogous
to the direct sum operation in Clebsch-Gordon decomposition. While Smith et al. (2017); Schiit et al.
(2017): Satorras al. (2021) operates on 1 = 0 tensors only, the spatial atfention mechanism
propose (Section Equation 6) and velocity/position update (Section 5 Equation 12) corresponds to
1 + 1 and @ 0 type networks, respectively. Klicpera al. (2021a); Villar et al. (2021) have discovered that = 0, are the complete levels of tensor to universally describe geometry on E(3), while higher-order tensors are not necessary completely describe the node environment.
Kovács et al. (2021) also discusses the concept of density projection where body-order functional forms can recovered by lower-order terms.
EXPERIMENTS
As discussed in Section 2 SAKE provides mapping from and to the joint space of geometry and embedding X' x H. while being equivariant on geometric space and invariant on embedding space. We are therefore interested characterize the performance of SAKE on two types of tasks: invariant modeling (Section 6.I. where we model some scalar property of a physical system, namely potential energy; and equivariant modeling (Section 6.2) where we predict coórdinates conditioned on initial position, velocity, and embedding.
On both classes of tasks, SAKE displays competitive performance while requiring significantly less inference time compared to current state-of-the-art models. See Appendix Section for experimental details and settings.
INVARIANT TASKS: MACHINE LEARNING POTENTIAL CONSTRUCTION
MD17 potential energy (Chmiela et al. 2017 tests the capacity of the model in the extreme small data regime. Its training set contains merely 1000 configurations and quantum chemical energies and forces of single small molecules in vacuum computed using density functional theory (DFT). As

---

Published as a conference paper at ICLR 2023
Table I: Inference time (ms) and energy (E) and force (F) mean absolute (MAB) (meV and V/A) the MD17 quantum chemical dataset.
SchNet DimeNet
sGDML
PaINN
NequlP
SAKE
Klaper p.0le
Inference time
huich of 32
88/376
206
31
38/99
197
4
E
16.0
8.8
8.2
6.9
5.3
5.91
5.915.2
Aspirin
F
58.5
21.6
29.5
14.7
9.4
8.2
8.09
E
3.5
2.8
3.0
2.7
2.2
2.20 20
Ethanol
2.20
F
16.9
10.0
14.3
9.7
3.7
3.8
2.75
E
5.6
4.5
4.3
3.9
3.3
3.22
Malonaldehyde
F
28.6
16.6
17.8
13.8
6.7
5.8
4.32 4.3
E
6.9
5.3
5.2
5.0
4.9
4.91 4.92
Naphtalene
F
25.2
9.3
4.8
3.3
2.2
1.6
2.25 2.205.21
E
8.7
5.8
5.2
4.9
4.0
4.67 167
Salicylic acid
4.66
F
36.9
16.2
12.1
8.5
5.4
3.9
4.29 1.28
E
5.2
4.4
4.3
4.1
4.0
4.00
Toluene
24.7
2.10
E
6.1
5.0
4.8
4.5
4.5
4.51 4.50
Uracil
F
24.3
13.1
10.4
6.0
4.2
3.3
4.26 4.25
summarized in Table T SAKE consistently outperforms benchmarked models with the exception of NequlP (Batzner al. 2021), which has comparable performance but is noticeably slower. (Note that, to avoid inaccurate report on the suboptimal software configuration, we only report runtime data directly quoted from their original publications and replicate the hardware environment by ourselves. The same applies hereafter.) In terms of training cost, most state-of-the-an models requires days of completed within hours. This significant advantage in
training, whereas the MD17 experiment speed would allow SAKE to be more rapidly trained and deployed on realistic applications involving molecular dynamics simulations.
(Schütt et al. 2017) goes beyond the single-molecne regime. It involves slightly more diverse chemical space containing 5000-step ab initio molecular dynamics simulation trajectories (with energies and forces) of 129 molecules with the same formula C7H10O2. The test set of ISO17 split into known (which Kovács et al. (202T) argues to be very close to training set) and unknown molecules, based on the chemical identity (topology) at the beginning of the simulation, which could be regarded as interpolative and extrapolative tasks, respectively. As shown in Table SAKE significantly outperforms other models the unknown molecules in the test set, indicating that SAKE is capable of extrapolating and generalizing onto unseen chemical spaces when trained on limited data.
IS017
QM9 (Ramakrishnan et al. 2014) tests the transductive generalizability across distinct small chemical graphs. It entails a very diverse chemical space of 134k molecules with annotated phsyical properties calculated with SELLP/6-31G(2d555 cquilibrium (low energy) conformations. In Table 3) SAKE achieves statc-of-the- performance at predicting HOMO, LUMO, and Ac properties class of most crucial molecular properties closely related to reactivity. Interestingly. SAKE performs competitively on predicting extensive physical properties but not intensive (also see discussions in [Pronobis et al. (2018)). We hypothesize that
level of quantum chemistry, albeit all with at-
this will be mitigated by choosing size-invariant pooling functions- we leave this for future study.
6.2 EQUIVARIANT TASKS
Charged N-body dynamical system forecasting (Kipf et al. 2018 Fuchs et al. 2020 tests if model can predict the evolution of physical system sufficiently long after initial conditions. This simple system consists of charge-caryying particles with initial positions and velocities, and the position at given moment is predicted. As shown in Table 4 although the interactions (Coulomb forces) are entirely pairwise, we see here that the additional expressiveness SAKE affords lead to competitive performance this demonstrative task.

---

SAKE Published as a conference paper at ICLR 2023
Table 2: Test set energy (E) and force (F) mean absolute enror (MAE) (meV and meV/A) on known and unknown molecules IS017.
ACE
SchNet
PhysNet
E
16
16
12.17 18
known
F
43
43
12.33
2
unknown
E
85
104
127
53.37
13-15
03.02
F
85
95
60
39.46
39.59
59.35
Table QM9 lest set performance (mean absolute error).
Åe
HOMO
LUMO
C2
Bohr
meV
meV
meV
cal/mol
SchNet
0.235
63
41
34
0.033
0.033
DimeNet+H 0.044
33
25
20
0.030
0.023
SE(3)-TF
0.142
53
35
33
0.051
0.054
EGNN 0.071
48
29
25
0.029
0,031
PaiNN CEEEP
0.059
36
46
20
0.012
0.024
TorchMD-Net 0.059
36
20
17
0.011
0.023
SphereNet Camr 2007
0.030
31
19
23
0.025
0.022
SAKE
0.068
23
16
13
0.014
0,087
MD17 forecast (Chmiela et al. 2017 Huang ct al. 2022) involves a simulation forecast task with slightly more complicated systems com- pared to Table 4 uses the same dataset in Table but predicts the time evolution of the molecular dynamics simulation directly, rather than predicting the mapping from the geome- try to potential energy. Following the proto- col Huang et al. (2022), we predict the atom position based on the velocity and coördinate 3000 steps prior. As shown in T SAKE achieves superior performance on out of 8 sys- tems, without leveraging spherical harmonics (as in TFN (Thomas et al. 2018) or SE(3)-TF
Table 4: Mean Squared Error (MSE) and inference time (ms)for charged particle dynamic system fore- casting.
Architecture
MSE
Inference time
SE(3-TF
0.244
0.1346
TFN
Thrmee HIr
0.155
0.0343
GNN
0.0107
0.0032
EGNN
Sulma sta. JO
0.0071
0.0062
SAKE
0.0049
0.0079
SEGNN Brondhatiee stal, MAT
0.0043
0.0260
(Fuchs et al. 2020i) hand-coded edges (as in
GMN (Huang et al. 2022)).
Table 5: Mean squarred error (MSE) (10-2 A on MD17 trajectory forecast.
Aspirin Benzene Ethanol Malonaldehyde Naphthalene Salicylic Toluene Uracil
TFN
12.37 58.48...
4.81
13.62....
0.49
1.03.=
10.89
0.84
Naw. LIX
SH(3-TF
11.12ks:
68.1 1mp
4.74ma
13.89
0.52
1.13am
10.88
0.79
EGNN
14.4 62.40mm.
4.64
13.64
0.47
1.02az
11.78
0.64am
GMN
9.76mmn
48.12
4.63Aam
12.82
0.40amm
0.88
10.22...
0.59
SAKE
9.33 am
137.20
4.63
12.81
0.38
0.82....
10.98
0.53
Walking motion capture (CMU 2003 has higher system complexity and noise and adopted
to demonstrate SAKE's general capacity to forecast dynamic systems beyond microscopic scale. In
this task, again closely following the experiment setting of Huung et al. (2022); Kipf et al. (2018). we predict the position of walking person (subject 35 in CMU motion capture database (CMU 2003)) based on their initial position. Again, we observe that SAKE outperform other models by a large margin (Table and is significantly faster.
8

---

Published as a conference paper at ICLR 2023
6.3 ABALATION STUDY
Table 7: SAKE performance on MDI7-Aspirin (also see Table 1 with various components included
(Y) or excluded (N).
Spatial Semantic
Speed and
attention attention
position update Energy RMSE Force RMSE
Eq.
Eq. 10
Eq 12
(me V)
(meV/A)
Y
5.915.
8.09 10
Y
8.08
17.48
17.47
Y
Y
Ñ
6.21
10.31
10:33
N
Y
N
8.15
16.5816:51
Y
N
Y
7.88735
16.21
T6.22
16.19
10.78
17.90
Y
N
Ñ
6.20
10.52 1034-
N
N
N
8.21 8.24
16.62 16.64
To elucidate each component's contribution to- wards the final performance, we perform an ab-
lation study on one of the most popular tasks studied so far MD17 potential energy model- ing (Section 6.1 [T) More specifically, we focus on the most complicated molecular system in the dataset, aspirin. We inspect three compo- nents proposed in the paper-spatial attention (Equation 6). which we argue is the main novelty of the paper. well semantic attention (Equa- tion I0) and velocity and position update (Equation 12). As is shown in Table([] spatial attention improves the performance regardless of the rest of the configuration. To clearly demonstrate the utility of these auxiliary modules, we also add these components (except spatial attention) to the backbone of EGNN (Satorras al. 2022), and summarize the results in Appendix Table qualitatively, with these modules, EGNN achieves decent performance on par with SchNet (Schütt et al. 2017). with the performance gap comes primarily from its inability conceive the node angular environments.
Table 6: Walking motion capture performance.
GNN
EGNN
GMN
SAKE
MAE 67.321.
59.1#2.1
43.9+1.1
14.59
=1.6
Epoch time
5.66
1.81
DISCUSSION
Conclusions In this paper we have introduced an invariant/equivariant functional form termed sputial attention that uses neurally parametrized linear combinations of edge vectors to equivariantly yet universally characterize node environments at a fraction the cost state-of-the-an equivariant approaches that makes use of spherical harmonics. Equipped with this spatial attention module, we use it to build a spatial attention kinetic network (SAKE) architecture which is permutationally, translationally, rotationally, and reflectionally equivariant. We have demonstrated the utility of this model in n-body physical modeling tasks ranging from potential energy prediction to dynamical system forecasting.
Limitations Theoretical: The universality condition is only discussed w.r.t. node geometry, without considering node embeddings; moreoever, the inequality condition in Theorem I although rarely violated in physical modeling (even when system is highly symmetrical), can be potentially over- restricting. We plan to generalize this framework to consider the expressive power of functions on the joint space of node embedding and geometry in future works. Experimental: Herein, apart from the novel functional form spatial attention (Equation 6). the rest of the architecture has not been thoroughly optimized and analyzed in the context of the growing design space of equivariant neural networks.
Future directions We plan to conduct more thorough experiments on (bio mole)ual systems to explore the potential of SAKE in building general protein/small molecule force fields and enhanced sampling methods that could facilitate large-scale simulations useful in therapeutics and material discovery.

---

Published as a conference paper at ICLR 2023
Social impact and ethics statement This work provide accurate and extremely efficient way to approximate properties and dynamics molecular systems and physical states. may advances research in a wide range of disciplines, including physics, chemistry, biochemistry, biophysics, and drug and material discovery. As with all molecular machine learning methods, negative implications may be possible if used in the design of explosives, toxins, chemical weapons, and overly addictive recreational narcotics.
Reproducibility statement The software package containing the algorithm proposed here is distributed open source under MIT license. All necessary code, data, and details to reproduce the experiments can be found in Appendix Section
Funding Research reported in this publication supported by the National Institute for General Medical Sciences of the National Institutes of Health under award numbers RO1GM132386 and R01GM140090. YW acknowledges funding from NIH grant G0133122 and the Sloan Kettering Institute. JDC acknowledges funding from NIH grants ROIGMI 32386 and RM1GM1 40090.
Disclaimer The content is solely the responsibility of the authors and does necessarily represent the official views of the National Institutes of Health.
Disclosures JDC current member of the Scientific Advisory Board of OpenEye Scientific Software, Redesign Science, Ventus Therapeutics, and Interline Therapeuties, and has equity interests in Redesign Science and Interline Therapeutics. The Chodera laboratory receives or has received funding from multiple sources, including the National Institutes of Health, the National Science Foundation, the Parker Institute for Cancer Immunotherapy. Relay Therapeutics, Entasis Therapeutics, Silicon Therapeutics, EMD Serono (Merck KGaA), AstraZeneca, Vir Biotechnology. Bayer, XtalPi, Interline Therapeutics, the Molecular Sciences Software Institute, the Starr Cancer Consortium, the Open Force Field Consortium, Cycle for Survival, Louis V. Gerstner Young Investigator Award, and the Sloan Kettering Institute. A complete funding history for the Chodera lab can be found at ht attp://choderre org. funding
The authors would like to thank the ICLR reviewers for providing constructive feedback that substantially improved the quality and clarity of this manuscript. YW thanks Theofanis Karaletsos, Insitro, and Andrew D. White, University of Rochester, for useful discussions and Leo Klein, Freic Universität Berlin, for catching an embarrassing bug.
Acknowledgments
REFERENCES
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. CoRR, Nel11669.02990 2016.
/669 02907
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful graph neural networks? arXiv preprint 11X:0:1810.008866, 2018.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. arXiv preprint arXiv:1704.012/11 2017.
Will Hamilton, Zhitao Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Advances in neural information processing systems, pages 1024-1034, 2017.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzaler Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: arXiv:/806.01 /806,0/21, 2018.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, E(n) equivariant graph neural networks. CoRR, abs/2102.098444, 2021. http::/ /ttp::///aaii..aa/ wtwaaaivvvvv...../ /agg/ab/// /2102/12.20 09844
Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Claus-Robent Müller, Schnet: continuous-filtet convolutional neural network for modeling quantum interactions, 2017.
10

---

Published as a conference paper at ICLR 2023
Nathaniel Thomas, Tess Smidt, Steven M. Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-qquivariaat neural networks for 3d point clouds. CoRR, abs/1802.08219, 2018. URL/attp://axxv..oo
abs aae////aaa//// 11/008800000 08201.2219 08219
Johannes Klicpera, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional
graph neural networks for molecules. In Conference on Neural Information Processing Sys- tems (NeurIPS), 2021a.
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. te(3)-transformmrs:: 3d roto-translation equivariant attention networks. In Advances in Neural Information Processing Systems 34 (NeurIPS). 2020.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Korblluhh, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. Eu)evvairariian graph neural networks for data-efficient and accurate interatomic potentials, 2021.
Brandon Anderson, Truong-Son Hy, and Risi Kondor. networks, 2019.
Cormorant: Covariant molecular neural
Zonghan Wu, Shirui Pan. Fengwen Chen, Guodong Long, Chenggi Zhang, and Philip S. Yu. A comprehensive survey on graph neural networks. CoRR, abs/1901.00556, 2019. URL http:
//axxiv.org/abs.. 901 00596
Jiaqi Han, Yu Rong, Tingyang Xu, and Wenbing Huang. Geometrically equivariant graph neural networks: A survey, 2022. URL https::/ /tttp:://axii...
abs/ www.ab////.2.....
S. Smith, O. Isayev, and A. E. Roitberg. Ani-l: an extensible neural network potential with dft accuracy at force field computational cost. Chemical Science, 8(4):3192-32033 2017. ISSN 2041- 6539. doi: 10. 0o.139//6x005720a. URL /ttp::/p./.....oooooo///////////
doi wwtp:/..oo1100000000..000.00...0
Jörg Behler. Atom-centered symmetry functions for constructing high-dimensional neural network potentials. The Journal of Chemical Physics, 134(7):074106, 2011. doi: 10.1063//.355371 URL
https: /tt:::/doooooo .or.i..org//oo 10
Johannes Klicpera, Shankari Giri, Johannes Margraf, and Stephan Günnemann. and uncertainty-awar directional message passing for non-equilibrium molecules. CoRR, abs/2011.14115, 2020a. URL.jttps::/aaxix http:///xxxvv.ogaa////////0/////
Fast
14115
Yuanqing Wang. Josh Fass, Benjamin Kaminow, John E. Herr, Dominic Rufa, Ivy Zhang, Iván Pulido, Mike Henry, Hannah E. Bruce Macdonald, Kenichiro Takaba, and John D. Chodera. End-to-end differentiable construction of molecular mechanics force fields. Chem. Sci., 13:12016 12033, 2022. doi: 10. 1033///2CC02739..
BRLphtt:: /ttp:://xx doi 1ttp://0.1119/222220 wttp./....///..110////2222339390
Yuanqing Wang, Iván Pulido, Kenichiro Takaba, Benjamin Kaminow, Jenke Scheen, Lily Wang, and John D. Chodera. Espalomacharge: Machine learning-enabled ultra-fast partial charge assignment, 2023. URL httl:ttts:
abs /alob//23330 22022 06758
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference Learning Representations, 2022. URRpsttts:///peereee//
Johannes Klicpera, Chandan Yeshwanth, and Stephan Günnemann. Directional message passing on molecular graphs via synthetic coordinates. CoRR, abs/2111.04718,, 2021b. URL https
fariiv. www./bg..a......o. htagg//bs/200 /211/ 04718
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling, Geometric and physical quantities improve E(3) equivariant message passing. CoRR, abs/2110.02905.. 2021. URL
Soledad Villar, David W. Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum-Smith. Scalars universal: Equivariant machine learning, structured like classical physics, 2021. URL https
w/arxiv.orr wwtpp.wwwwww..././...oo 06610
11

---

Published as a conference paper at ICLR 2023
Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra, 2021.
Philipp Thölke and Gianni De Fabritiis. netthmd-net: Equivariant transformers for neural network abs/ www.:eeeoorrmmooooiooitee..eeee
based molecular potentials. CoRR, abs/2202.02541, 2022. URL htttt:///arxxi///
F2C22.0541
Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equiv- ariant graph mechanics networks with constraints. In International Conference on Learning Representations, 2022. URL
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael J. L. Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons, 2021.
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Velickovic. Principal neighbourhood aggregation for graph nets. CoRR, abs/2004.05718, 2020. URL https://
a arxiv. arxiv.org/abe.. abs 2004.0077
Vijay Prakash Dwivedi, Chaitanya K. Joshi, Thomas Laurent, Yoshua Bengio, and Xavier Bresson. Benchmarking graph neural networks. CoRR, abs/2003.00982 2020. URL https: https::/arxiv /arxiv
org/ arg/abs/200 2003 00982
Oliver T. Unke and Markus Meuwly. PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of Chemical Theory and Computation, 15(6): 3678-3693 may 2019. doi: 10.102//acsjeet::90000 URL http::::///oo.oggg////.o 10218 2Facs. jote. 9b00181
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks, 2018.
Jonathan T. Barron. Continuously differentiable exponential linear units, 2017. URL https
7/arxiv.
Dávid Péter Kováes, Cas van der Oord, Jiri Kucera, Alice EA Allen, Daniel J Cole, Christoph Ortner, and Gábor Csányi. Linear atomic cluster expansion force fields for organic molecules: beyond rmse. Journal of chemical theory and computation, 17(12)7696-7711, 2021.
Johannes Klicpera, Janck Groß, and Stephan Günnemann. Directional message passing for molecular
graphs. CoRR, abs/2003.031 2020b. URL
03123
Stefan Chmiela, Huziel E. 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-44 jul 2019. doi: 10. 00.10166..cp22211..0.0000 URL
/tttp:://di..oz/10..11
2019 02. 007
Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science Advances, 3(5), May 2017. ISSN 2375-2548, doi: 10. 10.1126/siiad URL
http://dxx doi www.orgg//.... 1126/ http://saaaav...1111.11. 1603015
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177, 2014.
Wiktor Pronobis, Kristof T Schütt, Alexandre Tkatchenko, and Klaus-Robert Müller. Capturing intensive and extensive dft/tddft molecular properties with machine learning. The European Physical Journal B, 116 2018
Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems, 2018. URL /tRp/tttpp///aai.ooo///aaa////
04687)
Victor Garcia Satorras, Emiel Hoogeboom, Fabian B. Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows, 2022.
CMU. Cmu motion capture database. http://mocap.cs. cmu .dd 2003.
12

---

Published as a conference paper at ICLR 2023
Timothy F Havel. Distance geometry: Theory, algorithms, and chemical applications. Encyclopedia of Computational Chemistry, 120:723-742, 1998.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neural Networks, 2(5):359-366 1989. ISSN 0893-6080. doi: https:
URL https:/ /tt:://wwwwwwe www.cclenedireettttomm com/
Roberto Bolaño. 2666. Anagrama, Barcelona, 2004.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows, 2016.
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji
Lakshminarayanan. Normalizing flows for probabilistic modeling and inference, 2021.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: Exact likelihood generative learning for symmetric densities, 2020.
Danilo Jimenez Rezende, Sébastien Racanière, Irina Higgins, and Peter Toth. Equivariant hamiltonian flows, 2019. /tt.attss///axxt.or/////////////
Changyou Chen, Chunyuan Li, Liqun Chen, Wenlin Wang. Yunchen Pu, and Lawrence Carin. Continuous-time flows for efficient inference and density estimation, 2018.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations, 2019.
Chin-Wei Huang, Laurent Dinh, and Aaron Courville. Augmented normalizing flows: Bridging the gap between generative flows and latent variable models, 2020.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp, 2017.
PROOFS
8.1 PROOF FOR THEOREM 1
Proof. (Informal) An appropriate set of choices to satisfy the requirement is:
I, and takes the set of [N(v)] + (2)) elements:
(he(ee. =i = -]} U
(14)
where [-] is the Iverson bracket. The first (w)) set of one-indexed functions give the node-to- neighbor distances whereas the second (o))) 2 -et of two-indexed functions give the distances among neighbors. This recovers the distance matrix (x among the node and its neighbors
Ax(u)= | X., ll and Ax(v); ||x,; x., + j, < i, < [W(v))- As such, the relative positions of the node and its neighbors are uniquely defined up to E(n) symmetry (Havel) arbitrary coördinate system which g is evaluated. Since {사] is neural, by the universal approximation theorem (Hornik et al. 1989). (x(v) can still be recovered, least through recovering the set of (사} discussed above. Considering again the universal approximation of /4, any function the geometric node environment be approximated by @SA.
1998). In other words, x, and X,, be uniquely embedded
DETAILED METHODS
9.1 CODE AVAILABILITY
The corresponding software package and all scripts used to conduct the experiments in this paper distributed open source under MIT license at: This package be installed via: pip install sake-gnn
com/ choderalab sake
13

---

Published as a conference paper at ICLR 2023
9.2 HARDWARE CONFIGURATION
All models are trained on NVIDIA Tesla V100 GPUs. Following the settings reported in the publications of baseline models. the inference time benchmark experiments (Section 6.2 6.1 done on NVIDIA GeForce GTX 1080 Ti GPU (For Table and Table 66) and NVIDIA GeForce GTX 2080 Ti GPU (For Table 4).
9.3 ARCHITECTURE AND OPTIMIZATION DETAILS
One-layer feed-forward neural networks are used as fr in Equation (edge update); two-layer feed-forward neural networks are used 6 in Equation (edge update), ø" Equation [ (node update), < in Equation 12 (velocity update), and in Equation 6 (spatial attention). SiLU is used
everywhere as activation, except in Equation [ (velocity update) where the last activation function constraint the velocity scaling to between 0 and 2 and in Equation 10
chosen y = 2 Sigmoid((z) where CeLU is used before attention; additionally, tanh is applied on the additive part of Equation 12 to constraint it to between -1 and 1. 4 attention heads are used with y in Equation TO spaced evenly between O and 5 Á. 50 RBF basis are used, spacedly evenly between O and 5 Á. All models are optimized with Adam optimizer. We summarize the hyperparameters used these experiments in All random seeds are fixed 2666, homage to Bolaño (004)..
Table 8
Experiment
Depth Width Learning Rate Epochs Batch Size 1.2 Regularinn degularization Cutoff
MD17 (Table
Aspirin
32
10 3+
5000
10
5.0
MD17 (Table
Ethonal
32
10-3+
5000
10-5
10.0
MD17 (Table
M
Fllonaldeyyde
8
64
10-3+
5000
4
10-5
10.0
MD17 (Table
Naphthalene
64
10-3+
5000
10
5,0
MD17 (Table
Salicylic Acid
10-3+
5000
10-5
5.0
MD17 (Table [ Toluene
32
10-3+
5000
10-5
5.0
MD17 T Uracil
64
10-3+
5000
4
10
5.0
MD17 Trajectory Forcast (Table
8
10-3
1000
10
ISO17 (Table
10-3+
100
10
QM9 (Table
32
10-4+
5000
32
10-10
N-Body Forecast 4able
32
5 x 10
1000
100
10-12
Table 8: Hyperparameters used in experiments ( A cosine warm up and annealing schedule is used, where the learning rate is gradually increased from 10-6 to the peak value the first 10% epochs and decreased in the rest 90%.)
9.4 DATA AVAILABILITY
The source and details of benchmark datasets summarized in Table.
Experiment
Split
MD17
Equat//unnt machune opant
8 Systems, TOOK IM snapshots
Random: IK Train Fixed
ISOI
129 Molecules; 5000 snapshots
QM9 inp-Squeee madare eeplaaune
135k molecules particles
Fixed
N-Body Forecast
Table
MIT
Fixed: 3K Train; 2K Valid: 2K Test
Table 9: Dataset details.
10 ABALTION STUDY EGNN
11 BRIEF INTRODUCTION OF EQUIVARIANT NORMALIZING FLOWS
Normalizing flows (Rezende and Mohamed
2016 Papamakarios et al., 2021 are family of learnable
bijections f=x 2 - X' that transform a tractable distribution on latent space (z( to another on the target space x fix(2: 8) (dropping dependency on parameter thereafter) whose density can be analytically written
log qu (x) = log qu (z) + log det Ofaz(x)
(15)
8x
14

---

Published as a conference paper at ICLR 2023
Table 10: EGNN (Satorras al. 2022
performance MD17-Aspirin (also see T with various
components included (Y) or excluded (N).
Distance Semantic
smearing attention Position update Energy RMSE
Force RMSE (meV)
Eq.
7
Eq. 10
Eq. 12
(meV/A)
Y
Y
Y
25.922%
42.83
89
78
VIGIA
MNNK
Y
VNND
N
YYVE
16.0216:08
31.57
15.7
29.59 29.69
36.9797
N
Y
NNNIN
589.34588 500
17.34 17.31
17.99
217.15
33.85 33.00
$2
INN
Y
588.95
218.73311
N
N
250. 250.182311
538.68
33..11
N
N
206.13200 205.79
882.32
329.32882.81
882.81
538
882.03
and conversely from a sample on the target space to the latent space z = fxx whose likelihood given by
log Pz (z) = log (x) + log det O5 (2)
(16)
Oz
where fax = fazl. To close the gap between the intractable ps and the tractable Qz. that one
sample on the target space efficiently, the flow then trained by maximizing either Equation T5 an unnormalized target distribution is given, Equation [6] if samples are given.
The concept of equivariant normalizing flow was first introduced in Köhler et al. (2020): Rezende al. (2019); Satorras et al. (2022), where they adopted the framework of continuous normalizing flow (Chen et al. 2018) to define the bijection fex as
x= f f2 (z()
(17)
and restrict f E(n)-equivariant w.r.t. Integration, the sum over infinitely many equivariant functions, does not alter the equivariance. To numerically approximate this integration Chen et al. (2019) involves evaluating f'. multiple times and is therefore expensive.
12 EQUIVARIANT EXACT LIKELIHOOD SAMPLING WITH SAKE FLOW.
Current equivariant normalizing flow models (aaooraas al. 2022 Köhler al.) 2020) (brielfy reviewed in Appendix Section [T]. relies on ODE-based numerical Integration(Cher et al. 2019), are computationally expensive. We propose much simpler invertible flow model that uses SAKE model, termed SAKE Flow. First, following scheme introduced in Huang et al. (202 (in which it argued that with flexible enough kernels Equation 17 could be approximated arbitrarily well), extend the space X with an auxiliary space A. Correspondingly, we extend the tractable distribution be qzz, a) = 9= (z(2)) (z)4a (a) and the target distribution be p(x, a) = Pz (z) We then change the problem statement of Equation 16 to: find a parametrized function fax (z, a 0) far (x, a f)
that is a bijection on the space
X to maximize the joint likelihood:
i = argmax Eawg, [log p(x, a))]
(18)
which is, up to a constant, an evidence lower bound for x:
log Pa
(19)
Egg p(x, a)] + DRL. (a(al/(a)lll(axx) + H.
(20)
E,g [logp(x, a)] Ha.
(21)
The equality holds when the nnon-negative) Kullback-Leible divergence between the tractable distribution and the conditional distribution given samples p(a x) is zero. As such, an unbiased estimate of the marginal likelihood can be given by
log pxx) log p(x, a) log (a).
(22)
15

---

Published as a conference paper at ICLR 2023
Similar to Huang al. (2020); Dinh al. (2017). we define fzx a, 0) as series of alternating affine coupling:
ta/a (z, a) z, exp(S* (2)) 3 + T
(23)
with analytic inverse. Composing these transformations gā ga o ...ga ga we get our bijection Lea (2, a; 0).
Now, it only remains to define the structure of the translation functions {T) and scaling functions {5). As is outlined in Satorras al. (2022). if fee is translatio-iivaaiiane is impossible have Ps x) d 1, so we drop the translation invariance requirement and design fzz, as well as the composing (T.S}; be only rotation and reflection equivariant. Correspondingly, we require all tractable distributions to be confined on (IV) 1) jnn-dimensional subspace with O gravity center. (Note that this reduces the symmetry group we work on from E(n) SO(n).) A valid choice both for 4s (z) and (a) with is a centered Gaussian distribution (Satorras et al. 2022): :
p(z) (2 (2 1|n/2
exp( 1ap(~2|1z]] 2 -||z|(²)
(24)
with Σ x, = 0. Moreover, to keep the gravity center at O, we require that (T,S} shall not change PEV the gravity center. A general recipe to construct such ST is to use the equivariant and invariant outputs of SAKE parametrize T and respectively.
h x SAKEModel(h, h, x);
MEEN( exp(M)AN()))))))
(25)
T(z) = z x
To preserve the center of gravity, we center the translation have zero center of gravity and enforce the same scalar to be used across particles as the scaling factor.
16


Reviews:
[{"review_id": "yIMNpaf9O1", "reviewer": "Reviewer_ULzV", "summary": "", "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": "", "summary_of_the_paper": "This paper exploits several kinds of attention mechanisms (Euclidean and semantic attention, spatial attention) upon EGNN [6]. Necessary theoretical claims have been made to justify the expressivity of the proposed spatial attention. Experiments are conducted on invariant tasks and equivariant tasks.", "summary_of_the_review": "The authors are tackling a variable problem on enhancing the expressivity of equivariant GNNs on physical modeling. However, given the limited novelty, technicality, and insufficient ablations, I suggest weak reject for its current version. ", "strength_and_weaknesses": "Strengths:\n\n1. The design of spatial attention in Eq. (5) is well motivated and novel. As illustrated in the proof of Theorem 1, if certain condition meets, the spatial attention is able to recover the local environment of each node. The proof of Theorem 1 is interesting by recovering the distance matrix.\n\n2. The part of Related Work is well organized, providing a good guidance of understanding of how recent progress is made.\n\n3. The experimental comparison with various methods have been conducted. \n\nWeaknesses:\n\n1. The condition $h_{e_{vu_i}}\\neq h_{e_{vu_j}}$ in Theorem 1 is over-restricted and easily broken when two neighbor nodes share the same embedding (the same node features and local context). The authors claimed that the inequality holds when the system is not strictly symmetrical. Unfortunately, many practical molecular systems are symmetric such as benzene and many other materials such as crystals. \n\n2. How to build a permutation-invariant and E(n)-invariant (or equivariant) function has been discussed in [A], which is not cited, unfortunately. In [A], Propositions 10-11 have presented the universal form of message passing from local neighbors up to the O(n) equivariant. More importantly, the results by [A] does not require the condition $h_{e_{vu_i}}\\neq h_{e_{vu_j}}$ that is indispensable in this paper. The authors are suggested to compare the difference with [A] which seems more elegant. \n\n3. Moreover, the proof of Theorem 1 only shows the results on invariant function other than equivariant function. In [B] and its citated references, the universal form of equivariant message passing can be derived from TFN and GemNet, both of which again do not require any restriction of node embedding. \n\n4. It is unclear why \"The distinctness condition corresponds to the full-rank requirement in GMN\".\n\n5. The authors have shown experimental comparisons on MD17 and QM9, which are already well explored, and the results of the proposed method are not clearly superior to other methods. The authors are strongly suggested to conduct performance on other recent and more challenging tasks, such as the trajectory prediction setting of MD17 and the benchmark OC20-22. \n\n6. The ablation studies are insufficient and unclear to justify the contribution of this paper. It seems from Algorithm 1 that the main novelty lies in the proposal of spatial attention Eq. 6 along with other minor techniques. The authors should detail how each proposed component affects the performance and provide necessary analyses.\n\n\n\n\n[A] Scalars are universal: Equivariant machine learning, structured like classical physics, NIPS 2021.\n\n[B] GemNet: Universal Directional Graph Neural Networks for Molecules, NIPS 2021.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "piVsoMu4ej", "reviewer": "Reviewer_YAqW", "summary": "", "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": "", "summary_of_the_paper": "The paper proposes an E(n) equivariant network acting on pointclouds with edges. The method is simple to compute and shows competitive performance on various tasks.\n\nThe authors claim that their attention method is universal, and that it is more expressive than \"dot product scalarization\" attention used in prior wok.", "summary_of_the_review": "I have serious concerns about both theoretical claims in this paper: that the spatial attention mechanism is universal, and that it is more expressive than different attentional methods. I thus recommend rejection of this paper.\n\n---\nAfter their revision, the theoretical issues are fixed. I now think the paper is a valuable contribution to the literature on equivariant pointcloud networks and should be accepted. I weakly recommend acceptance, as the novelty is still limited.", "strength_and_weaknesses": "## Strength\n- The method is evaluated on a wide variety of tasks.\n- The method performs competitively.\n\n## Weakness\n- The core spatial attention method is a small variation of / combination of pre-existing methods.\n\n### Representation of theorem 1\nI think that theorem 1 is misrepresented in the paper. The authors write\n> spatial attention is capable of universally approximating any functions defined on local node environment while preserving E(n)-invaraince/equivariance in arbitrary n-dimensional space.\n\nI would interpret this as saying that any invariant/equivariant function from the neighbourhood features and coordinates can be represented by spatial attention. I would thus expect a theorem like:\n\n> Given $N$ neighbours, $C$-dimensional features and $n$ spatial dimensions, let $V \\subset \\mathbb R^{C^N}$ be the subset of the edge features where all edges have distinct values. Then for any function $g : V \\times \\mathbb R^{n^N} \\times \\mathbb R^n \\to \\mathbb R$, that is invariant under permutations and $E(n)$, there is a spatial attention layer that matches $g$ arbitrarily well on all inputs.\n\nHowever, this is not what is shown in the proof. The proof states that for **one given value of edge features** $\\mathbb R^{C^N}$, and one invariant function of the coordinates $g: \\mathbb R^{n^N} \\times \\mathbb R^n \\to \\mathbb R$, there is a corresponding spatial attention layer. The spatial attention is thus not shown to be universally express any function from the node neighbourhood with node features. Other than the paper suggests, the theorem thus also has no implication about whether the SAKE network can represent any equivariant function of pointclouds with node features. I think the claim that the layer is a universal approximator is thus incorrect.\n\nLike theorem 1, I think that remark 3 wrongly claims universality.\n\n### Remark 2\nThe DPS method is not so clearly defined, but for equation $\\phi^{DPS}(v; \\lambda_k, \\lambda_q)=(\\phi^{SA}(v; \\lambda_+)^2 - \\phi^{SA}(v; \\lambda_-)^2)/2$ (Sec 8.2) to make sense, it seems like:\n\n$\\phi^{DPS}(v; \\lambda_k \\lambda_q)_i = \\sum\\_{uu'}\\lambda\\_{q,i,u}\\lambda\\_{k,i,u'} \\langle f\\_u , f\\_{u'} \\rangle$\n\nwhere $f_u=f(e_{uv})\\in \\mathbb R^n$ and $\\lambda\\_{q,i,u}=\\lambda\\_{q,i}(h\\_{e\\_{uv}}) \\in \\mathbb R$ is the intended defintion.\n\nIf so, the claim in Sec 8.2 that with orthogonal edges $\\phi^{DPS}(v)=0$ is clearly false. For example, if there is one neighbour with edge $e$, $N_\\lambda=1, \\lambda=1, f(e)=e, \\mu(c)=c$, then it looks like $\\phi^{DPS}(v)=||e||^2$. Thus the claim that the spatial attention is more expressive than DPS, doesn't appear correct.\n\nA forteriori, it appears to me that a construction similar to the proof in Sec 8.1 can be used for $\\phi^{DPS}$ to show that it is just as universal as the spatial attention method as defined in Theorem 1. Appropriately choosing $\\lambda$, one can construct a list of all inner products $\\langle x_v - x_u , x_v - x_{u'} \\rangle$, from which all pairwise distances between the neighbours can be constructed via $||x_u - x_{u'}||^2=||x_v - x_{u}||^2+||x_v - x_{u'}||^2 - 2 \\langle x_v- x_u, x_v - x_{u'} \\rangle$. Subsequently, the same argument  as made in sec 8.1 can be made to show that this construction is universal as defined in theorem 1.\n\nAlso with a definition of $\\phi^{DSA}$ that's more similar to conventional attention, like $\\phi^{DPS}(v; \\lambda_k \\lambda_q)_{uu'} = \\sum\\_{i}\\lambda\\_{q,i,u}\\lambda\\_{k,i,u'} \\langle f\\_u , f\\_{u'} \\rangle$, choosing $\\lambda=1,f\\_u=e\\_{vu}$, it is universal as defined in theorem 1.\n\nIn conclusion, regardless of how exactly $\\phi^{DSA}$ is defined, it doesn't appear any less expressive than the the proposed spatial attention method. I thus don't think Remark 2 is correct.\n\n## Typo\n- p2, first paragraph \"invaraince\"\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "K042JYc0Sj4", "reviewer": "Reviewer_g5jj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "", "summary_of_the_paper": "This paper introduces spatial attention into E(n)-equivariant neural networks to describe the node environments for the modeling of many-body systems. The architecture builts on the EGNN introduced by Satorras et al. and introduces spatial attention to improve the performance in several  invariant and equivariant prediction tasks.", "summary_of_the_review": "Overall, the paper introduces a novel equivariant spatial attention layer, extending the EGNN architecture. The authors demonstrated improved performance in a broad range of tasks. However, the novelty of the spatial attention layer seems marginal. More detailed ablation study is needed to clarify the usefulness of the introduced spatial attention layer. ", "strength_and_weaknesses": "Strength:\n- The introduced spatial attention improves the expressiveness of the EGNN and preserves the E(n)-Equivariance of the original model.\n- The paper conducts extensive experiments to demonstrate the improved performance in invariant property prediction, equivariant property prediction, and density estimation in equivariant normalizing flows.\n- Model inference is significantly faster than architectures based on spherical harmonics.\n- Detailed ablation study is performed to compare SAKE and EGNN. \n\nWeakness:\n\n- The ablation results are confusing. EGNN is significantly worse than SAKE (for Aspirin, SAKE has MAE of 9.9 and EGNN has an MAE of 298.05). It looks like the EGNN is not properly trained. Further, the authors need to provide more detailed descriptions for the ablation study. What does “No update” mean? \n- What is the equivariant function f in Eq. 6? I cannot find a detailed explanation of the function f in the paper. Is it an identity function or other equivariant function?\n- The Proof 8.1 only considers the special case of f = I. Will Theorem 1 still hold if f is not the identity function?\n- Some inference time comparisons are incomplete. In Table 1, the inference time of SchNet, sGDML, PaiNN is missing. In Table 5, the inference time of GNN, EGNN is also missing. How much overhead does the attention mechanism introduce compared with GNN and EGNN? \n- In QM9, the model performs less competitively in some intensive property prediction tasks. The authors hypothesize that it can be mitigated by size-invariant pooling functions. This should be a straightforward improvement by using a mean pooling rather than sum pooling. I suggest providing additional results using the mean pooling.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "37rk_IefnrS", "reviewer": "Reviewer_u1rJ", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "", "summary_of_the_paper": "Some of the earlier successful methods for physical modelling (like Schutt et al., 2017 and Bartok et. al, 2013) are unable to fully account for the geometry of the system. To model particle arrangements, they rely on using radial information (which captures inter-particle distances). It should be straightforward to see that such methods can't fully capture the geometry of node environments. For instance, one might have particular arrangements which are quite different relative to each other, even when the distance from a node remains unaltered.  More recent methods use a truncated series (resembling a multipole expansion) that uses spherical harmonics to encode higher-order interactions between particles -- and then incorporate it in a message-passing formulation. Such methods have proven to be quite powerful and have been developing quite rapidly over the past few years. However, they do tend to be somewhat expensive computationally. This paper attempts to address this by proposing a method (SAKE) that relies on a simple attention-based formulation that uses linear combinations of edge vectors while being equivariant. More specifically, SAKE uses the _norm_ of linear combinations of edge features to specify the node environment. The method and some of the contributions are summarized below: \n\nThe background on equivariance and a high-level description of the message-passing formalism is provided in sections 2.1 and 2.2. The problem is formulated as implementing a function f_\\theta: \\mathcal{X} x \\mathcal{H} \\mapsto \\mathcal{X} x \\mathcal{H}. The space \\mathcal{X} x \\mathcal{H} forms a joint space of n-dimensional coordinates, and D-dimensional semantic embeddings. The approach that is most similar in terms of applicability and import is that of Satorras et al. The framework for spatial attention is defined in section 4. It uses an equivariant function on the edge embeddings. These embeddings are used to generate attention weights \\lambda. The authors suggest taking N_{\\lambda} such linear combinations, computing their norms and then further concatenating them (multiple heads), and mapping it to the \\mathcal{H} space (by function \\mu). Both the \\lambda and \\mu functions are modelled via MLPs. This simple spatial attention module is E(n) invariant. It is also universal for E(n) invariant functions, as encapsulated in theorem 1. The edge embedding is computed by borrowing ideas from Schutt et. al. and Satorras et al. (equation 7). In addition, there are two more attention modules to promote anisotropy (described on page 5). Finally, there is a similar term as in Satorras et al. that models updating a fictitious velocity (eq 10 and 11). \n\nSAKE is equivariant on the geometric space and invariant on the embedding space. The experiments are divided into three parts. For invariant modelling, the authors consider standard tasks on MD17, ISO17, and QM9.  On MD17, the approach is competitive, and beats all but one of the baselines, but with a significant time gain. While only three baselines are considered for ISO17, SAKE is reported to be better than the competition, while again reportedly faster. Then there are two equivariant tasks (charged N-body and walking motion) where a similar trend is reported. Finally, the method is used to describe a normalizing flow method and is shown to be competitive with the approaches of Kohler et al. and Satorras et al. 2022. \n", "summary_of_the_review": "See the above summary. ", "strength_and_weaknesses": "Strengths: \n- The paper addresses an important problem. The method is well-motivated and quite simple and elegant.\n- Experimental results are quite encouraging. The presented approach is competitive with various methods that were proposed recently -- although it is not the best across the board -- but has a significant advantage in terms of run-time. \n- I think the formulation also suggests approaches to optimize for the architecture, which can improve the experimental results further. \n\nWeaknesses:\n- I think the writing of the paper is a little bit clunky. \n- The novelty of the method is somewhat limited. But I think the experimental results make the case for it as far as I am concerned.\n- The code is not properly anonymized. \n\nMinor comments:\n- At the beginning of page 3, \\mathcal{H} is D dimensional, but at the beginning of section 4, it is written to be C dimensional. Might want to remove this inconsistency. \n- Please re-write the abstract. For instance, the first line is too long and seems to have some punctuation missing. \n\nSome typos\n- Page 2, line 2: invaraince --> invariance.\n- Page 2, line 2 in paragraph 2: \"keep track\" --> keeps track. \n- Page 3, first paragraph of section 3: \"Compared to SAKE, these models are not capable in equivariant modeling\" --> \"Compared to SAKE, these models are not capable of equivariant modeling\".\n- Page 4, below equation 6: \"function that operate on the edge vector\" --> \"function that operates on the edge vector\"\n- Page 4, below remark 1: \"With all local degrees of freedom are incorporated in spatial attention\" --> \"With all local degrees of freedom incorporated in spatial attention,\"\n- Page 5, above section 6: \"In relation to spherical harmonics-based models.\" --> \"Relation to spherical harmonics-based models.\"\n- Page 7, section 6.2: \"model can predict the evolving of a physical system sufficiently long after initial conditions\" --> \"model can predict the evolution of a physical system sufficiently long after initial conditions\"\n- Page 8, next to table 5: \"invaraince\" --> invariance. \n- Page 8, next to table 5: \"as the composing\" --> \"as the composition\"\n- Page 8: \"Note that this reduce\" --> \"Note that this reduces\"", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster