Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs
Authors: Yi-Lun Liao
Tess Smidt
Venue: NeurIPS 2022 Submitted (2022)
Keywords: equivariant neural networks
graph neural networks
computational physics
transformer networks

Abstract: 3D-related inductive biases like translational invariance and rotational equivariance are indispensable to graph neural networks operating on 3D atomistic graphs such as molecules. Inspired by the success of Transformers in various domains, we study how to incorporate these inductive biases into Transformers. In this paper, we present Equiformer, a graph neural network leveraging the strength of Transformer architectures and incorporating SE(3)/E(3)-equivariant features based on irreducible representations (irreps). Irreps features encode equivariant information in channel dimensions without complicating graph structures. The simplicity enables us to directly incorporate them by replacing original operations with equivariant counterparts. Moreover, to better adapt Transformers to 3D graphs, we propose a novel equivariant graph attention, which considers both content and geometric information such as relative position contained in irreps features. To improve expressivity of the attention, we replace dot product attention with multi-layer perceptron attention and include non-linear message passing. We benchmark Equiformer on two quantum properties prediction datasets, QM9 and OC20. For QM9, among models trained with the same data partition, Equiformer achieves best results on 11 out of 12 regression tasks. For OC20, under the same setting of training with IS2RE data only, Equiformer improves upon state-of-the-art models.

Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs
Anonymous Author(s)
Affiliation
Address
email
Abstract
3D-related inductive biases like translational invariance and rotational equivariance are indispensable to graph neural networks operating 3D atomistic graphs such as molecules. Inspired by the success of Transformers we study how to incorporate these inductive biases into Transformers. In this paper, we present Equiformer, graph neural network leveraging the strength of Transformer architectures and incorporating (E based on irreducible representations (irreps). Irreps features encode equivariant information in channel dimensions without complicating graph structures. The simplicity enables us to directly incorporate them by replacing original operations with equivariant counterparts. Moreover, to better adapt Transformers 3D graphs, we propose a novel equivariant graph attention, which considers both content and geometric information such relative position contained in irreps features. To improve expressivity of the attention, we replace dot product attention with multi- layer perceptron attention and include non-linear message passing. We benchmark Equiformer on two quantum properties prediction datasets, QM9 and OC20. For QM9, among models trained with the same data partition, Equiformer achieves best results on 11 out of 12 regression tasks. For OC20, under the setting of training with IS2RE data and optionally IS2RS data, Equiformer improves upon state-of- oate-of-thee ff-the-an art models.
3
various domains,
5
features
11
13
15
18
17
18
19
Introduction
20
Machine learned models accelerate the prediction of quantum properties of atomistic systems
like molecules by learning approximations of ab initio calculations [29, 87, 37, 25, 4, 10, 49, 76, 69, 54, 51J. In partieular, graph neural networks (GNNs) have gained increasing popularity due their performance. By modeling atomistic systems as graphs, GNNs naturally treat the set-like nature of 25 collections of atoms, encode the interaction between atoms in node features and update the features
22
23
24
28 by passing messages between nodes. One factor contributing to the success of neural networks is the
27 ability incorporate inductive biases that exploit the symmetry data. Take convolutional neural
28 networks (CNNs) for 2D images as an example: Patterns in images should be recognized regardless
29 of their positions, which motivates the inductive bias of translational equivariance. for atomistic
30 graphs, where each atom has its coordinate in 3D Euclidean space, consider inductive biases related 3D Euclidean group E(3), which include equivariance to 3D translation, 3D rotation, and inversion. Concretely, some properties like energy of an atomistic system should be constant regardless of
31
33 how we shift the system; others like force should be rotated accordingly if we rotate the system. To incorporate these inductive biases, equivariant and invariant neural networks have been proposed. The 35 former leverages geometric tensors like vectors for equivariant node features [71, 79, 43, 23, 4, 5, 51], and the latter augments graphs with invariant information such distances and angles extracted from
37 3D graphs [63, 26, 25, 48, 67, 42].
Submitted to 36th Conference on Neural Information Processing Systems (NeurIPS 2022). Do not distribute.

---

38
A parallel line of research focuses on applying Transformer networks [77] to other domains like computer vision [9. 16, 72] and graph [18, 44, 84, 65] and has demonstrated widespread success. However, as Transformers were developed for sequence data [15, 3. 7]. it is crucial to incorporate domain-related inductive biases. For example, Vision Transformer [16] shows that adopting a pure Transformer to image classification cannot generalize well and achieves worse results than CNNs when trained on only ImageNet [60] since it lacks inductive biases like translational invariance. Note that ImageNet contains over 1. 1.28M images and the size is already larger than that of many quantum properties prediction datasets [59, 56, 10). Therefore, this highlights the necessity of including correct inductive biases when applying Transformers to the domain of 3D atomistic graphs.
39
40
41
42
43
44
45
47
In this work, we present Equiformer, an equivariant graph neural network utilizing (3) equivariant features built from irreducible representations (irreps) and equivariant attention mecha- nisms to combine the 3D-related inductive bias with the strength of Transformer. Irreps features encode equivariant information in channel dimensions without complicating graph structures. The simplicity enables us to directly incorporate them into Transformers through replacing original operations with equivariant counterparts and introducing an additional equivariant operation called tensor product. Moreover, we propose a novel equivariant graph attention, which considers both content and geometric information such as relative position. Equivariant graph attention improves upon typical attention in Transformers by replacing dot product attention with theoretically stronger multi-layer perceptron attention and including non-linear message passing. With these innovations, Equiformer demonstrates the possibility of generalizing Transformers to 3D atomistic graphs and achieves competitive results two quantum properties prediction datasets, QM9 [59, 56] and OC20 [10]. For QM9, compared to models trained with the same data partition, Equiformer achieves the best results on 11 out of 12 regression tasks. For OC20, under the setting of training with IS2RE data and optionally IS2RS data, Equiformer improves upon state-of-the- models.
48
49
50
51
53
54
55
57
58
59
61
62 2
Related Works
63
Here, we focus on equivariant neural networks and discuss other works in Sec. B in appendix.
64
Equivariant GNNs Equivariant neural networks [71, 43, 79, 23, 50, 73, 4, 38, 64, 62, 76, 5, 70, 46, 51] operate geometric tensors like type-I. vectors to achieve equivariance. The central idea
65
66 is to use functions of geometry built from spherical harmonics and irreps features to achieve 3D rotational and translational equivariance as proposed in Tensor Field Network (TFN) [71], which generalizes 2D counterparts [81, 12, 13] to 3D Euclidean space [71, 79, 43). Previous works differ 69 in equivariant operations used in their networks. TFN [71] and NequlP [4] use graph convolution with linear messages, with the latter utilizing extra equivariant gate activations [79]. SEGNN [5] introduces non-linear messages [29, 61] for irreps features, and the non-linean messages use the same gate activation and improve upon linear messages. (E(3)-Transforme [23] adopts an equivariant version of dot product (DP) attention [77, 39] with linear messages, and the attention can support vectors of any degree (type) L. Subsequent works on equivariant Transformers [70, 46) follow the practice of DP attention and linear messages but use more specialized architectures considering only type-0 and type-1 vectors. The proposed Equiformer incorporates all the advantages through combining MLP attention with non-linear messages and supporting vectors of any type. Compared to [23], the proposed combination of MLP attention and non-linear messages is more expressive than pure linear or non-linear messages and
67
68
70
71
72
73
74
75
76
77
78
TFN [71], NequlP [4], SEGNN [5], and
79
pure MLP dot product attention. Compared to other equivariant Transformers [70, 46), in addition
81 to being more expressive, the proposed attention mechanism can support vectors of higher degrees and involve higher order tensor product interactions, which can lead to better performance [4, 5].
82
3
Background
83
84
3.1
E(3 Equivariance
85
Atomistic systems are often described using coordinate systems. For 3D Euclidean space, we can freely choose coordinate systems and change between them via the symmetries of 3D space: 3D translation, rotation and inversion ( (F)) The groups of 3D translation, rotation and inversion form Euclidean group E(3), with the first two forming (3), the second being SO(3), and the
86
87
88
89 last two forming (33). The laws of physics are invariant to the choice of coordinate systems and
90 therefore properties of atomistic systems are equivariant, e.g., when we rotate our coordinate system,
91 quantities like energy remain the same while others like force rotate accordingly. Formally, a function mapping between vector spaces X and Y is equivariant to a group of transformation G if for any

---

input I E X. output E Y and group element 9 E G, we have f(Dx(g)z) = (y())f(z), where (9) and (g) are transformation matrices parametrized by g in X and Y.
93
94
85
Incorporating equivariance into neural networks as inductive biases is crucial as this enables generaliz- ing unseen data in a predictable manner. For example, 2D convolution f is equivariant the group of 2D translation, and thus, CNNs can identify patterns at any location even if they have never seen the patterns at that specific location before. For 3D atomistic graphs, we consider the group of E33). Features and learnable functions should be (33)-equivarias
87
98
99
to geometric transformation acting on position F.. In this work, following previous works [71, 43, 79] implemented in e3nn [28], we achieve S (33))EE(3 equivariance by using equivariant features based on vector spaces of irreducible representations and equivariant operations like tensor product for learnable functions.
100
101
102
103
3.2
Irreducible Representations
104
group representation [17, 85] defines the transformation matrices (g) of group elements g that act on vector space X. For 3D Euclidean group E(3), two examples of vector spaces with different transformation matrices are scalars and Euclidean vectors in R i.e., vectors change with rotation while scalars do not. To address translation symmetry, we simply operate on relative positions. Below we focus our discussion on 0(3). The transformation matrices of rotation and inversion are separable and commute, and we first discuss irreducible representations of SO(3).
105
106
107
108
109
110
Any group representation of SO(3) on given vector space can be decomposed into concatenation of provably smallest transformation matrices called irreducible representations (irreps). Specifically, for group element g E SO(3). there are (2L+ 11-++-)--y-(-++++ 1) irreps matrices (g) called Wigner-D matrices acting on (2L + 1) 1)-dimensional vector spaces, where degree L is a non-negative integer. L can be interpreted as an angular frequency and determines how quickly vectors change when rotating coordinate systems. (9) of different L act on independent vector spaces, Vectors transformed by
111
112
113
114
115
116
(g) are type-L vectors, with scalars and Euclidean vectors being type-0 and type- vectors. is common to index elements of type-L vectors with an index m called order, where - L S m S L.
117
118
The group of inversion Z, only has two elements, identity and inversion, and two irreps, even e and odd o. Vectors transformed by irrep e do not change sign under inversion while those by irrep o do. We create irreps of (3)) by simply multiplying those of SO(3) and Zg and we introduce parity p to type-L vectors to denote how they transform under inversion. Therefore, type-L vectors in SO(3) are extended to type-(L, p) vectors in O()). where p is or o. In the following, we use Lype-L vectors for the ease of discussion, but we can generalize to type-(L, vectors, unless otherwise stated.
119
120
121
122
123
Irreps Features. We concatenate multiple type-L vectors to form SE(3) eqquivaraat irreps features. Concretely, irreps feature f has CL type-L vectors, where 0 < L < Limax and Cz is the number of channels for Lype-L vectors. We index irreps features f by channel c, degree L, and order and denote as f!t)). Different channels of type-L vectors are parametrized by different weights but are transformed with the same Wigner-D matrix Dg(g). Regular scalar features correspond to including only type-0 vectors. This can generalize to E33) by including inversion and extending L to (L, p).
125
126
127
129
Spherical Harmonies. Euclidean vectors in g can be projected into type-L vectors (L) by using spherical harmonics (SH) y(L); ((L) (L)
130
SH are -quu-vaauva with Dr(g))
131
( Dr(g) ). SH of relative position ris generates the first set of irreps features. Equivariant Di
132
133
information propagates to other irreps features through equivariant operations like the tensor product.
134
3.3
Tensor Product
135
We use tensor products to interact different type-L vectors and first discuss the tensor product for S3(3). The tensor product denoted as @ uses Clebsch-Gordan coefficients combine type-L1 vector ((i) and type-L2 vector g gfaaa) and produces type-Ly vector (4a) as follows:
L Σ
La
[La)) = (1( 60 (La) mS
(la.ma) (brm)))aamm)) gma
(L))
m
m2=-L2 L2
138 where my denotes order and refers to the my-th element of f(L)), Mebscch-Gordal coefficients c(La.ma) (2a are non-zero only when [ - Lel < Ls S 1L1 + Lzl and thus restrict output vectors to be of certain types. For efficiency, we discard vectors with L > Lmazz where Lmaa is
139
3

---

Input Graph
Input 3D Graph
Atom
Type
Constant One
1g
Lineer
Linse
Einzedding
Layer Nomm
Equivarintt
Gouph Altention
Stalass
Imees Risture.
Lesky ReL
Gaie
Layer Norm
(c)
Embedkling (Atom + Edge-Degree)
Feed Forwani
Network
Layer Norm
I Lince
Rechape
Feed Forward Network:
Cole
Ouiiut
iince
Head
(4) Equiformer Architecture
(b) Equivariant Graph Attention
(d) Feed Forward Network
Figure 1: Architecture of Equiformer. We embed input 3D graphs with atom and edge-degree embeddings and process them with Transformer blocks, consisting of equivariant graph attention and feed forward networks. In this figure, "s" denotes multiplication, "g"" denotes addition, and "DTP" stands for depth-wise tensor product. E within circle denotes summation over all neighbors. Gray cells indicate intermediate irreps features.
hypereparameter, to prevent vectors of increasingly higher dimensions. The tensor product is an DL, (g'))a) for g' E SO(3).
cquivariant operation, with (Dz, f())) @ (DL,
142
We call each distinct non-trivial combination of L1 I.2 L,3 path. Each path is independently equivariant, and we can assign one learnable weight to each path in tensor products, which is similar to typical linear layers. We generalize Eq. to irreps features and include multiple channels of vectors of different types through iterating over all paths associated with channels of vectors. In this way, weights are indexed by (c,,4, C2, 12 c32 la), where e₁ is the ci-th channel of type-l1 vector in input irreps feature. We use @w to represent tensor product with weights w. Weights can be conditioned on quantities like relative distances. Please refer Sec. A.4 appendix for discussion on inversion in tensor products and Sec. D.I and E. for additional results of including inversion.
144
146
147
148
149
150
Equiformer
We incorporate -eeeuvvrraaa irreps features into Transformers [77] and use equivariant operations. better adapt Transformers to 3D graph structures, we propose equivariant graph attention. The overall architecture of Equiformer is illustrated in Fig. 1.
153
4.1 Equivariant Operations for Irreps Features
156
Here we discuss equivariant operations used in Equiformer that serve as building blocks for equivariant graph attention and other modules. They include the equivariant version of the original operations in Transformers and the depth-wise tensor product as illustrated in Fig. 2.
157
158
159
Linear layers are generalized to irreps features by transforming different type- vectors separately. Specifically, we apply separate linear operations to each group of type-I. vectors. We remove bias terms for non-scalar features with including biases for type-L vectors with L > 0 can break cquivariance.
Linear.
161
as biases do not depend on inputs, and therefore,
165
Layer Normalization. Transformers adopt layer normalization (LN) [2] to stabilize training. Given input & XC with N being the number of nodes and the number of channels, LN calculates the linear transformation of normalized input as LN(x)
164
o7 3, where PC,0c E 2-pC ac
R NxI mean and standard deviation of input x along the channel dimension, 7, RxC learnable parameters, and o denotes element-wise product. By viewing standard deviation as the root mean square value (RMS) of L2-norm of type-L vectors, LN can be generalized irreps features. Specifically, given input XXCxX (CXx(2++1 of type-L vectors, the output is LN(x)
167

---

Input ireess Teatures
Oulput irreps festures
Type Type Type-0 vectors vectons sectoos
Satunc
Divide by RMS
SILU
neaL
Sigmoid
by RMS
(a) Linear
(b) Layer Normalization
(c) Gate
(d) Depth-wise Tensor Product
Figure 2: Equivariant operations used in Equiformer. (a) Each gray line between input and
output irreps features contains one learnable weight. (b) "RMS" denotes the root mean square value along the channel dimension. For simplicity, we have removed multiplying by here. (c) Gate layers are equivariant activation functions where non-linearly transformed scalars are used to gate non-scalar irreps features, (d) The left two irreps features correspond to two input irreps features, and the rightmost one the output irreps feature. The two gray lines connecting two vectors in the input inreps features and one vector in the output irreps feature form a path and contain one learnable weight. An alternative visualization of depth-wiss wisee tensor products can be found in Fig. in appendix. We only show SE(3) -equuiuaarari operations here, and they can be directly generalized to F(3)) features.
RMSc (2) o7, where (2)) RNxxCx1 calculates the L2-norm of each type-L vectors in æ, and RMS (norm(z)) Nx calculates the RMS of L2- norm with mean taken along the channel dimension. We remove mean and biases for type-L vectors with L * 0 following linear layers.
170
171
172
173
We use the gate activation [79] for equivariant activation function as shown in Fig. 2(c). Typical activation functions are applied to type-0 vectors. For vectors of higher L, we multiply them with non-linearly transformed type-0 vectors for equivariance. Specifically, given input x containing non-scalar CL type-L vectors with 0 < S Lmax and (Co + Ei=1 Cz) type-O vectors, we apply SiLU [19, 55] to the first Co type-O vectors and sigmoid function to the other Ehmas Cy type-0 vectors to obtain non-linear weights and multiply each type-L L vector with corresponding non-linear weights. After the gate activation, the number of channels for type-O vectors is reduced to Co-
Gate.
174
175
176
178
180
Depth-wise Tensor Product.
The tensor product defines interaction between vectors of different L. the depth-wise tensor product (DTP), which restricts one type- I
To improve its efficiency, vector in output irreps features depends only on one type-I! vector in input irreps features, where L, can be equal to or different from L'. This similar to depth-wise convolution [34], where one output channel depends on only one input channel. Weights w in the DTP can be input-independent or DTP conditioned on relative distances, and the DTP between two tensors æ and y is denoted &DTP y.
182
183
184
185
186
4.2
Equivariant Graph Attention
187
Self attennienio [77, 78, 23, 39, 84, 6] transforms features sent from one spatial location to another with input-dependent weights. We use the notion from Transformers [77] and message passing networks [29, 61, 62, 5] and define message maj sent from node to node i as follows:
188
189
mij aij Uij
190
where attention weights aij depend on features on node and its neighbors (t) and values Usj transformed with input-ineependen weights. In Transformers and Graph Attention Networks (GAT) [78, 6], Uyy depends only on node j. In message passing networks [29, 61, 62, 5], Viy depends on features nodes and with constant aij. The proposed equivariant graph attention adopts tensor products to incorporate content and geometric information and utilizes multi-layer perceptron attention for asj and non-linear message passing for Viy as illustrated in Fig. 1(b).
192
194
195
196
Incorporating Content and Geometric Information. Given features and I, on target node i and source node j. we combine the two features with two linear layers obtain initial message Tzj Linear (I:)) + Linear,re(sr Tiy is passed to a DTP layer and a linear layer to consider geometric information like relative position contained different type-L L vectors in irreps features:
197
198
199
d/1 zij 2 DTP
@ (Fy
SH(F and fiz Linear(x(.))
33)
5

---

200
where z'y is the tensor product of Fay and spherical harmonics embeddings (SH) of relative position Fij. with weights parametrized by ||, Jij considers semantic and geometric features on source and target nodes in linear manner and is used to derive attention weights and non-linear messages.
202
203
Multi-Layer Perceptron Attention. Attention weights az capture how each node interacts with neighboring nodes. Gij are invariant to geometric transformation [23]. and therefore. we only
204
use type-0 vectors (scalars) of message fúj denoted as f( for attention. Note that f encodes directional information, as they are generated by tensor products of type-L vectors with L 2 O.
205
208
207
Inspired by GATv2 [6]. we adopts multi-layer perceptron attention (MLPA) instead of dot product attention (DPA) used in Transformers [77, 39j. In contrast dot product, MLPs are universal
208
209
approximators [33, 32, 14] and can theoretically capture any attention patterns [6]. Similar to
GAT [78, 6). given (( uses one leaky ReLU layer and one linear layer for aiz:
210
exp(215)
žtj a. ) and azj softmax, ΣKEN(O) exp(zu)
(4)
(0)) where a is a learnable vectors of the same dimension as fij and 2ij is a single scalar. The output of attention is the sum of value f4 multipled by corresponding dzy over all neighboring nodes j E (i)), where Viy can be obtained by linear or non-linear transformations of fy as discussed below.
211
213
214
Non-Linear Message Passing. Values Biy are features sent from one node to another, transformed
with input-independent weights. We first split fiş into f.f/)) and s(j)) where the former consists of type-L vectors with 0 S L < Lmar and the latter consists of scalars only. Then, we perform
216
(LL)
217
non-linear transformation fij obtain non-linear message:
Mij = Gate( (1()))
and Vij = Linear((uus 8 DTP SH(Fi)))
(5)
218
We apply gate activation to fig obtain pij. We use one DTP and linear layer to enable interaction between non-linear type-L vectors, which is similar to how transform Ziy into fij- Weights
(LL)
here are input-independent. We can also use p(() directly as vij for linear messages.
220
Head Attention. Following Transformers [77], we can perform h parallel equivariant graph attention functions given fiz. The h different outputs are concatenated and projected with a linear layer, resulting in the final output y illustrated in Fig. 1(b). Note that parallelizing attention functions and concatenating can be implemented with "Reshape".
222
224
225
4.3
Overall Architecture
226
For completeness, we discuss other modules in Equiformer here.
227
This module consists of atom embedding and edge-degree embedding. For the former, we use a linear layer to transform one-hot encoding of atom species. For the latter, depicted in the right branch in Fig. Icc), we first transform a constant one vector into messages encoding local geometry with two linear layers and one intermediate DTP layer and then sum aggregation to encode degree information [83, 65]. The DTP layer has the same form as that in Eq. 3. We scale the aggregated features by dividing with the squared root of average degrees in training sets that standard deviation of aggregated features would be close to 1. The two embeddings are summed to produce final embeddings of input 3D graphs.
Embedding.
229
231
233
Radial Basis and Radial Function. Relative distances || parametrize weights in some DTP layers. To reflect subtle changes in || we represent distances with Gaussian radial basis with learnable mean and standard deviation [63, 67, 42, 65] radial Bessel basis [26, 25]. We transform radial basis with a learnable radial function to generate weights for those DTP layers [63, 23, 4]. The function consists of two MLPs with layer normalization [2] and SiLU [19, 55] and final linear layer.
236
238
239
Feed Forward Network. Similar to Transformers, use two equivariant linear layers and an intermediate gate activation for the feed forward networks in Equiformer.
241
242
Output Head. The last feed forward network transforms features on each node into a scalar. We perform sum aggregation over all nodes to predict scalar quantities like energy. Similar to edge-degree embedding, we divide the aggregated scalars with the squared root of average numbers of atoms.
244

---

Task
<3ROOO meV
FLUMM meV
11
C2 callmol K
H
R²
Co meVV
ZPVE
Methods
Units
bohr'
mn
D
meV
meV
bokar
me
meV
NMP [29]
092
43
38
.180
SchNet [63]'
235
63
41
34
.033
.03
14
14
.773
19
14
1.70
Cormorant [1]
085
34
38
0038
026
20
.611
21
22
2.03
LieCony [21]
084
25
8001
DimeNett+ [25]
044
33
25
20
030
023
8
7
.31
6
1.21
TFN [7L]*
223
Transorra (23)
.122
-3
35
33
051
054
BONN [62]
0771
-88
29
25
0129
033
12
12
106
12
11
1.55
SphereNer [48]*
046
23
022
292
SEGNN [5]
060
42
24
21
0023
031
15
16
666
13
15
1.62
RQGAT [46]
063
44
26
22
014
027
12
13
257
13
13
1,50
Equiformer
056
10
Table 1: MAE results on QM9 testing set. f denotes using different training, validation, testing data partitions as mentioned SEGNN [5]. denotes results from STanS-Tranterrmer [23].
Energy (AV
EWT (5)
Cat
Average
Average
SchNot 163]
0,6465
0.7074
0.6475
0.6626
0.6660
2.96
2.12
3.03
2.38
2.65
DimncNet+ [25]
0.5636
0.7127
0,5612
0.6492
0.6217
4.25
2.48
4.40
2.56
3.42
T [42]
0,5561
0.7342
0.5659
0.6964
0.6382
4.51
2.24
4.37
2.38
3.38
SphereNci [48]
0.5632
0.6024
4.59
(SIEGNN [5]
0.5497
0.6851
0.5519
0.6102
0.5992
4.99
2.50
4.71
2.88
3.77
SEGNN [S]
0,5310
0,6432
0.5341
0.5777
0.5715
5.32
2,80
4,89
3,09
4.03
Equiformer
0.5088
Table 2: Results on OC20 IS2RE yalidation set. i denotes results reported by SphereNet [48].
Energy (
EWT (%6
Mothods
ID)
000 Ads
001 Cat
OOD Both
Average
11)
OD1 Ads
OOD Cat
(OOI) Both
Average
CGCNN [82]
0.6219
0.8511
0.7509
1.93
SchNet [63]
0.6387
0.7342
0.6615
0.7037
0.6846
2.96
2.33
2.94
2.21
2.61
DimeNet-r [25)
0.5621
0.7252
0.5756
0.6613
0.6311
4.25
2.07
4,10
2.41
3.21
SpinConv [67]
0.5583
0.7230
0.5687
0.6738
0.6310
2.26
3,82
2.33
3.12
SphoreNet [48]
0.5625
0.7033
0,5708
0.6378
0.6186
4.47
2.29
4.09
2.41
3.32
SEGGN [5]
0,5327
0.6921
0,5369
0.6790
0.6101
5.37
2.46
4.91
2.63
3.84
Equiformer
0.5057
0.6881
0.5213
0.6301
2.41
2.69
Table 3: Results on OC20 IS2RE testing set.
245 5
Experiment
246
Our implementation based on PyTorch [52] (Modified BSD license), PyG [20] (MIT license), e3nn [28] (MIT license), timm [80] (Apache-2.0 license), and ocp [10] (MIT license).
5.1
QM9
Dataset. The QM9 [59, 56) dataset (CCBY-NC SA 4.0 license) consisting of 134k small molecules, and the goal is to predict their quantum properties such as energy. We follow the data partition used by Cormorant [1], which has 100k, 18k and 13k molecules in training, validation and testing sets. We minimize mean absolute error (MAE) between prediction and normalized ground truth.
250
252
253
Setting. Please refer to Sec. D in appendix for details on architecture and syper-parameters.
We mainly compare with methods trained with the same data partition and summarize the results in Table 1. Equiformer achieves the best results on 11 out of 12 tasks among models trained with same data partition. The comparison to SEGNN [5], which uses irreps features as Equiformer, demonstrates the effectiveness of combining non-lienar messages with MLP attention. Additionally, Equiformer achieves better results for most of tasks when compared to other equivariant Transformers [23, 46), which suggests a better adaption of Transformers to 3D graphs. Besides, the different data partition as denoted by in Table has 10% more molecules in the training set and less data in the testing set, and this can benefit some tasks that are more dependent on data partitions.
Result.
255
257
259
261
5.2
OC20
263
Dataset. The Open Catalyst 2020 (OC20) dataset [10] (Creative Commons Attribution 4.0 License) consists of larger atomic systems, each composed of small molecule called the adsorbate placed on a large slab called catalyst. The average number of atoms system is more than 70, and there are over 50 atom species. The goal is to understand interaction between adsorbates and catalysts through relaxation. An adsorbate first placed on top of catalyst to form initial structure (IS). The positions of atoms are updated with forces calculated by density function theory until the system is
264
265
257
7

---

Energy MAE (eV)
EWT (5)
Methods
ID
OOD Ads
OOD Cat
ODD Both
Ascrage
ID
OOD
OOD
OOD Beth
Average
GNS [30]
0.65
0.55
0.59
0.5825
GNS Noisy Nodes [30]
0.47
0.51
0.48
0.46
0,4800
Graphormer (65)
0.4329
0.5850
0.4451
0,5299
0.4980
Equiformer
0,4222
0,5420
0,4231
0,4751
0.1657
7.23
3.77
5.56
Equiformer+ + Noisy Nodes
0.4155
0.4976
0.4165
0.4344
0.4410
7.47
4.64
7.19
4.84
6.04
Table 4: Results on OC20 IS2RE validation set when IS2RS node-Ievel auxiliary task is adopted during training. "GNS" denotes the 50-layer GNS trained without Noisy Nodes data augmentation, and "GNS + Noisy Nodes" denotes the 100-layer GNS trained with Noisy Nodes. "Equiformer + Noisy Nodes" uses data augmentation of interpolating between initial structure and relaxed structure and adding Gaussian noise as described by Noisy Nodes [30].
Energy MAE (eV) L
ENT (%)
Methods
OOD
OOD
Average
OOD
OOD Cat
OOD Both
Average
GNS + Noisy Nodes [30]
0.4219
0.5678
0.4366
0.4651
0.4728
9.12
4.25
8.01
4.64
6.5
Graphormer (65)'
0.3976
0.5719
0:4166
0.5029
0.4722
8.97
3.45
6.1
Equiformer+ + Noisy Nodes
0.4171
0,5479
0,4248
0.4741
0.4660
7,71
3.70
7.15
4.07
5.66
Table 5: Results on OC20 IS2RE testing set when IS2RS node-level auxiliary task is adopted during training. denotes using ensemble of models trained with both IS2RE training and validation sets. In contrast, we use the same single Equiformer model in Table 4, which is trained with only the training set, for evaluation on the testing set.
289
stable and becomes relaxed structure (RS). The energy of RS, or relaxed energy (RE), is correlated with catalyst activity and therefore a metric for understanding their interaction. We focus on the task of initial structure to relaxed energy (IS2RE), which predicts relaxed energy (RE) given an initial structure (IS). There are 460k, 100k and 100k structures in training, validation, and testing sets, respectively. Performance is measured in MAE and energy within threshold (EwT)), the percentage in which predicted energy is within 0.02 eV of ground truth energy. In validation and testing sets, there are four sub-splits containing in-distribution adsorbates and catalysts (ID), out-of-distribution adsorbates (OOD-Ads), out-of-distribution catalysts (OOD-Cat), and out-of-distribution adsorbates and catalysts (OOD-Both).
271
272
273
274
275
276
277
278
Setting.
We consider two training settings based on whether a node-level auxiliary task [30] is adopted. In the first setting. we minimize MAE between predicted energy and ground truth energy without any node-level auxiliary task. In the second setting, we incorporate the task of initial structure to relaxed structure (IS2RS) as node-level auxiliary task [30]. In addition to predicting energy, we predict node-wise vectors indicating how each atom moves from initial structure to relaxed structure. Please refer to Sec. E in appendix for details on Equiformer architecture and hyper-paramedte parameters.
279
280
281
282
283
284
IS2RE Result without Node-Level Auxiliary Task. We summarize the results under the first setting in Table 2 and Table 3. Compared with state-of-the-arr models like SEGNN [5] and SphereNet [48], Equiformer consistently achieves the lowest MAE for all the four sub-splits in validation and testing sets. Note that energy within threshold (EwT) considers only the percentage of predictions close enough to ground truth and the distribution of errors, and therefore improvement
286
286
289
in average errors (MAE) would not necessarily reflect that in error distributions (EwT). Similar phenomena can be observed in Table 3. where for "OOD Both" sub-split, SphereNet [48] achieves lower MAE yet lower EwT than SEGNN [5]. We also note that models in Table 2 and 3 are trained by minimizing MAE and therefore comparing MAE in validation and testing sets could mitigate the discrepancy between training objectives and evaluation metrics and that OC20 leaderboard ranks the relative performance of models mainly according to MAE.
290
292
293
294
295
IS2RE Result with IS2RS Node-Level Auxiliary Task. We report the results on validation and testing sets under the second setting in Table 4 and Table 5. As of May 20, 2022, Equiformer achieves the best results on IS2RE task when only IS2RE and IS2RS data are used. We note that the proposed Equiformer in Table 5 achieves competitive results even with much less computation. Specifically, training "Equiformer + Noisy Nodes" takes about 24 GPU-days when A6000 GPUs are
296
297
298
299
300
used. The training time of "GNS + Noisy Nodes" [30] is 56 TPU-days. "Graphormer" [65] uses ensemble of 31 models and requires 372 GPU-days to train all models when A100 GPUs are used.
301
The comparison to GNS demonstrates the improvement from invariant message passing networks to equivariant Transformers. Compared to Graphormer [65], Equiformer demonstrates the effectiveness of cquivariant features and the proposed equivariant graph attention. Note that Equiformer, with
303
8

---

Methods
Non-lincar message passing
MLP attenton
Dot product attention
Task
Aa
C.
Index
bohrr
cal/moll
D556
33
17
16
.014
.025
.611
18
17
.015
025
Table 6: Ablation study results on QM9.
Methods
Energy MAE (cV)
Non-linear message passing
MLP attention
Dot product
Index
aitention
00D Ads
OOD Car
OOD Both
Average
0,5088
0.6271
0,5051
0.5545
0,5489
0,5168
0.6308
0,5088
0.5657 0,5692
0.5555
0,5386
0.6322
0,5297
0 0.5689
Methods
EWT (96
Non-lincar message passing
MLP attention
Dot productt attention
00D
OOD
Average
1.88
2.93
4.92
2.98
3.93
4.59
2.82
4.79
3.02
3.81
2.60
4.36
Table 7: Ablation study results on OC20 IS2RE validation set.
18 Transformer blocks, is relatively shallow GNS trained with Noisy Nodes has 100 blocks and Graphormer has 48 Transformer blocks and that deeper networks can typically obtain better results when IS2RS auxiliary task is adopted [30].
306
308
5.3
Ablation Study
309
We conduct ablation studies on the improvements brought by MI.P attention and non-linear messages in the proposed equivariant graph attention. We modify dot product (DP) attention [77, 23] so that it only differs from MLP attention in how attention weights ai are generated from faj- Please refer
310
311
to Sec. C.3 in appendix for details on DP attention. For experiments QM9 and OC20, unless otherwise stated, we follow the hyper-parameters used previous experiments.
313
Result on QM9. The comparison is summarized in Table 6. Non-linear messages improve upon linear messages when MLP attention is used. Similar to what is reported by GATv2 [6]. the improvement of replacing DP attention with MLP attention is not very significant. We conjecture that DP attention with linear operations is expressive enough to capture common attention patterns as the numbers of nighboring nodes and atom species are much smaller than those in OC20. However,
315
316
317
318
319
MLP attention is roughly 7% faster as it directly generates scalar features and attention weights from
320
fiy instead of producing additional key and query irreps features for attention weights.
Result on OC20. We consider the setting of training without IS2RS auxiliary task and use smaller learning rate 1.5 10-4 1 for DP attention as this improves the performance. We summarize the comparison in Table 7. Non-linear messages consistently improve upon linear messages. In contrast to the results on QM9, MLP attention achieves better performance than DP attention. We surmise this is because OC20 contains larger atomistic graphs with more diverse atom species and therefore requires more expressive attention mechanisms.
322
324
326
Conclusion and Broader Impact
327
328
In this work, we propose Equiformer, a graph neural network (GNN) combining the strengths of Transformers and equivariant features based on irreducible representations (irreps). With irreps features, we build upon existing generic GNNs and Transformer networks [77, 16, 84, 45, 47] by incorporating equivariant operations like tensor products. We further propose equivariant graph attention, which incorporates multi-layer perceptron attention and non-linear messages. Experiments on QM9 and OC20 demonstrate both the effectiveness of Equiformer and the advantage of equivariant graph attention over typical dot product attention.
329
331
333
335
The broader impact lies in two aspects. First, Equiformer demonstrates the possibility of adapting Transformers to domains such as physics and chemistry, where data can be represented as 3D atomistic graphs. Second, Equiformer achieves more accurate approximations of quantum properties calculation. We believe there is much more to be gained by hamnsssig these abilities for productive investigation of molecules and materials relevant to application such as energy, electronics, and pharmaceuticals [10], than to be lost by applying these methods for adversarial purposes like creating hazardous chemicals. Additionally, there are still substantial hurdles to go from the identification of useful harmful molecule to its large-scale deployment.
336
337
338
339
341

---

References
343
[1]
Brandon Anderson, Truong-Son Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. In Conference on Neural Information Processing (NeurIPS), 2019.
345
[2] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arxiy preprint arxiv: 064500 2016.
347
[3] Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. In Conference on Neural Information Processing (NeurIPS), 2020.
349
350
351
[4] Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger. Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data- efficient and accurate interatomic potentials. Nature Communications, 13(1) 3()), May 2022.
352
353
354
[5] Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik J Bekkers, and Max Welling, Geometric and physical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations (ICLR), 2022.
355
357
[6] Shaked Brody, Uri Alon, and Eran Yahav. How attentive are graph attention networks? In International Conference on Learning Representations (ICLR), 2022.
359
[7] Tom Brown, Benjamin Mann. Nick Ryder. Melanie Subbiah. Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Bemer, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodci. Language models are few-shot learners. In Conference on Neural Information Processing (NeurIPS), 2020.
361
362
363
364
385
[8] Dan Busbridge, Dane Sherburn, Pictro Cavallo, and Nils Y. Hammerla. Relational graph attention networks. arxiv preprint arxiv: pexiiv:1000 1904.05811, 1, 2019.
366
387
19] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In European Conference Computer Vision (ECCV), 2020.
368
369
370
[10]
Lowik Chanussot : Abhishck Das*, Siddharth Goyal*, Thibaut Lavril Muhammed Shuaibi Morgane Riviere, Kevin Tran, Javier Heras-Domingo, Caleb Ho, Weihua Hu, Aini Palizhati, Anuroop Sriram, Brandon Wood, Junwoong Yoon, Devi Parikh, C. Lawrence Zitnick, and Zachary Ulissi. Open catalyst 2020 (oc20) dataset and community challenges. ACS Catalysis, 2021.
371
372
373
374
[11] Zhengdao Chen, Lisha Li, and Joan Bruna. Supervised community detection with line graph neural networks. In International Conference on Learning Representations (ICIR),, 2019.
375
376
[12] Taco Cohen and Max Welling. Group equivariant convolutional networks. In International Conference on Machine Learning (ICML), 2016.
377
378
[13] Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical CNNs. International Conference on Learning Representations (ICLR), 2018.
379
380
[14] George Cyhenko. Approximation by superpositions of a sigmoidal function, Mathematics of Control, Signals and Systems, 2:303-314, 1989.
381
382
[15] Jacob Devlin, Ming-Wer Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirec- tional transformers for language understanding. arxiv preprint arxiv: 2019.
383
384
[16] Alexey Dosovitskiy, Lucas Beyer. Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakoh Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021.
385
387
388
[17] Mildred S Dresselhaus, Gene Dresselhaus, and Ado Jorio. Group theory. Springer, Berlin, Germany, 2008 edition, March 2007.
389
[18] Vijay Prakash Dwivedi and Xavier Bresson. A generalization of transformer networks to graphs. 2020.
391
[19] Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Sigmoid weighted linear units for neural network function approximation in reinforcement learning. arXiv preprint arx.v:1702.031118 2017.
392
10

---

393
[20] Matthias Fey and Jan E. Lenssen. Fast graph representation learning with PyTorch Geometric. In ICLR Workshop on Representation Learning on Graphs and Manifolds, 2019.
[21]
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neu- ral networks for equivariance to lie groups arbitrary continuous data, arXiv preprint arXiv:2002. 2020.
386
397
396
[22]
Nathan Frey, Ryan Soklaski, Simon Axelrod, Siddharth Samsi, Rafael nomezz-oobbarlli, Connor Coley, and Vijay Gadepally, Neural scaling of deep chemical models. ChemRxiv, 2022.
399
400
[23] Fabian Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Se(3))transformes: 3d roto-translation equivariant attention networks. In Conference on Neural Information Processing (NeurlPS), 2020.
401
402
[24] Hongyang Gao and Shuiwang Ji. Graph representation learning via hard and channel-wise attention networks. Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, 2019.
403
405
[25] Johannes Gasteiger, Shankari Giri, Johannes T. Margraf, and Stephan Günnemann. Fast and uncertainty- aware directional message passing for non-equilibrium molecules. In Machine Learning for Molecules Workshop. NeurIPS, 2020.
406
407
408
[26] Johannes Gasteiger, Janek GroB, and Stephan Günneemann. Directional message passing for molecular graphs. In International Conference on Learning Representations (ICLR), 2020.
409
410
[27] Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C. Lawrence Zitmick, and Abhishek Das. How do graph networks generalize to large and diverse molecular systems? arxiv preprint arxiv:2204.02746 2022.
411
412
413
[28] Mario Geiger. Tess Smidt, Alby M.. Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice, Kostiantyn Lapehevskyi, Maurice Weiler, Michal Tyszkiewiez, Simon Batzner, Dylan Madisetti, Martin Uhrin, Jes Frellsen, Nuri Jung, Sophia Sanhorn, Mingjian Wen, Josh Rackers, Marcel Rød, and Michael Bailey. c3nn/e3nn: 2022-04-13, 13, April 2022.
414
415
416
417
[29] Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In International Conference on Machine Learning (ICML), 2017.
418
419
[30] Jonathan Godwin, Michael Schaarschmidt, Alexander I. Gaunt, Alvaro Sanchez-Gonzalez, Yulia Rubanova, Petar Veličković, James Kirkpatrick, and Peter Battaglia. Simple GNN regularisation for 3d molecular property prediction and beyond. In International Conference on Learning Representutions (ICLR), 2022.
420
421
422
[31] William L. Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Conference on Neural Information Processing (NeurIPS), 2017.
[32]
Kurt Homik. Approximation capabilities of multilayer feedforward networks. Neural Networks, 4(2):251-
257, 1991.
[33] Kurt Homik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural Networks, 2(5):359- 366, 1989.
427
428
[34] Andrew G. Howard, Menglong Zhu, Bo Chen, Dmitry Kalenichenko, Weijun Wang, Tobias Weyand, Marco Andreetto, and Hartwig Adam. Mobilenets: Efficient convolutional neural networks for mobile vision applications. ArXiv, abs/ 0ttss/704.0486.11 2017.
429
430
431
[35] Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-Isc: A large-scale challenge for machine learning on graphs. arXiv preprint arv: 103.00943, 2021.
432
433
[36]
Weihua Hu, Matthias Fey, Marinka Zitnik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint 2rXiv2005500068 2020.
434
435
436
[37] Weile Jia, Han Wang, Mohan Chen, Denghui Lu, Lin Lin, Roberto Car, Weinan E, and Linfeng Zhang Pushing the limit of molccular dynamics with ab initio accuracy to 100 million atoms with machine learning. In Proceedings of the International Conference for High Performance Computing, Networking. Storage and Analysis, SC '20. IEEE Press, 2020.
437
438
439
440
[38] Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations (ICLR), 2021.
441
11

---

443
[39]
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah. Transformers in vision: A survey. 2021.
445
[40]
Dongkwan Kim and Alice Oh. How to find your friendly neighborhood: Graph attention design with self- supervsueri In International Conference on Learning Representations (ICLR), 2021.
447
[41] Thomas N. Kipf and Max Welling. Semi-supervies classification with graph convolutional networks. In International Conference on Learning Representations (ICLR), 2017.
449
[42] Johannes Klicpera, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. In Conference Neural Information Processing (NeurIPS), 2021.
451
[43] Risi Kondor, Zhen Lin, and Shubhendu Trivedi, Clebsch-gordan nets: a fully fourier space spherical convolutional neural network. In Advances in Neural Information Processing Systems 32, pages 10117. 17- 10126, 2018.
452
454
[44] Devin Kreuzer, Dominique Beaini, William L. Hamilton, Vincent Létourneau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. In Conference on Neural Information Processing (NeurlPS), 2021.
456
457
[45] Cheng-I I Lai, Yang Zhang, Alexander H. Liu, Shiyu Chang, Yi-Lun Liao, Yung-Sung Chuang, Kaizhi Qian, Sameer Khurana, David Daniel Cox, and James R. Glass. PARP: Prune, adjust and re-prune for self-supervised speech recognition. In A. Beygelzimer, Y. Dauphin. P. Liang, and J. Wortman Vaughan, editors, Conference on Neural Information Processing (NeurIPS), 2021.
458
459
480
461
[46] Tuan Le, Frank Noé, and Amn Clevert. Equivariant graph attention networks for molecular property prediction. arXiv preprint 220.::222220991111 2022.
482
463
[47] Yi-Lun Liao, Sertae Karaman, and Vivienne Sze. Searching for efficient stalti-stageeann vision transformers. arxiv preprint 100i::209,,0642,, 2021.
464
465
[48] 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 on Learning Representations (ICLR), 2022.
466
487
468
[49] Denghui Lu, Han Wang, Mohan Chen, Lin Lin, Roherto Car, Weinan E, Weile Jia, and I Linfeng Zhang. 86 pflops deep potential molccular dynamics simulation of 100 million atoms with ab initio accuracy. Computer Physics Communications, 259:107624, 2021.
469
470
471
[50] Benjamin Kurt Miller, Mario Geiger, Tess E. Smidt, and Frank Noé. Relevance of rotationally equivariant convolutions for predicting molecular properties. arxiv preprint arxiv:2008.08461, 2020.
472
473
[51] Albert Musaclian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local cquivariant representations for large-scale atomistic dynamics. arxiv preprint 220.:::22040024449 2022.
474
476
[52] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, performefre deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Newral Information Processing Systems 32, pages 8024-800000 8035. Curran Associates, Inc., 2019.
479
480
481
482
[53] Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and Thomas F. Miller. OrbNet: Deep learning for quantum chemistry using yymetry--daptedd atomic-orbital features. The Journal of Chemical Physics, 153(12):124111, sep 2020.
483
485
[54] Joshua A. Rackers, Lucas Tecot, Mario Geiger, and Tess E. Smidt. Cracking the quantum scaling limit with machine learned electron densities. arxiv preprint 20kt::2001.03726. 2022.
487
[55] Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions. arXiv preprint 2017.
488
489
[56] Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Duta. 1, 2014.
490
491
[57] Bharath Ramsundar, Peter Eastman, Patrick Walters, Vijay Pande. Karl Leswing, and Zhenqin Wu. Deep Learning for the Life Sciences. O'Reilly Media, 2019. https://www sttps://wwwaamazon. amazon. com/
492
-eep-Learning- Life- wew---eeeneee-iicooooopoooo Microscopy/dpp////ooccccc 1492099877
12

---

494
[58] Yu Rong. Yatao Bian. Tingyang Xu, Weiyang Xie, Ying WEI, Wenbing Huang, and Junzhou Huang. Sule-superviee graph transformer large-scale molecular data. In Conference on Neural Information Processing (NeurIPS), 2020.
496
487
[59] Lars Ruddigkeit, Ruud van Deursen, Lorenz C. Blum, and Jean-Louis Reymond. Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17. Journal of Chemical Information and Modeling, 52(11) 1)22444288877 2012. PMID: 23088335.
488
439
[60]
Olga Russakovsky, Jia Deng. Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):11-222.. 2015.
501
502
503
504
[61]
Alvaro Sanchez- Gonzzlcz, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning (ICML), 2020.
505
506
507
[62] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In International Conference on Machine Learning (ICML), 2021.
508
509
[63] K. T. Schütt, P.-.. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, and K.-R. Müller. Schnet: A continuous-fiten convolutional neural network for modeling quantum interactions. In Conference on Neural Information Processing (NewrIPS), 2017.
510
512
[64]
Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In International Conference on Machine Learning (ICML), 2021.
513
515
[65] Yu Shi, Shuxin Zheng, Guolin Kc, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Yan Liu. Benchmarking graphormer large-scale molecular modeling datasets. arxiv preprint arxiv:2203.048// 2022.
518
[66] Yunsheng Shi, Zhengjie Huang, Shikun Feng, Hui Zhong, Wenjia Wang, and Yu Sun. Masked label predic- tion: Uniiied message passing model for semi-supervised classification. arxiv preprint 2ati::2000.035099 2020.
520
[67]
Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C. Lawrence Zitnick. Rotation invariant graph neural networks using spin convolutions. arxiv preprint arxiv:2 2021.
522
523
[68] Tess E. Smidt, Mario Geiger, and Benjamin Kurt Miller. Finding symmetry breaking order parameters with euclidean neural networks. Physical Review Research, 3(1), jan 2021.
525
[69] Anuroop Sriram, Abhishek Das. Brandon M Wood, and C. Lawrence Zitnick. Towards training billion parameter graph neural networks for atomic simulations. In International Conference on Learning Representations, 2022.
527
528
$29
[70] Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations (ICLR), 2022.
530
[71] Nathaniel Thomas, Tess E. Smidt, Steven Keamne,, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and ielllatiion-quvaiiaait neural networks for 3d point clouds. arxiv preprint ar / / 19, 2018.
532
533
[72] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. arXiv preprint arXiv:2012.1287,, 2020.
535
536
537
[73] Raphael J. L. Townshend, Brent Townshend, Stephan Eismann, and Ron O. Dror. Geometric prediction: Moving beyond scalars. arXiv preprint arXiv:2006. 14163, 2020.
538
539
[74] Raphael John Lamarre Townshend, Martin Vögele, Patricia Adriana Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Bowen Jing, Brandon M. Anderson, Stephan Fismann, Risi Kondor, Russ Altman, and Ron O. Dror. ATOM3d: Tasks on molecules in three dimensions. In Thirty-fifth Conference on Neural Information Processing Sysiems Dutasets and Benclmarks Track (Round 1). 2021.
540
541
542
5M4
[75] 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 3666333,,, may 2019.
13

---

548
[76] Oliver Thorsten Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus Robert Muller. SE(3)--quivariant prediction of molecular wavefunctions and electronic densities. In A. Beygelz- imer, Dauphin. P. Liang. Jnd Wortman Vaughan, editors, Conference on Neural Information Processing (NeurIPS), 2021.
548
549
550
[77] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Conference on Neural Information Processing (NeurIPS), 2017.
551
[78] Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Reprexentations (ICLR), 2018.
554
555
[79] Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen, 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data. In Advances in Neural Information Processing Systems 32, pages 1402-10413,2200 2018.
556
557
558
[80] Ross Wightman. Pytorch image models. https: http:://gtthbb..mm//giggggeeeeee wtw///wwgghhaaaayyyorroo---oeeee wwwp:/orrrr-iiaa-mmmdllee. 2019.
[81] Daniel E. Worrall, Stephan J. Garbin, Daniyar Turmukhambetov, and Gabriel J. Brostow. Harmonic networks: Deep translation and rotation equivariance. arxiv preprint arxiv: 1612.0444,, 2016.
561
562
[82] Tian Xic and Jeffrey C. Grossman. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Physical Review Letters, (20(14), apr 2018.
563
584
[83]
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference Learning Representations (ICLR), 2019.
[84] Chengxuan Ying. Tianle Cai, Shengjie Luo, Shuxin Zheng. Guolin Ke, Di He. Yanming Shen, and Tie- Yan Liu. Do transformers really perform badly for graph representation? In Conference on Neural Information Processing (NeurIPS), 2021.
587
568
569
[85] A. Zec. Group Theory in a Nutshell for Physicists. Princeton University Press, USA, 2016.
570
[86] Jiani Zhang, Xingjian Shi, Junyuan Xie, Hao Ma, Irwin King, and Dit-Yan Yeung. Gaan; Gated attention networks for learning on large and spatiotemporal graphs. In Proceedings of the Thirty- Fourth Conference Uncertainty in Artificial Intelligence, pages 339-349, 2018.
571
[87] Linfeng Zhang, Jiequn Han, Han Wang. Roberto Car, and Weinan E. Deep potential molecular dynamics: A scalable model with the accuracy of quantum mechanies. Phys. Rev, Lett., 120:14300: Apr 2018
574
Checklist
575
576
1. For all authors...
577
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
578
contributions and scope? [Yes)
579
(b) Did you describe the limitations of your work? [Yes]
580
(c) Did you discuss any potential negative societal impacts of your work? [Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
582
583
2. If you are including theoretical results...
584
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
585
(b) Did you include complete proofs of all theoretical results? [N/A]
586
3. If you ran experiments..
587
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material URL)? [Yes] We include code
588
589
related to experiments on QM9 in supplementary. Code reproducing all experiments will be available.
590
591
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
14

---

593
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [No] As each experiment takes much time, we are not able to run several times. However, the proposed method improves upon previous works by clear margin, and therefore the improvement not within noise ranges.
594
586
597
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, cloud provider)? [Yes]
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
600
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any assets either in the supplemental material as URL? [N/A]
603
604
(d) Did you discuss whether and how consent was obtained from people whose data you're re using/curating? [N/A]
605
606
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
608
5. If you used crowdsourcing conducted research with human subjects...
609
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
612
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
614
15

---

Appendix
615
618
A
Additional Mathematical Background
617
In this section, we provide additional mathematical background on group equivariance helpful for the discussion of the proposed method. Other works [71, 79, 43, 23, 5] also provide similar background. We encourage interested readers to see these works [85, 17] for more in-depth and pedagogical presentations.
618
619
620
A.1
Group Theory
622
Definition of Groups. A group is an algebraic structure that consists of set G and a binary
623
operator o : G G and is typically denoted G. Groups satisfy the following four axioms:
1. Closure: g o G for all g, G.
625
2. Identity: There exists an identity element E G such that goe= e for all g G.
3. Inverse: For each g E G. there exists an inverse element g- E G such that g o gog-
628
627
g =c
4. Associativity: o fogo (f g) oh f (g o h) for all f.g. EG
628
629
In this work, we focus on 3D rotation, translation and inversion. Relevant groups include:
630
1. The Euclidean group in three dimensions F(3): 3D rotation, translation and inversion.
2. The special Euclidean group in three dimensions SE(3): 3D rotation and translation.
632
3. The orthogonal group in three dimensions 03)): 3D rotation and inversion.
633
4. The special orthogonal group in three dimensions SO(3): 3D rotation.
Group Representations. The actions of groups define transformations. Formally, transformation acting on vector space X parametrized by group element E G is an injective function T3 X + X. powerful result of group representation theory is that these transformations can be expressed as matrices which act on vector spaces via matrix multiplication. These matrices are called the group representations. Formally, group representation D G - GL((N) is mapping between group G and set of N x N invertible matrices, The group representation D(g) X + X maps an N --ddimnsional vector space X onto itself and satisfies D(h) D(g h) for all g, h E G.
835
637
639
640
641 How a group is represented depends on the vector space it acts on. If there exists a change of basis P in the form of an N x N matrix such that P- D (9 for all G, then say the two group representations are equivalent. independent subspaces of the vector space, the representation D(g) is reducible. A particular class of representations that are convenient for composable functions are irreducible representations or "imeps"" which cannot be further reduced. We express any group representation of SO(3) direct sum (concatentation) of irreps [85, 17, 28]:
643
D')g is block diagonal, which means that acts on
645
647
(Diolt)
)
+ Dr.Go) p P-P-1
D(g) p-1
Dr, (g)
P
(6)
648
where (g) are Wigner-D matrices with degree as metnioned in Sec. 3.2.
649
A.2
Equivariance
650
Definition of Equivariance and Invariance. Equivariance a property function f:A.A. X + Y mapping between vector spaces X and Y. Given group G and group representations (g(g) and Dg) input and output spaces X and Y. f equivariant to G if (g) = T(Dx(g)z) for all x e X and g G. Invariance corresponds to the where Dy (g) is the identity I for all g G.
652
16

---

654
Equivariance in Neural Networks. Group equivariant neural networks are guaranteed to to make equivariant predictions on data transformed by group. Additionally, they are found to be data- efficient and generalize better than non-symmetr--wwan and invariant methods [4, 54, 22]. For 3D atomistic graphs, we consider equivariance to the Fuclidean group F(3), which consists of 31) rotation, translation and inversion. For translation, we operate on relative positions and therefore
655
656
657
658
659
our networks are invariant to 3D translation. We achieve equivariance to rotation and inversion by representing our input data, intermediate features and outputs in vector spaces of ((3) irreps and acting on them with only equivariant operations.
661
A.3
Equivariant Features Based on Vector Spaces of Irreducible Representations
663
Irreps Features. As discussed in Sec. 3.2 in the main text, we use type-L vectors for SE(3)- equivariant irreps features' and type-(L,p) vectors for -qquivarra irreps features. Parity p denotes whether vectors change sign under inversion and can be either (even) (odd). Vectors with p = o change sign under inversion while those with p = e do not. Scalar features correspond to type-0 vectors in the case of S (333)-qquivarianee and correspond to type-(0, e) in the case of E 3) equivariancce whereas type-(0, o) vectors correspond to pseudo-scalars. Euclidean vectors in 3 correspond to type-1 vectors and type- o) vectors whereas type-(1, e) vectors correspond
684
665
666
667
668
670 to pseudo-vectors. Note that type-(L, e) vectors and type-(L, o) vectors are considered vectors of different types in equivariant linear layers and layer normalizations.
Spherical Harmonics. Euclidean vectors i' in R$3 can be projected into type-L L vectors f(() by 672
using spherical harmonics Y(L); f(L) ( ( [ [68]. This equivalent to the Fourier transform of the angular degree of freedom | which can be optionally weighted by ||7]]. In the case of 674
673
S (E) -equivariance, f(4) transforms in the same manner as type-L. vectors. For -qqievanie, f behaves type-(T,,p) p) vectors, where p e if L. is even and p if L is odd.
675
676
677
Vectors of Higher and Other Parities. Although previously we have restricted concrete ex- amples of vector spaces of 0(3) irreps to commonly encountered scalars (type-(0, e) vectors) and Euclidean vectors (type-(1, (1, o) vectors), vector of higher L and other parities are equally physical. For example, the moment of inertia (how an object rotates under torque) transforms as a 3 x 3 symmetric matrix, which has cammetric-traceles
678
679
680
681
components behaving as type-(2, e) vectors. Elasticity (how an object deforms under loading) transforms as rank-4 or 3 x 3 x 3 x 3 symmetric tensor, which includes components acting as type-(4, e) vectors.
683
A.4
Tensor Product
685
Tensor Product for O(3). We use tensor products to interact different type-(L, p) vectors. We extend our discussion in Sec. 3.3 in the main text to include inversion and type-(L, p) vectors. The tensor product denoted as & uses Clebsch-Gordan coefficients to combine (ype-(Tn,P1) P1) vector
686
687
f (1)) and type-(L2, P2) vector g ((l1a.22) and produces type- (L3, vector (Ia.pa) as follows:
688
La
(ss.p)) = (f(bxip)) @ ( mg
Σ m2
(La.mma))
fth Jmnz
p(L..P)) (22.
L1 mg Lg
P3- = P1 X P2
(8)
The only difference of tensor products for 0(3) described in Eq. from those for SO(3) described in Eq. I is that we additionally keep track of the output parity ps as in Eq. 8 and use the following multiplication rules: e x e, o e, and x o o e o. For example, the tensor product of a type-(1, o) vector and a type-(1, e) vector can result in one type-(0, o) vector, one (1p vector, and one (2pe o) vector.
680
691
692
693
In SEGNN [5]. they are also referred to as steerable features. We use the term "irreps features" to remain consistent with e3nn [28] library.
17

---

694 Clebsch-Gordang Coefficients. The Tebscch-Goddan coefficients for SO(3) are computed from
695 integrals over the basis functions of given irreducible representation. e.g. the real spherical harmonics, as shown below and are tabulated to avoid unnecessary computation.
C ma) (L),,))))) [iimi L2ma) (Lymal
(4a)) (1)
(9)
For many combinations of L1. L2, and L3, the Clebsch-Gordan coefficients are zero. The gives rise to the following selection rule for non-trivial coefficients: -|ا + L S L3 S اا + L2).
698
Examples of Tensor Products. Tensor products generally define the interaction between different type-(I., () vectors in a syymmetry-preservinn manner and consist of common operations as follows:
700
701
1. Scalar-scalar multiplication: scalar (L = 0, e) scalar (L = 0.p = e) + scalar
702 (L 0.p e).
2. Scalar-vector multiplication: scalar (L = 0, e) vector (L = -1,p 0) + vector
704
(L = 1, = o).
705
3. Vector dot product: vector (L - @ vector (L = 1, 0) - scalar (L 0,p
706
e).
707
4. Vector cross product: vector (L 1,p = o) @ vector (L = 1, = o) + pseudo-vector
(L 1, e).
B Related Works
710
B.1
Graph Neural Networks for 3D Atomistic Graphs
711
Graph neural networks (GNNs) are well adapted to perform property prediction of atomic systems because they can handle discrete and topological structures. There are two main ways to represent atomistic graphs [74], which are chemical bond graphs, sometimes denoted 2D graphs, and 3D spatial graphs. Chemical bond graphs use edges to represent covalent bonds without considering 3D geometry. Due to their similarity to graph structures in other applications, generic GNNs [31, 29, 41,
712
714
83, 78, 6] can be directly applied to predict their properties [59, 56, 57, 36, 35]. On the other hand, 716
3D spatial graphs consider positions of atoms in 3D spaces and therefore 3D geometry. Although
718
3D graphs can faithfully represent atomistic sy systems, one challenge of moving from chemical bond graphs to 3D spatial graphs is to remain invariant equivariant to geometric transformation acting on atom positions. Therefore, invariant neural networks and equivariant neural networks have been proposed for 3D atomistic graphs, with the former leveraging invariant information like distances and angles and the latter operating on geometric tensors like type-L L vectors.
720
722
723
B.2
Invariant GNNs
724
Previous works [63, 82, 75, 26, 23, 48, 67, 42] extract invariant information from 3D atomistic graphs and operate on the resulting invariant graphs. They mainly differ in leveraging different geometric information such as distances, bond angles (3 atom features) dihedral angles (4 atom features). SchNet [63] uses relative distances and proposes continuous-filte convolutional layers to learn local interaction between atom pairs. DimeNet series [26, 25] incorporate bond angles by using triplet representations of atoms. SphereNet [48] and GemNet [42, 27] further extend to consider dihedral angles for better performance. In order to consider directional information atoms. In addition to being memory- ieemory-ittensive intensive [69], they also change graph structures by introducing higher-order interaction
725
727
729
731
contained in angles, they rely on triplet or quadruplet representations
terms [11], which would require non-trivial modifications to generic GNNs in order to apply them 739
to 3D graphs. In contrast, the proposed Equiformer uses equivariant irreps features to consider directional information without complicating graph structures and therefore can directly inherit the
735
736
design of generic GNNs.
B.3
Attention and Transformer
738
Graph Attention. Graph attention networks (GAT) [78, 6] use multi-layer perceptrons (MIP) to calculate attention weights in a similar manner to message passing networks. Subsequent works
18

---

740
using graph attention mechanisms follow either GAT-like MLP attention [8, 40] or Transformer-lic dot product attention [86. 24, 66, 18, 40, 44]. In particular, Kim et al. [40] compares these two types of attention mechanisms empirically under self-supervised setting. Brody el al. [6] analyzes their theoretical differences and compares their performance general settings.
741
743
Graph Transformer. A different line of research focuses adapting standard Transformer works to graph problems [18, 58, 44, 84, 65]. They adopt dot product attention in Transformers [77] and propose different approaches to incorporate graph-related inductive biases into their networks. GROVE [58] includes additional message passing layers or graph convolutional layers to incorporate
745
747
748
local graph structures when calculating attention weights. SAN [44] proposes learn position embeddings of nodes with full Laplacian spectrum. Graphormer [84] proposes to encode degree information in centrality embeddings and encode distances and edge features in attention biases. The
749
750
proposed Equiformer belongs to one of these attempts to generalize standard Transformers to graphs and is dedicated 3D graphs. To incorporate 3D-related inductive biases, we adopt an equivariant version of Transformers with irreps features and propose novel equivariant graph attention.
752
753
C
Details of Architecture
754
Iapus imappe
lepas impe featuns
Ieput imqus featano
Trpai imone fosteos
Impal imcope fasturcs
Inpu imcpe fasurcs 2
Ecauurs
(n) Paths contributing type-O vectors.
(b) Paths contriburing type-/ vectors.
(c) Paths contriburing to type-2 vectues.
Figure 3: An alternative visualization of the depth-wise tensor product. We follow the visualiza- tion of tensor products e3nn [28] and separate paths into three parts based on the types of output vectors. We note that one vector in the output irreps feature depends only one vector in each input irreps feature.
C.1
Equivariant Operation Used in Equiformer
756
We illustrate the equivariant operations used in Equiformer in Fig. and provide an alternative visualization of depth-wise tensor products in Fig. 3.
758
C.2
Equiformer Architecture
759
For simplicity and because most works we compare with do not include equivariance to inversion, we adopt (E(3) -eeeevarrrai irreps features in Equiformer for experiments in the main text and note that -qui--quuraiiant irreps features can be easily incorporated into Equiformer.
760
We define architectural hyper--aramett like the number of channels in some layers in Equiformer, which are used to specify the detailed architectures in Sec. and Sec. F.
763
We dembed to denote embedding dimension, which defines the dimension of most irreps features. Specifically, all irreps features x,, y/ in Fig. have dimension dembed unless otherwise stated. Besides, use deh represent the dimension of spherical harmonics embeddings of relative positions in all depth-wise wise tensor products.
765
767
768
For equivariant graph attention in Fig. 1(b), the first two linear layers have the same output dimension dembest The output dimension of depth-wise tensor products (DTP) are determined by that of input irress features. Equivariant graph attention consists of h parallel attention functions, and the value vector in each attention function has dimension dhead- We refer to and dhead as the number of heads and head dimension, respectively. By default, we set the number of channels in scalar feature
769
771
19

---

Linear
Linear
Linear
DIP
Linear
fy
Dox Producc
Scale
Sofbrrs
f
dy
Reshsipe
Linear
y//
Figure 4: Architecture of equivariant dot product attention without non-linear message passing. In this figure, "ø" denotes multiplication, "Ø" denotes addition, and "DTP" stands for depth-wise tensor product. Σ within a circle denotes summation over all neighbors. Gray cells indicate intermediate irreps features. We highlight the difference of dot product attention from multi-layer perceptron attention in red. Note that key kiz and value V/j are irreps features and therefore fiy in dot product attention typically has more channels than that in multi-layer perceptron attention.
Jig to be the same the number of channels of type-0 type-(0, e) vectors in Cij. When non-linear messages are adopted in Vij, we set the dimension of output irreps features in gate activation to two hyper-parametent h and dhead to specify the detailed architecture of equivariant graph attention.
(0)
774
be h x daead- Therefore,
776
777 As for feed forward networks (FFNs), we denote the dimension of output irreps features in gate activation d fn- The FFN in the last Transformer block has output dimension dicatures and we set fn of the last FFN, which is followed by output head, to be d feature well. Thus, two hyper- parameters fn and d'feature are used to specify architectures of FFNS and the output dimension
778
779
780
after Transformer blocks.
Irreps features contain channels of vectors with degrees up Lmaz- We denote Cr type-L vectors as (CL, L) and C(r..)) type-(L, p) vectors as (C(r.p) L. p) and use brackets to represent concatenations of vectors. For example, the dimension irreps features containing 256 type-O vectors and 128 type-1 vectors can be represented as [[256, (),6)))),1))) (128, 1)].
784
785
786
C.3
Dot Product Attention
787
We illustrate the dot product attention without non-lineat message passing used in ablation study in Fig. 4. The architecture is adapted from 3E(3)-Transforma [23]. The difference from multi-layer perceptron attention lies in how we obtain attention weights aij from fiy- We split fy into two irreps features, key kij and value waj. and obtain query with linear layer. Then, we perform scaled dot product [77] between qe and kig for attention weights.
789
791
C.4
Discussion Computational Complexity
We discuss the computational complexity of the proposed equivariant graph attention here.
20

---

794
First, we compare dot product attention with MLP attention when linear messages are used for value Vij. Dot product attention requires taking the dot product of two irreps features, query qi and key kaj, for attention weights, and both d and kij have the same dimension value vig. In contrast,
795
MLP attention uses only scalar features Jij for attention weights. The dimension of scalar features
797
798
fi is the same as that of the scalar part of Vtj. Therefore, MLP attention generates less and smaller intermediate features for attention weights and is faster than dot product attention.
800
Second, compared to linear messages, using non-linear messages increases the number of tensor product operations from to 2. Since tensor products are compute-intensive this inevitably increases training and inference time.
802
803
Please refer to Sec. D.2 and Sec. E.2 for the exact numbers of training time on QM9 and OC20.
D Details of Experiments on QM9
805
D.I
Additional Comparison between SE(3) and E(3) Equivariance
806
We train two versions of Equiformers, one with - features denoted as "Equiformer" and the other with -e3)-qquivariant features denoted as " *3(3)-Bu EEquiforrr""""" and we compare them
807
808
Table 8. Including equivariance inversion further improves the performance on QM9 dataset.
809
As for Table 1, we compare "Equiformer" with other works since most of them do not include equivariance to inversion.
810
Task
a
Ar
C₂ cal/mol K
Methods
Units
bohr3
meVV
meV
meV
Equiformer
056
33
17
16
.025
E(3)- Equiformer
054
32
16
16
.013
.04
Table 8: Ablation study of (E(3))((3 equivariance QM9 testing set. "Equiformer" operates on (3))))cquivarri features while (3(())-qqffommer"" uses E(3) features. Including inversion further improves mean absolute errors.
811
D.2
Training Details
812 We normalize ground truth by subtracting mean and dividing by standard deviation. For the task of U, Up. G. and H. where single-atom reference values are available, we subtract those reference values from ground truth before normalizing.
813
815
We train Equiformer with 6 blocks with Lmax following SEGNN [5]. We choose Gaussian radial basis [63, 67, 42. 65] for the first six tasks in Table 1 and radial Bessel basis [26, 25] for the others. Table 9 summarizes the hyper-pue pppee-aarseeters for the QM9 dataset. Further details will be provided the future. The detailed description of architectural aaper-parameters can be found in Sec. C.2.
816
817
818
We use A6000 GPU with 48GB train each model and summarize the computational cost of training for one epoch as follows. Training F(33 Eqqiformer for one epoch takes about 14.75 minutes. The time of training Equiformer, Equiformer with linear messages (indicated by index 2 in Table 6), and Equiformer with linear messages and dot product attention (indicated by index 3 in Table for one epoch is 11 minutes, 6.6 minutes and 7.1 minutes, respectively.
820
822
823
824
E
Details of Experiments on OC20
825
E.1
Additional Comparison between SE(3) and E(3) Equivariance
826
We train two versions of Equiformers, one with SE(3) features denoted as "Equiformer" and the other with Ee3)-qquivariant features denoted as (E(33)Eqqiformm))
827
and we compare them in Table 10. Including inversion improves the MAE results on ID and OOD Cat sub-splits but
829 degrades the performance on the other sub-splits. Overall, using -qqiivaauva features results in slightly inferior performance. We surmise the reasons are as follows. First, inversion might not be the
21

---

sapee--prameeers
can be Value description
Optimizer
AdamW
Learning rate scheduling
Cosine learning rate with linear warmup
Warmup epochs
Maximum learning rate
x 10
Batch size
128
Number epochs
300
Weight decay
10-3
Cutoff radius (A)
Number of radial bases
128 for Gaussian radial basis, 8 for radial bessel basis
Hidden sizes radial functions
64
Number hidden layers radial functions
Equiformer
Number of Transformer blocks Embedding dimension dlonted Spherical harmonics embodding dimension dah
(128. 0), (64, 1), (32, (1.
Numher attention heads Attention head dimension dionnd
(32, 0j. (16, 1), (8. 384,0) (192, (96.2) (512.0)
Hidden dimension in feed forward networks dzfm
Output feature dimension deaatur
Eqquioorer
Number of Transformer blocks Embedding dimension dembeed Spherical harmonics embedding dimension deh.
(128, (32 (32,0,0) (2), (1.,,) (1.2,e)
Numher attention heads
Attention head dimension diood
(32, (.00000000000..1000000.000000000 (384,0,e). 0),), (96, 0). (96) 1,0), (18. (8).2.)), (48.2
Hidden dimension in feed forward networks dyfm
Output feature dimension dycature
(512,0,e)
Table 9: (ype-(I,,)) vectors as (C(L.p): p. and use brackets to represent concatenations of vectors.
for QM9 dataset. We denote CL type-L vectors as (C2, L) and C(r.p)
key bottleneck. Second, including inversion would break type-1 vectors into two parts, type-(1, e) 832 and type-(1, vectors. They are regarded as different types equivariant linear layers and layer normalizations, and therefore, the directional information captured in these two types of vectors can only exchange depth-wise tensor products. Third, we mainly tune hyper-parameters for Equiformer with S E lequivrrra features, and it is possible that using =equivaraaa features would favor different hyper parammttareeeers.
834
836
837
For Table 2, 3, 4, and 5, we compare "Equiformer" with other works since most of them do not include equivariance to inversion.
838
Energy MAE (
EwT (%)
Methods
ID
OOD Ads
OOD Cat
OOD Both
Average
ID
OOD Ads
OOD Cat
OOD Both
Average
Equiformer
0.5088
0.6271
0.5051 0.5034
0.5545
0.5489
4.88
2.93
4.92
2.98
3.93
E(3) Equiformer
0.5035
0.6385
0.5658
0.5528
5.10
2.98
5.10
3,02
4.05
Table 10: Ablation study of (E(3)E(() "Equiformer" operates (quui-equivariant features while "E(3)-Equiformoe) cquivariant features.
equivariance on OC20 IS2RE validation set.
E(3)-
E.2
Training Details
840
IS2RE without Node-Level Auxiliary Task. We use
similar to those for QM9 dataset and summarize in Table 11. The detailed description of architectural found in Sec. C.2.
842
843
IS2RE with IS2RS Node-Level Auxiliary Task. We increase the number of Transformer blocks
to 18 deeper networks benefit more from IS2RS node-level auxiliary task [30]. We follow
22

---

845
the same hyper- parameters in Table 11 except that we increase maximum learning rate to 5 x 10-4 and set feoture to [(512, 0), (256, 1)]. Inspired by Graphormer [65], we add an extra equivariant graph attention module after the last layer normalization to predict relaxed structures and use linearly decayed weight for loss associated with IS2RS, which starts 15 and decays to For Noisy Nodes [30] data augmentation, we first interpolate between initial structure and relaxed structure and then add Gaussian noise as described by Noisy Nodes [30]. When Noisy Nodes data augmentation is used, we increase the number of epochs to 40. Further details will be provided in the future.
846
848
850
852
We use two A6000 GPUs, each with 48GB, train models when IS2RS is not included during training. Training Equiformer and Equi-fqqef takes about 43.6 and 58.3 hours. Training Equiformer with linear messages (indicated by index 2 in Table 7) and Equiformer with linear messages and dot product attention (indicated by index 3 in Table 7) takes 30.4 hours and 33.1 hours, respectively. We use four A6000 GPUs to train Equiformer models when IS2RS node-level auxiliary task is adopted during training. Training Equiformer without Noisy Nodes [30] data augmentation takes about 3 days and training with Noisy Nodes takes days. We note that the proposed Equiformer in Table 5 achieves competitive results even with much less computation. Specifically, training "Equiformer + Noisy Nodes" takes about 24 GPU-days when A6000 GPUs are used. The training time of "GNS Noisy Nodes" [30] is 56 TPU-days. "Graphormer" [65] uses ensemble of 31 models and requires 372 GPU-days to train all models when A100 GPUs are used.
853
854
855
856
857
858
860
862
Slyper-prrametes
Value or description
Optimizer
Adam
Learning rate scheduling
Cosine leurning rate with linear warmup
Warmup epochs
2
Maximum learning rale
2 x 10
Batch size
32
Number of epochs
20
Weight decay
1 10-3
Cutoff radius (A)
5
Number of radial hasis
128
Hidden size of radial function
64
Number hidden layers in radial function
2
Equiformer
Number of Transformer blocks Emhedding dimension denieed Spherical harmonics embedding dimension d.i.
[256. 0), (128, 1)]
(1. (1. 1)]
Number of attention heads h
8
Attention head dimension dacad Hidden dimension feed forward networks dp
(32, 0). (16, 1)]
(768, (384, 1)
Output feature dimension fe
(512, 0j]
13 -33--oqufformer
Number of Transformer blocks Emhedding dimension demteed Spherical harmonics embedding dimension daa.
(256. 0,6,0,e0 (64.0, o), (64, (64, 1,4,
(1.0,))
Number attention heads h Attention head dimension daead Hidden dimension in feed forward networks drym
e... (8.0,)), (8, (8, (8). 1,
0....0 (192.0.)) (.)) 1,0]
Output feature dimension festure
[312,0.e]]
Table 11: Hyper-paree parameerssss for OC20 dataset under the setting of training without IS2RS aux- iliary task. We denote C₁ type-L vectors (C',, L) and C(L.p) type-(L,p) vectors as (C(L,)) L.p) and use brackets to represent concatenations of vectors.
863
E.3
Error Distributions
864
We plot the error distributions of different Equiformer models on different sub-splits of OC20 IS2RE validation set in Fig. 5. For each curve, we sort the absolute errors in ascending order for better visualization and have a few observations. First, for each sub-split, there are always easy examples, for which all models achieve significantly low errors, and hard examples, for which all models have high errors. Second, the performance gains brought by different models are non-uniform among different sub-splits. For example, using MLP attention and non-linear messages improves the errors
885
866
868
23

---

870
on the ID sub-split but is not that helpful on the OOD Ads sub-split. Third, when IS2RS node-level auxiliary task is not included during training. using stronger models mainly improves errors that are beyond the threshold of 0.02 eV. which is used to calculate the metric of energy within threshold (EwT). For instance, on the OOD Both sub-split, using non-linear messages, which corresponds
871
872
873
874
to red and purple curves, improves the absolute errors for the 15000th through 20000th examples. However, the improvement in MAE does not translate to that in EwT as the errors are still higher than
875
the threshold of 0.02 eV. This explains why using non-linear messages in Table 7 improves MAE from 0.5657 to 0.5545 but results in almost the same EWT.
877
878
F
Limitations
We discuss several limitations of the proposed Equiformer and equivariant graph attention below.
880
First, Equiformer is based on irreducible representations (irreps) and therefore can inherit the limitations common to all equivariant networks based on irreps and the library e3nn [28]. For example, using higher degrees L can result in larger features and using tensor products can be compute-intensive. Part of the reasons that tensor products can be computationally expensive are that the kernels have not been heavily optimized and customized as other operations in common libraries like PyTorch [52]. But this is the issue related to software, not the design of networks. While tensor products of irreps naively do not scale well, if all possible interactions and paths are considered, some paths in tensor products can also be pruned for computational efficiency. We leave these potential efficiency gains to future work and in this work focus on general equivariant attention if all possible
882
884
886
887
888
889
paths up Lmaz in tensor products are allowed.
890 Second, the improvement of the proposed equivariant graph attention can depend on tasks and datasets. For QM9, MLP attention improves not significantly upon dot product attention as shown in Table 6. We surmise that this is because QM9 contains less atoms and less diverse atom types and therefore linear attention is enough. For OC20, MLP attention clearly improves upon dot product attention as shown in Table 7. Non-linear messages improve upon linear ones for the two datasets.
892
894
895
Third, equivariant graph attention requires more computation than typical graph convolution. It includes one softmax operation and thus requires one additional sum aggregation compared to typical message passing. For non- non-linear message passing, it increases the number of tensor products from one to two and requires more computation. We note that if there is a constraint on training budget, using stronger attention (i.e., MLP attention and non-linear messages) would not always be optimal because for some tasks or datasets, the improvement is not that significant and using stronger attention slow down training. For example, for the task of C₂ QM9, using linear (index 2) non-linear messages (index 1) results in the same performance as shown in Table 6. However, non- non-lincar messages increase the training time of one epoch from 6.6 minutes to 11 minutes.
896
897
896
900
902
904
Fourth, the proposed attention has complexity proportional to the products of numbers of channels and numbers of edges since the the attention is restricted to local neighborhoods. In the context of 3D atomistic graphs, the complexity the same as that of messages and graph convolutions. However, in other domains like computer vision, the memory complexity of convolution is proportional to the number of pixels or nodes, not that of edges. Therefore, it would require further modifications in order to use the proposed attention in other domains.
905
906
908
24

---

10
10
Absokite Error
0.02 Threshald
0.02 eV Threshold
10-
Absolite
10"
10"
10
10000 Examples Sorted
20000
25000
10000
15000
26300
Examples by Error
Equiformer with IS2RS and Noisy Nodes
Equifomer with IS2RS and Noley Nodes
Equiformer with IS2RS
Equifomer with IS2RS
Equifermer
Equiformer with linear messages
Equifomer with Ineer messages
Equiformer with linear messages and dol product attention
Equifomer with ineer messages and dot product attention atteuctattnntionn
(a) ID sub-split.
(b) OOD Ads sub-split.
10'
10
10"
Absolute Eror
Absolute Enor
10
0.02 aV Trreshald
0.02 WV Threshold
10
10
10
5000
10000
19300
30000
25000
0
6003
10020
16000
20300
29000
Examples Sorled Absolute Error
Examples Sorted by Absolute Error
Equiformer with IS2RS and Noisy Nodes Equiformer
Equiforer with IS2RS and Noisy Nodes
Equifermer
Equiformr
Equifermer with linear messages
Equiformer with inear messages
Equiformer with linear messages and dot product attention
Equiformr with Inear messages and dot product attention
(c) OOD Cat sub-split.
(d) OOD Both sub-split.
Figure 5: Error distributions of different Equiformer models on different sub-splits of OC20 IS2RE validation set.
25


Reviews:
[{"review_id": "om_ojFrDpUD", "reviewer": "Reviewer_5yNV", "summary": "The paper proposes an irreps-based Transformer architecture for atomistic property prediction, named Equiformer. The model leverages irreps to interact between features with different degrees, very much similar to previous literature like TFN, SE(3)-Transformer, or SEGNN. Equiformer further incorporates a GAT-like attention mechanism to enhance the capacity of the model. Equiformer has been experimentally evaluated on two datasets QM9 and the OC20 IS2RE task, obtaining competitive performance against other equivariant graph neural networks.", "questions": "The questions below are proposed according to the weaknesses specified in the previous section.\n\n**Q1:**\nAre the authors able to systematically discuss the insights on why Equiformer is advantageous over the existing irreps-based methods and equivariant transformer architectures? Potential perspectives would include theoretical analysis (e.g., on the expressivity or universality) or other systematical comparisons and discussions towards this aspect.\n\n**Q2:**\nIt is strongly recommended to report the existing state-of-the-art results on QM9, like PaiNN and TorchMD-Net (ET). The authors can find these references given in the previous section. By considering these results, the performance of Equiformer on QM9 seems to be not that competitive.\n\n**Q3.**\nIt would be great to involve running-time comparison/parameter complexity of both Equiformer and the most competitive baseline on both QM9 and OC20. From the current paper, it is very hard to tell where the benefit comes from. One potential concern is related to fairness that Equiformer might be leveraging much higher computational complexity with slower training/inference speed as well as a larger network size.", "limitations": "The authors claimed in the checklist that they have discussed the limitations. However, to the best of my effort, I am not able to find them in the paper. The authors are strongly recommended to discuss some limitations, potentially, as specified in Q3, from the high computational complexity of irreps and attention module.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "## Strengths\n1. Clear presentation with the necessary background on group theory and the irreps-based equivariant construction approach.\n2. Detailed depiction of the model architecture and the proposed attention mechanism.\n3. Nice ablation studies on both QM9 and OC20 to investigate the role of the proposed message-passing and attention module.\n\n## Weaknesses\n1. The technical contribution is somewhat limited given the rich literature on irreps-based equivariant GNNs like TFN [1] and SE(3)-Transformer [2]. There are also equivariant Transformers in this domain, e.g., SE(3)-Transformer [2], TorchMD-Net (ET) [3], and EQGAT [4]. Though some of them have been discussed in the related work, there is still a lack of insightful discussions on why the proposed Equiformer is superior over the so many existing ``Equivariant Transformer'' networks, where the differences lie, and what the key contributing parts are.\n2. The enhancement on QM9 is observed to be limited. Some important state-of-the-art models for atomistic property prediction are missing Table 1 (MAE result on QM9 test set), including PaiNN [5] and TorchMD-Net [3]. In particular, TorchMD-Net is also a Transformer-like architecture, which, I believe, is relevant to the paper here and should be compared in Table 1.\n3. Lacking discussions about modeling complexity and computational overhead. Generally, the irreps-based methods have already been observed to be quite computationally expensive, compared to other approaches for achieving equivariance, such as the scalar-based networks (e.g., EGNN [6], GMN [7], to name a few). Further incorporating attention will bring extra complexity to the model that is already very heavy to train and inference. There are no comparisons on the running time and number of parameters consumed compared with other baselines.\n\nRefs:\n\n[1] Thomas et al. Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds. 2018.\n\n[2] Fuchs et al. SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks. NeurIPS 2020.\n\n[3] Thölke et al. Equivariant Transformers for Neural Network based Molecular Potentials. ICLR 2022.\n\n[4] Le et al. Equivariant graph attention networks for molecular property prediction. 2022.\n\n[5] Schutt et al. Equivariant message passing for the prediction of tensorial properties and molecular spectra. ICML 2021.\n\n[6] Satorras et al. E(n)-equivariant graph neural networks. ICML 2021.\n\n[7] Huang et al. Equivariant Graph Mechanics Networks with Constraints. ICLR 2022.\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": "Zzah3p80u0x", "reviewer": "Reviewer_QY2z", "summary": "This paper proposes an E(3)/SE(3) equivariant transformer on 3D molecular graphs. The central point is to apply a non-linear MLP attention mechanism based on the type-0 irreps features. Then the attention weights are multiplicated with other irreps features with type >0. The evaluations are carried out on QM9 and OC20, which somehow supports the benefit of the proposed method. ", "questions": "1.\tIt is suggested to move the related work right after the introduction part, and provide more comparisons between different equivariant transformer methods.\n\n2.\tLine 107-118: does the type of inversion actually involved in the current implementation?\n\n3.\tLine 179-184: how to determine the pair (L, L’)? By random?\n\n4.\tLine 66-67: Not all GNNs are based on message passing based models.\n\n5.\tLine 254-258: the comparisons here are too shot to justify the contributions of this work.\n", "limitations": "The limitations and potential negative societal impact are not explicitly discussed. ", "rating": 5, "confidence": 4, "soundness": 1, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1.\tThis paper is overall compactly written and easy to follow.\n\n2.\tBesides QM9 (which is actually a well-explored dataset and could be overfitted by many recent methods), the authors evaluate their method on OC20, a new but more desirable dataset for performance comparison. It is valuable to see that the proposed method achieves good results on IS2RE. Necessary ablation studies are also performed. \n\n\nWeaknesses:\n\n1.\tThe biggest concern is that the methodological novelty is weak. It has been a common practice to explore transformers for graph data, such as Grover and Graphormer. And, even when considering equivariance, SE3-transformer and the works by [53,35] have investigated how to conduct attentions between equivariant features. While this paper does utilize certain slightly different components (such as LN, DTP), it is hard to see something essentially new in this submission compared to current papers. The authors claim the benefit of using MLP attention against the traditional dot-product counterpart. Yet, this replacement of attention computation just seems a trick other than a valuable contribution, particularly given that the MLP attention mechanism have already explored before in the graph learning field, such as GATv2 as cited by the authors. \n\n2.\tIt is nice to see the evaluations on OC20. Yet, it seems Equiformer obtains better performance in terms of MAE, but not regarding EwT. Does this mean Equiformer is over-parametric and fits the data well in some cases, but not that good on average? Experimentally, it is still unknown why the proposed attention can help promote the performance, and why it is better than other equivariant transformers such as SE3-transformer and [53,35]. \n\n3.\tSEGNN is a recent paper that can be regarded as a generalization version of many equivariant methods including tensor-field, LieConv, NequIP, EGNN, SE3 transformer, etc. Clearly, this paper has followed the notations and presentations in SEGNN (see for example the denotation in Eq. (6)). When compared to SEGNN, the equivariant attention in this paper actually degenerates to a specific case of SEGNN, for example, the attentions are given by Eq. (14) in SEGNN paper when \\alpha(f_i, f_j) becomes \\alpha(f_ij). \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": "DiVlrae79l", "reviewer": "Reviewer_qLqa", "summary": "This paper proposes a novel Equiformer model for DFT-level scalar quantum properties prediction based on 3D conformation of molecules or molecules and catalysts system.\nEquiformer has permutation equivariant and SE(3)/E(3) equivariant architecture, a novel transformer-like architecture operating on irreducible representations(Irreps) feature.\nThis paper proposes a novel Depth-wise Tensor Product method, creatively adopts MLP attention for Irreps and borrows the successful Non-Linear Message from SEGNN as Value in attention mechanism.\nExperiments are conducted on QM9 and OC20-IS2RE, exhibiting the effectiveness of Non-Linear Message, MLP attention and overall architecture, achieving SOTA performance in both tasks.\n", "questions": "+ All interactions between different type vectors and different elements within vector depend on DTP with $SH(r_{ij})$, while no experiment shows Self (Depth-wise) Tensor Product or Cross (Depth-wise)  Tensor Product cannot improve the performance. \n+ Comparison with Dot Product Attention could be more complete, e.g., key can be obtained from $x_j$ with a linear layer and attention bias can be obtained from the scalar part of $f_{ij}$ with a MLP before Reshape (splitting heads) or learnable radial function(mentioned in line 235-239) encoding $||r_{ij}||$, etc.\n", "limitations": "The authors didn’t describe the limitations of their work.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n+ Proposes Novel Depth-wise Tensor Product\n+ Adopts MLP attention for Irreps and shows its effectiveness\n+ Clear Figure 2 in appendix describes how each component works.\n+ Proposed S×SE(3) Equivariant Equiformer achieves SOTA performance on QM9 and OC20-IS2RE\n\nWeaknesses:\n- Attention Score component could cooperate with Attention Value component, ablation studies only conducted on MLP attention + Non-Linear Message, MLP attention + Linear Message and Dot Product attention + Linear Message. It seems that is enough to show the effectiveness of each component, but the core contribution of this paper is the combination of MLP attention and Non-Linear Message. Thus it might be necessary to compare MLP attention to Dot Product attention when Non-Linear Message is used.\n- No experiment or explanation shows the effectiveness/efficiency of DTP.\n- Equiformer is only applied to scalar prediction tasks (except auxiliary task), which only need SE(3) invariance or trivial equivariance(i.e. identity representation). The experiments cannot show the significance of the equivariant nature of Equiformer.\n- Equiformer architecture is very complicated, while there isn’t a clear motivation or intuition or explanation or thorough experimentation to support it. If the embedding module and the parameterization of DTP weights are considered, Equiformer could be more complicated.\n- Section 3.1(line 154-184) is less informative than Figure 2, while occupies more space and harder to understand. Figure 2 is more important than the verbose background introduction in Section 2(Section 2.3 is good, while other parts could be more compact or omitted). And given line 114 has already said $C_L$ is the number of channels for type-L vectors, line 174-175 says “input x containing $(C_0+\\sum_{L=1}^{L_{max}}C_L)$ type-0 vectors” is confusing. It might be better to replace this paragraph with Figure 2(c).\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": "5ajWccxDjL", "reviewer": "Reviewer_YZ7s", "summary": "The paper proposes a transformer architecture which is equivariant to the Euclidean group SE(3). While other transformer architectures which are equivariant exists, the main novelty of this paper consists in the ability to incorporate higher dimensional irreps, thus, not limiting itself to the l=0,1 representations (scalars, vectors).\n\nThe authors also propose a novel attention mechanism which involves any dimensional irrep based on the tensor product. The authors test their models in the QM9, and OC20 datasets, and show that it achieves state-of-the-art performance in most tasks.\n", "questions": "It would have also been interesting to validate the model on some molecular dynamics datasets. For example, MD17 is nowadays a standard benchmarks for equivariant models. \n\nAlso, I think the relevant work section should belong in its entirety in the main text. I would suggest the authors to restructure the manuscript such that all the relevant information is present in the main text.\n", "limitations": "\nI wish the authors would discuss the limitation of their approach, for example whether the model becomes too parameter intensive for larger graphs. \n\nOverall, I think the current paper is a nice work, and with a more extensive experimental section, involving also vectorial/tensorial tasks (like force prediction tasks in MD17) could qualify to be accepted in a venue like NeurIPS.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper is clear and well written. The assumptions are clearly stated and the algorithm is clearly described. The paper addresses an interesting topic, that is, how to extend the inductive bias of equivariance to transformers. \n\nThe main weakness of the paper consists perhaps in the fact that the experimental section could have been more extensive and details. For instance, it would have been interesting to report the number of parameters for the various experiments as well as the training time, since it is often the case that transformers are quite heavy. Also, all the experiments have “scalar” tasks, and it would be very interesting how the model would perform for vectorial predictions. \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": {}}]

Author rebuttal:

Decision: Reject