Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs
Authors: Yi-Lun Liao
Tess Smidt
Venue: ICLR 2023 notable top 25% (2023)
Keywords: equivariant neural networks
graph neural networks
computational physics
transformer networks

Abstract: Despite their widespread success in various domains, Transformer networks have yet to perform well across datasets in the domain of 3D atomistic graphs such as molecules even when 3D-related inductive biases like translational invariance and rotational equivariance are considered. In this paper, we demonstrate that Transformers can generalize well to 3D atomistic graphs and 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). First, we propose a simple and effective architecture by only replacing original operations in Transformers with their equivariant counterparts and including tensor products. Using equivariant operations enables encoding equivariant information in channels of irreps features without complicating graph structures. With minimal modifications to Transformers, this architecture has already achieved strong empirical results. Second, we propose a novel attention mechanism called equivariant graph attention, which improves upon typical attention in Transformers through replacing dot product attention with multi-layer perceptron attention and including non-linear message passing. With these two innovations, Equiformer achieves competitive results to previous models on QM9, MD17 and OC20 datasets.

Published as a conference paper at ICLR 2023
EQUIFORMER: EQUIVARIANT GRAPH ATTENTION TRANSFORMER FOR 3D ATOMISTIC GRAPHS
Yi-Lun Liao, Tess Smidt
Massachusetts Institute of Technology
{ylliac, tsmi dt. @mit.
comp omi carchi www.cccttttttteeiiiiiiooe former
ABSTRACT
Despite their widespread success various domains, Transformer networks have yet to perform well across datasets in the domain of 3D atomistic graphs such as molecules even when 3D-related inductive biases like translational invariance and rotational equivariance are considered. In this paper, we demonstrate that Trans- formers can generalize well to 3D atomistic graphs and present Equiformer, a graph neural network leveraging the strength of Transformer architectures and incorporat- ing (E(3)E -qquivarrai features based on irreducible representations (irreps). First, we propose simple and effective architecture by only replacing original operations in Transformers with their equivariant counterparts and including tensor products. Using equivariant operations enables encoding equivariant information in channels of irreps features without complicating graph structures. With minimal modifications to Transformers, this architecture has already achieved strong empir- ical results. Second, we propose a novel attention mechanism called equivariant graph attention, which improves upon typical attention in Transformers through replacing dot product attention with multi-layer perceptron attention and including non-linear message passing. With these two innovations, Equiformer achieves competitive results to previous models on QM9, MD17 and OC20 datasets.
INTRODUCTION
Machine learned models can accelerate the prediction of quantum properties of atomistic systems like molecules by learning approximations of ab initio calculations (Gilmer et al., 2017; Zhang et al., 2018b; Jia et al., 2020; Gasteiger et al., 2020a; Batzner et al., 2022; Lu et al., 2021; Unke et al., 2021; Sriram et al., 2022; Rackers et al., 2023). In particular, graph neural networks (GNNs) have gained increasing popularity due to their performance. By modeling atomistic systems as graphs, GNNs naturally treat the set-like nature of collections of atoms, encode the interaction between atoms in node features and update the features by passing messages between nodes. One factor contributing to the success of neural networks is the ability to incorporate inductive biases that exploit the symmetry of data. Take convolutional neural networks (CNNs) for 2D images as an example: Patterns in images should be recognized regardless of their positions, which motivates the inductive bias of translational equivariance. As for atomistic graphs, where each atom has its coordinate in 3D Euclidean space, we consider inductive biases related to 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 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 former leverages geometric tensors like vectors for equivariant node features (Thomas et al., 2018 8: Weiler et al., 2018; Kondor al., 2018; Fuchs et al., 2020; Batzner et al., 2022; Brandstetter et al., 2022; Musaelian et al., 2022), and the latter augments graphs with invariant information such distances and angles extracted from 3D graphs (Schütt et al., 2017; Gasteiger et al., 2020b;a; Liu et al., 2022; Klicpera et al., 2021).
A parallel line of research focuses on applying Transformer networks (Vaswani et al., 2017) to other domains like computer vision (Carion al., 2020; Dosovitskiy et al., 2021; Touvron al., 2020) and graph (Dwivedi & Bresson, 2020; Kreuzer et al., 2021; Ying et al., 2021; Shi et al., 2022) and has demonstrated widespread success. However, as Transformers were developed for sequence data (Devlin et al., 2019; Baevski et al., 2020; Brown et al., 2020), it is crucial to incorporate domain- related inductive biases. For example, Vision Transformer (Dosovitskiy et al., 2021) shows that adopting a pure Transformer to image classification cannot generalize well and achieves worse results than CNNs when trained on only ImageNet (Russakovsky et al., 2015) since it lacks inductive biases like translational invariance. Note that ImageNet contains 1, 1.28M images and the size

---

Published as a conference paper at ICLR 2023
is already larger than that of many quantum properties prediction datasets (Ruddigkeit et al., 2012; Chmiela et al., 2017; Chanussot* et al., 2021). Therefore, this highlights the necessity of including correct inductive biases when applying Transformers to the domain of 3D atomistic graphs.
Despite their widespread success in various domains, Transformers have yet to perform well across datasets (Fuchs et al., 2020; Thölke & Fabritiss, 2022; Le et al., 2022) in the domain of 3D atomistic graphs even when relevant inductive biases are incorporated. In this work, we demonstrate that Transformers can generalize well to 3D atomistic graphs and present Equiformer, an equivariant graph neural network utilizing SE(3) (3)())))vaariam features built from irreducible representations (irreps) and a novel attention mechanism to combine the 3D-rellted inductive bias with the strength of Transformer. First, we propose a simple and effective architecture, Equiformer with dot product attention and linear message passing, by only replacing original operations in Transformers with their equivariant counterparts and including tensor products. Using equivariant operations enables encoding equivariant information in channels of irreps features without complicating graph structures. With minimal modifications to Transformers, this architecture has already achieved strong empirical results (Index 3 in Table 6 and 7). Second, we propose a novel attention mechanism called equivariant graph attention, which improves upon typical attention in Transformers through replacing dot product attention with multi-layer perceptron attention and including non-linear message passing. Combining these two innovations, Equiformer (Index 1 in Table 6 and 7) achieves competitive results on QM9 (Ruddigkeit et al.. 2012: Ramakrishnan et al., 2014), MD17 (Chmicla et al., 2017: Schütt et al., 2017; Chmiela et al., 201 8) and OC20 (Chanussot* et al., 2021) datasets. For QM9 and MD17, Equiformer achieves overall better results across all tasks or all molecules compared to previous models like NequlP (Batzner et al., 2022) and TorchMD-NET (Thölke & Fabritũi,, 2022). For OC20, when trained with IS2RE data and optionally IS2RS data, Equiformer improves upon state-of-the-art models such SEGNN (Brandstetter et al., 2022) and Graphormer (Shi et al., 2022). Particularly, as of the submission of this work, Equiformer achieves the best IS2RE result when only IS2RE and IS2RS data are used and improves training time by 2.3x to 15.5x compared to previous models.
2 RELATED WORKS
We focus on equivariant neural networks here. We provide detailed comparison between other equivariant Transformers and Equiformer and discuss other related works in Sec. B in appendix. SELuvaiaa GNNs. Equivariant neural networks (Thomas et al., 2018; Kondor et al., 2018; Weiler et al., 2018: 8: Fuchs et al., 2020: Miller al., 2020; Townshend et al., 2020; Batzner et al., 2022; Jing et al., 2021; Schütt et al., 2021; Satorras et al., 2021; Unke et al., 2021; Brandstetter et al., 2022: Thölke & Fabritiis, 2022; Le et al., 2022; Musaelian et al., 2022) operate on geometric tensors like type-L vectors to achieve equivariance. The central idea is to use functions of geom- ctry built from spherical harmonics and irreps features to achieve 3D rotational and translational equivariance as proposed in Tensor Field Network (TEN) (Thomas et al., 2018), which generalizes 2D counterparts (Worrall et al., 2016; Cohen & Welling, 2016; Cohen et al., 2018) to 3D Euclidean space (Thomas et al., 2018; Weiler et al., 2018; Kondor et al., 2018). Previous works differ in cquivariant operations used in their networks. TFN (Thomas et al., 2018) and NequlP (Batzner al., 2022) use graph convolution with linear messages, with the latter utilizing extra equivariant gate activations (Weiler et al., 2018). SEGNN (Brandstetter et al., 2022) introduces non-linear messages (Gilmer al., 2017: messages use the same gate activation and improve upon linear messages. SE(3)- Trans-Transormmer (Fuchs et al., 2020) adopts an equivariant version of dot product (DP) attention (Vaswani et al., 2017) with linear messages, and the attention can support vectors of any type I.. Subsequent works on equivariant Transformers (Thölke& & Fabritis, 2022; Le et al., 2022) follow the practice of DP attention and linear messages but use more specialized architectures considering only type-O and type-1 vectors.
et al., 2020) for irreps features, and the non-linear
The proposed Equiformer incorporates all the advantages through combining MLP attention with non-linear messages and supporting vectors of any type. Compared to TFN, NequlP, SEGNN and SE(3) Transorran,,,,,r,, the proposed combination of MI. attention and non-linear messages is more expressive than pure linear or non-linear messages and pure MLP or dot product attention. Compared to other equivariant Transformers (Thölke & Fabritis, 2022: Le al., 2022), in addition to being more expressive, the proposed attention mechanism can support vectors of higher degrees (types) and involve higher order tensor product interactions, which can lead to better performance.
3 BACKGROUND
3.1 E(3 EQUIVARIANCE
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
2

---

Published as a conference paper at ICLR 2023
translation, rotation and inversion (F- The groups of 3D translation, rotation and inversion form Euclidean group E33,, with the first two forming (3)3), the second being SO(3), and the last two forming (3)). The laws of physics are invariant to the choice of coordinate systems and therefore properties of atomistic systems are equivariant, e.g., when we rotate our coordinate system, quantities like energy remain the same while others like force rotate accordingly. Formally, function f mapping between vector spaces X and Y is equivariant to group of transformation G if for any input x E X, output y Y and group element g E G. we have fDxx (g)xx) = Dy (z), where Dx Dx(g) and Dy are transformation matrices parametrized by g in X and Y. For learning on 3D atomistic graphs, features and learnable functions should be -quiivaavr to geometric transformation acting position F. In this work, following previous works (Thomas et al., 2018: Kondor et al., 2018; Weiler al., 2018) implemented in e3nn (Geiger et al., 2022), we achieve by using equivariant features based on vector spaces of irreducible representations and equivariant operations for learnable functions. In the main text, we discuss S ( -euuevrrraaa and benchmark Equiformer with
We leave the discussion on inversion and equuvvvraaae in Sec. A and present results of -(3)-equivariance in Sec. D.2 and Sec. F.4 in appendix.
3.2 IRREDUCIBLE REPRESENTATIONS
A group representation (Dresselhaus et al., 2007: Zee, 2016) defines transformation matrices Dx (g) of group elements g that act 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 operate on relative positions. The transformation matrices of rotation and inversion are separable and commute. We discuss irreducible representations of (3) below and discuss inversion in Sec. A.3 in appendix.
Any group representation of SO(3) on given vector space can be decomposed into a concatenation of provably smallest transformation matrices called irreducible representations (irreps). Specifically, for group element g E SO(3). there are (2I+1)- (2L+1))bb--(2+11 irreps matrices Dz (g) called Wigner-D matrices acting on (2L + 1) 1)-dimensional vector spaces, where degree a non-negative integer. can be interpreted as an angular frequency and determines how quickly vectors change when rotating coordinate systems. Dí (g) of different L act on independent vector spaces. Vectors transformed by D2(9) are type- vectors, with scalars and Euclidean vectors being type-0 and type-1 vectors. It is common to index elements of type-L vectors with an index m called order, where -L S m S L.
Irreps Features. We concatenate multiple type-I. vectors to form -e((3--equivariann irreps features. Concretely, irreps feature / has Ct type-l. vectors, where 0 < T. < Tamas and C₁. is the number of channels for type-L vectors. We index irreps features f by channel c, degree L, and order m and denote as <(L). Different channels of type-L L vectors are parametrized by different weights but are transformed with the same (g) Regular scalar features correspond to only type-0 vectors.
Spherical Harmonics. Fuclidean vectors F'nn R33 can be projected into type-I. vectors ((4) by using SH are E33) (equ-vquiaa with (r(g)
spherical harmonics (SH) (4); f(L) = (L) |Fi
(L)
Di(g) ). SH of relative position rig generates the first set of irreps features. Equivariant [P
information propagates to other irreps features through equivariant operations like tensor products.
3.3 TENSOR PRODUCT
Tensor products can interact diffferent type-L vectors. We discuss tensor products for SO(3) below and those for ((3) in Sec. A.4. The tensor product denoted as @ uses Clebsch-Gorde coefficients to combine Lype-Li vector f(bn) and type-L2 vector g(db)) and produces type-L3 vector h (4a))
L
La
I -(rlt) 2 )
Σ
Σ
claama)
(lh g(La)
(Lamm)(Lamma))
mi
12=-L4
where my denotes order and refers to the mj-th element of f(Ln), Clebsch-Gordan coefficients ((....m))
(L) (L).m.m)))) are non-zero only when [L, L2l S La S [h + Lzl and thus restrict output vectors to be of certain types. For efficiency, discard vectors with L Lmmaz, where Lmaz is to prevent vectors of increasingly higher dimensions.
We call each distinct non trivial combination of L1 & Le L3 path. Each path is independently equivariant, and we can assign one learaalee weight to each path in tensor products, which is similar to typical linear layers. We can generalize Eq. to irreps features and include multiple channels
3

---

irreps features. The equivariant Published as a conference paper at ICLR 2023
Input Graph
Ingut 30 Graph
Aten Type ?
Constant One
D00
Emberding
DIP
Linew
Liny
Redtupe
Scale
Scalars
Impe fa
L Caiky RelE
Godc
Nece
(c) Emshekding (Atomi + Edge-DDgree)
Softeas
Lirest
Redhapa
Gatc
Network
Ougget
Linerr
Head
(x) Equiffrmer 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, "ø" denotes multiplication, "Ø" denotes addition, and "DTP" stands for depth-wisg -wise tensor product. Σ within circle denotes summation over all neighbors. Gray cells indicate intermediate irreps features.
of vectors of different types through iterating over all paths associated with channels of vectors, In this way, weights are indexed by (4 C2 C32 (z), where e is the e1-th channel of type-11 vector in input irreps feature. We use @ to represent tensor product with weights w. Weights can be conditioned on quantities like relative distances.
4 EQUIFORMER
First, we propose simple and effective architecture, Equiformer with dot product attention and linear message passing, by only replacing original operations in Transformers with their equivariant coun- terparts and including tensor products for S (3) operations are discussed in Sec. 4.1. The equivariant version of dot product attention can be found in Sec. C.3, and that of other modules Transformers can found in Sec. 4.3. Second, we propose a novel attention mechanism called equivariant graph attention in Sec. 4.2. The proposed Equiformer combines these two innovations and is illustrated in Fig. 1.
4.1 EQUIVARIANT OPERATIONS FOR IRREPS FEATURES
We discuss below equivariant operations, which serve building blocks for equivariant graph attention and other modules, and analyze how they remain equivariant in Sec. C.I. They include the equivariant version of operations in Transformers and depth-wise tensor products shown in Fig. 2.
Linear. Linear layers are generalized to irreps features by transforming different type- L vectors separately. Specifically, we apply separate linear operations each group type-L vectors. We remove bias terms for non-scalar features with I. 0 as biases do not depend on inputs, and therefore, including biases for type-T. vectors with I. > 0 can break equivariance.
Layer Normalization. Transformers adopt layer normalization (LN) (Ba et al., 2016) stabilize training. Given input x E R NxC, with N being the number of nodes and the number of channels, LN calculates the linear transformation of normalized input LN(x)
iー 0 7 + B, de
where μc. ac E Nx1 are mean and standard deviation of input i along the channel dimension, 7,8 E R are learnable parameters, and denotes element-wise product. By viewing standard deviation as the root mean square value (RMS) of L2-norm type-L vectors, LN can generalized to irreps features. Specifically, given input x XxCx(22 of type- vectors, the output is LN(x)= RMSSc< 7, where norm(x) NxCxI calculates the L2-norm of each type-L vectors in v, and RMS (norm(()) R calculates the RMS of L2-norm with mean taken along the channel dimension. We remove means and biases for type-L vectors with L * O.
4

---

Published as a conference paper at ICLR 2023
Input immps features
Ourput imreps festures
Typo
Sabane
Dosde
SILU
Signoid
Type-
vesues
RMS
(a) Linear
(b) Layer Normalization
(c) Gate
(d) Depth-wise Teasor 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 y here. (c) Gate layers cquivariant activation functions where non-linearly transformed scalars are used to gate non-scalar irreps features. (d) The left two irreps features correspond two input irreps features, and the rightmost one is the output irreps feature. The two gray lines connecting two vectors in the input irreps features and one vector in the output irreps feature form a path and contain one learnable weight. An alternative visualization of depth- Septh-wise tensor products can be found in Fig. 3 in appendix. show S (2)3)-equivarian operations here, which can be generalized to -quuivaaiaa features.
Gate. We use the gate activation (Weiler et al., 2018) for equivariant activation function as shown in Fig. 2(c). Typical activation functions are applied to type-0 vectors. For vectors higher L, we multiply them with non-linenn non-linearly transformed type-0 vectors for equivariance. Specifically, given input containing non-scalar Cz type-L vectors with 0 L - Lmaz and (Co + - CL) type-0 vectors, we apply SiLU (Elfwing et al., 2017; Ramachandran et al., 2017) to the first Co type-O vectors and sigmoid function to the other Σ Eimns CL 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-0 vectors is reduced Co-
Dist Tensor Product. The tensor product defines interaction between vectors of different L. To improve its efficiency, we use the depth-wise tensor product (DTP), where one type-L vector in output irreps features depends only on one type- vector in input irreps features as illustrated in Fig. 2(d) and Fig. 3, with L being equal to or different from L'. This is similar to depth-wise convolution (Howard al., 2017), where one output channel depends on only one input channel. Weights w in the DTP can be input-independent or conditioned on relative distances, and the DTP between two tensors x and y is denoted as x @ DTP y. Note that the one-to-one dependence of channels can significantly reduce the number of weights and thus memory complexity when weights are conditioned on relative distances. contrast, if one output channel depends on all input channels, in our case, this can lead to out-of-memory errors when weights are parametrized by relative distances.
4.2 EQUIVARIANT GRAPH ATTENTION
Self-attention (Vaswani et al., 2017: Veličković et al., 2018; Fuchs et al., 2020; Khan et al., 2021; Ying et al., 2021; Brody et al., 2022) transforms features sent from one spatial location to another with input-dependent weights. We use the notion from Transformers (Vaswani et al., 2017) and message passing networks (Gilmer et al., 2017) and define message mij sent from node j to node i as follows:
mij = aij X Ukj
where attention weights aiy depend on features on node and its neighbors Mii) and values Bij are transformed with input-independent weights. In Transformers and Graph Attention Networks (GAT) (Veličković et al., 2018; Brody et al., 2022), Vij depends only on node j. In message passing networks (Gilmer et al., 2017), dij depends on features nodes and with constant dij. The proposed equivariant graph attention adopts tensor products to incorporate content and geometric information and uses multi-layer perceptron attention for aı and non-lineat message passing for Us as illustrated in Fig. 1(b).
Incorporating Content and Geometrie Information. Given features z; and H, target node i and source node j, we combine the two features with two linear layers to obtain initial message Ta sineards((x;) + s Tiy is passed DTP layer and linear layer to consider
5

---

Published as a conference paper at ICLR 2023
geometric information like relative position contained in different type-L vectors in irreps features: Sw(
z = 24j @ W
(Fy SH(ri)) and fiz = Linear(11l))
(3)
where diy is the tensor product of may and spherical harmonics embeddings (SH) of relative position Fij. with weights parametrized by || (- fiz considers semantic and geometric features on source and target nodes a linear manner and is used to derive attention weights and non-linear messages.
Multi-Layer Perceptron Attention. Attention weights atj capture how each node interacts with neighhoring nodes, dij are invariant to geometric transformation, and thus, we only use type-0 vectors (scalars) of message Jij denoted (/g]) for attention. Note that 5! encodes directional information, as they are generated by tensor products of type-L vectors with 2 0. Inspired by GATv2 (Brody et al., 2022), we adopts multi-layer perceptron attention (MLPA) instead of dot product attention (DPA) used in Transformers (Vaswani et al., 2017). In contrast to dot product, ML.ss are universal approximators (Hornik et al., 1989; Hornik, 1991; Cybenko, 1989) and can theoretically capture any attention patterns. Given ((0)) we uses one leaky ReLU Rayur and one linear layer for a.y:
exp(zig)
Zij ) and aij softmax, ELENTO (xp(zit)
where is a learnable vectors of the same dimension as fij and 2zj is a single scalar. The output of attention is the sum of value fi multipled by corresponding aej over all neighboring nodes E (()), where Vay can be obtained by linear or non-linear transformations of fij as discussed below.
Non-Linear Message Passing. Values Vij are features sent from node to another, transformed
with input-independent weights. We first split fúj into p(L) and f where the former consists of type-L, vectors with 0 S L < Lmmar and the latter consists of scalars only. Then, we perform
non-linear transformation to fig to obtain non-linear message:
Maj = Gate( and
Mij = Linear([uug DTP
ghrP (H(F))))
(5)
We apply gate activation to fi] obtain Piy. We use one DTP and linear layer to enable interaction between non-linear type-L. vectors, which is similar to how we transform Ziy into fij- Weights
here are input-independent. We can also use ((.)) directly as viy for linear messages.
Multi-Head Attention. Following Transformers (Vaswani et al., 2017), we can perform h parallel equivariant graph attention functions given fej. The h different outputs concatenated and projected with a linear layer, resulting in the final output y as illustrated in Fig. 1(b). Note that parallelizing attention functions and concatenating can be implemented with "Reshape".
4.3 OVERALL ARCUITECTRRE
For completeness, discuss other modules in Equiformer here.
Embedding. 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 (Xu et al., 2019; Shi et al., 2022). The DTT 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 so that standard deviation of aggregated features would be close to 1.
Radial Basis and Radial Function. Relative distances ||a parametrize weights in some DTP layers. To reflect subtle changes in Wigll. we represent distances with radial basis like Gaussian radial basis (Schütt et al., 2017) and radial Bessel basis (Gasteiger et al., 2020b;a). We transform radial basis with a learnable radial function to generate weights for those DTP layers. The function consists of two-layer MLP, with each linear layer followed by LN and SiLU, and a final linear layer.
Feed Forward Network. Similar to Transformers, we use two equivariant linear layers and an intermediate gate activation for the feed forward networks in Equiformer.
Output Head. The last feed forward network transforms features on each node into a scalar. We perform sum aggregation over all nodes to predict sealar quantities like energy. Similar to edge-degree embedding, we divide the aggregated scalars with the squared root of average numbers of atoms.

---

Published as a conference paper at ICLR 2023
Tiskk Latts
Air
FLUMO meN
P
G
U/. meN
ZENEE
Methods
%
meVV
mcV
caltmol
meV
meV
meV
me/
NMP (Gilmer al., 2017)"
4192
69)
5%
1.50
SchNet
2017)
235
1.70
Cormorant (Andemon et al., 2019)'
.085
61
34
38
.038
.026
.961
21
22
2,03
LieConv (Finci al., 2020)*
.064
49
30
25
0032
033
22
24
800
19
19
2.28
DimeNete+ (Ciaseegger al., 201200
.044
33
25
20
0030
J025
.331
1:21
(Thomas al. 2018)
.223
58
SEaS--Tamffemme (Fuchs 2020y!
J42
53
33
1051
EGNN (Satoras al. 2021)"
.071
48
70
25
1020
0331
12
12
106
12
11
1.55
PaiNN (Schünt et al., 2021)
.015
46
28
20
.012
.024
7.35
5.98
.066
5,83
5.85
1.28
TurchMD- MiEiMMD-EETTTI (Thilke Fubeitsis 2022)
(059
36
20
18
7.62
6.16
.033
6.38
6.15
1.84
SphereNet (Liu al., 2022)
.466
32
18
1.12
SEGNN (Brandstetter
2022)?
880
42
21
1.62
EQGAT (Le et al., 2022)
.053
32
20
16
.011
.024
23
.382
25
25
2.00
Equiformer
.166
30
15
14
.011
.023
7.63
6.63
.251
6.74
6.59
1.26
Table I: MAE results on QM9 testing set. f denotes using different data partitions.
Aspirin
Benoene
Eihanol
Molonaidehyde
Naphihulene
Salieylic scid
Toluene
Unxill
Methods
Foroes
energy
forces
forces
energy
focces
energy
fonces
energy
forces
energy
forces
enerpy
faroes
SchNet (Schilt al., 2017)
16.0
58.5
3.5
16.9
5.6
6.9
25.2
36.9
5.2
24.7
24.3
DemeNet (Gasteiger al. 2020b)
21.6
3.4
10:0
933
13.1
PaiNN (Schin 2021)
14.7
9.7
500
4.1
6.0
TorchMD NET [Thollke& Fabritis, 2022)
11.0
2.5
8.5
2.3
4.7
3.3
37
2.6
4,0
5.6
3.2
2.9
4.1
NequlP (haus 3) (Batzeer al.. 2022)
5.7
8.0
2,2
3.1
3.3
5.6
4.9
1.7
4.6
3.9
4.0
2.0
4.5
333
Equilormer Uonss 2)
5.3
7.2
22
6.6
2.2
3.1
3.3
5.8
3.7
2.1
4.5
4.1
3.8
2.1
4.3
3,3
Equilormer (F
53
6.6
2.5
8.1
2.2
29
3.2
54
4.4
2.0
4.3
3,9
3.7
2.1
4.3
3,4
Table 2: MAE results on MD17 testing set. Energy and force are in units of meV and meV/
5 EXPERIMENT
We benchmark Equiformer on QM9 (Sec. 5.1), MD17 (Sec. 5.2) and OC20 (Sec. 5.3) datasets. Moreover, ablation studies (Sec. 5.4) are conducted to demonstrate that Equiformer with dot prodcut attention and linear message passing has already achieved strong empirical results on QM9 and OC20 datasets and verify that the proposed equivariant graph attention improves upon typical dot product attention in Transformer as well as dot product attention in other equivariant Transformers. Additional results of including inversion can be found in Sec. D.2 and Sec. F44
5.1 QM9
Dataset. The QM9 dataset (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014) (CC BY-NC SA 4.0 license) consists of 134k small molecules, and the goal is to predict their quantum properties. The data partition we use has 110k, 10k, and 1kk molecules in training. validation and testing sets. We minimize mean absolute error (MAE) between prediction and normalized ground truth.
Training Details. Please refer to Sec. D.I in appendix for details architecture, hyper-parameters and training time.
Results. We summarize the comparison to previous models in Table 1. Equiformer achieves overall better results across 12 regression tasks compared to each individual model. The comparison to SEGNN, which uses irreps features as Equiformer, demonstrates the effectiveness of combining non-lienar message passing with MLP attention. Additionally, Equiformer achieves better results for most tasks when compared to other equivariant Transformers. which are SE(3)- Transformer, TorchMD-NET and EQGAT. This demonstrates better adaption of Transformers to 3D graphs and the effectiveness of the proposed equivariant graph attention. We note that for the tasks of μa and R², PaiNN and TorThMD-NET use different architectures, which take into account the property of the task. In contrast, we use the same architecture for all tasks. We compare training time in Sec. D.3.
5.2 MD17
Dataset. The MD17 dataset (Chmiela et al., 2017; Schütt al., 2017; Chmiela et al., 2018) (CC BY-NC) consists of molecular dynamics simulations of small organic molecules, and the goal is to predict their energy and forces. We use 950 and 50 different configurations for training and validation sets and the rest for the testing set. Forces are derived as the negative gradient of energy with respect to atomic positions. We minimize MAE between prediction and normalized ground truth.
Training Details. Please refer to Sec. E. in appendix for details on architecture, papaeepas and training time.
Results. We train Equiformer with Lmaz = and and summarize the results in Table 2. Equiformer achieves overall better results across 8 molecules compared to each individual model. Compared to TorchMD-NV which is also an equivariant Transformer, the difference lies in the pro- posed equivariant graph attention, which is more expressive and can support vectors of higher degree
7

---

Published as a conference paper at ICLR 2023
Energy (GV)
(5)
Methods
1D
QOD Ads
0OD Cat
OOD Both
Averaee
ID
000 Ads
OOD Cat
OOD Both
Average
CGCNN OXi & Grussman, 2018)
0.5149
0.9155
0.6219
0.8511
0.7509
3.40
1.93
3,10
2,00
2.61
SchNet (Schilt et al. 2017)
0.6387
0.7542
0.6616
0.7637
0.6846
2.96
2.33
2.94
2.21
2.61
DimeNa++ (Gasteiger al. 2020s)
0.5621
0,7252
0.5756
0.6613
0.611
4.25
2,07
4.10
2.41
3,21
PaiNN (Schurt
2021)
0.783
0.604
0.663
1,97
3,46
2.79
SpinCony
2021)
0.7230
0,5687
0.6738
0.6310
2.26
3,82
2.33
3.12
SpliereNet (1.iu et 2022)
0.5625
0.7033
0.5708
0.6378
0.6166
2.29
4.09
2.41
3.32
SEGNN (Brmdstetter el al., 2022)
0.5327
0.6921
0.5369
0.6790
0.6101
5.37
2.46
4.91
2.63
3.84
Equifocmer
0.5037
0.6881
0.5213
0.6301
0.5858
5.14
2,41
4.67
2.69
3.73
Table 3: Results on OC20 IS2RE testing set.
Energy (eV)
EWI (8)
Methock
TD
ODD Ads
0OD Cat
ODD Both
Avenyr
ID
OOD Adk
OOD Cat
OOD Both
Average
GNS (Giudwin sl., 2022)
0.54
0.65
0.55
0.59
0.5825
GNS Nousy (Godwin al., 2022)
0.47
0.51
0.46
0,8800
Graphommer (Shi al., 2022)
0:4329
0.5850
0.4111
0.5299
0.4980
Equiformer
0.4222
0.5420
0.4231
0.7544
0.4657
7:23
3.77
7.13
4.10
5.56
Equilormer+ + Nuisy Nodes
0.4156
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 is adopted during training.
Enery MAE (eVi
EWI (96)
Training time
Mothods
ID
OO1 Ads
OOD Car
OOD Both
Average
11
001 Ads
OUD Cat
OOD Hoth
Average
(GPU-dayy)
Noisy Nodes (Goodwin 2022)
0.4219
0.5678
0.4306
0,4651
0.4728
4.25
8.01
6.5
(IPU)
Graphormer (Shá et al., 2022)'
0.3976
0.5719
0.4166
0.5029
0,4222
3.45
3.79
6.1
(72()100)
Equiformer- Noisy Nodes
0.4171
0.547V
0,4248
0.441
0.4660
7.71
3.70
7.15
4.07
5,66
(A6000)
Table 5: Results on OC20 IS2RE testing set when IS2RS is adopted during training. f denotes using ensemble of models trained on both IS2RE training and validation sets. In contrast, we use the same single Equiformer model in Table 4, which is trained only on the training set. Note that Equiformer achieves better results with much less computation.
I. (i.e., we use Tmmaz = and 3) instead of restricting to type-0 and type-1 vectors (i.e., Tmaz = 1). For the last three molecules, although Equiformer with Lmaz 2 achieves lower force MAE but higher energy MAE, we can adjust the weights of energy loss and force loss so that Equiformer achieves lower MAE for both energy and forces as shown in Table 12 in appendix. Compared to NequlP, which also uses irreps features and Lmaz = 3, Equiformer with Lmazz 2 achieves overall lower MAE although including higher Linex can improve performance. When using Lmaa = 3, Equiformer achieves lower MAE results for most molecules. This suggests that the proposed attention can improve upon linear messages even when the size of training sets becomes small. We compare the training time of NequlP and Equiformer in Sec. E.3. Additionally, for Equiformer, increasing Lmaz from 3 improves MAE for most molecules except benzene, which results from overfitting.
5.3 OC20
Dataset. The Open Catalyst 2020 (OC20) dataset (Chanussot et al., 2021) (Creative Commons Attribution 4.0 License) consists of larger atomic systems, each composed of molecule called adsorbate placed on a slab called catalyst. Each imput contains more atoms and more diverse atom types than QM9 and MD17. We focus the task of initial structure relaxed energy (IS2RE), which is to predict the energy of a relaxed structure (RS) given its initial structure (IS). Performance is measured in MAE and energy within threshold (Ew'T), 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-distribttion adsorbates (OOD-Ads), out-of-dio -ff-of---ttriutiinn catalysts (OOD-Cat), and cut-of-distribution adsorbates and catalysts (OOD-Both) Please refer to Sec. F.1 for the detailed description of OC20 dataset.
Setting. We consider two training settings based on whether node-level auxiliary task (Godwin et al., 2022) 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 a node-level auxiliary task.
Training Details. Please refer to Sec. F.2 in appendix for details on Equiformer architecture, hyper-parameters and training time.
IS2RE Results without Node-Level Auxiliary Task. We summarize the results on validation and testing sets under the first setting in Table 15 in appendix and Table 3. Compared with state-of-the- art models like SEGNN and SphereNet, Equiformer consistently achieves the lowest MAE for all the four sub-splits in validation and testing sets. Note that EwT considers only the percentage of predictions close enough to ground truth and the distribution of errors, and therefore improvement in average emros (MAE) would not necessarily reflect that in error distributions (EwTT). A more detailed
8

---

Published as a conference paper at ICLR 2023
Methods
Non lincar message passing
MLP atteniinn
Dot product
Task
4.
Ar
FLLND meV
C2 calimol K
Training time (minutes:ppuch)
Number of parameters
Index
attention
Unit
6]
meV
meV
D
0116
30
15
14
0011
.023
12.1
3.53M
051
32
16
7.2
501M
053
16
3.35M
Table 6: Ablation study results on QM9.
Mothods
Energy (eN)4
EWT (1811
Non-lincar
MLP
Dut product
Training time (minaee/epoch))
Number of parameters
Index
messuge passing
amcotion
attesziinn
113
OOD Ads
DOL Car
OOD Both
Average
113
GDD Ads
0OD Cat
OOD Both
Average
0,5088
0.6271
0.5051
0.5545
0,5489
4.88
2.95
4.92
2.98
3,93
130.8
9:12M
0.5168
0,6308
0.5088
0.5657
0.5555
4.59
2,82
4.79
2.02
3.81
91.2
7.84M
0.5586
0,6382
0.5692
2,60
2,86
8.72M
Table 7: Ablation study results on OC20 IS2RE validation set.
discussion can be found in Sec. F.6 in appendix. We also note that models are trained by minimizing MAE, and therefore comparing MAE could mitigate the discrepancy between training objectives and evaluation metrics and that OC20 leaderboard ranks the relative performance according to MAE. Additionally, we compare the training time of SEGNN and Equiformer in Sec. F.5.
IS2RE Results with IS2RS Node-Level Auxiliary Task. We report the results on validation and testing sets in Table 4 and Table 5. As of the date of the submission of this work, Equiformer achieves the best results on IS2RE task when only IS2RE and IS2RS data are used. Notably, the result in Table 5 is achieved with much less computation. We note that under this setting, greater depths and thus more computation translate to better performance (Godwin et al., 2022) and that Equiformer demonstrates incorporating equívariant features and the proposed equivariant graph attention can improve training efficiency by 2.3x to 15.5x compared to invariant message passing networks and invariant Transformers.
5.4 ABLATION STUDY
We conduct ablation studies to show that Equiformer with dot product attention and linear message passing has already achieved strong empirical results and demonstrate the improvement brought by MLP attention and non-linear messages in the proposed equivariant graph attention. Dot product (DP) attention only differs from MLP attention in how attention weights dzy are generated from fij- Please refer to Sec. C.3 in appendix for further details. For experiments on QM9 and OC20, unless otherwise stated, we follow the hyper-parameters used in previous experiments.
Result on QM9. The comparison is summarized in Table 6. Compared with models in Table 1, Equiformer with dot product attention and linear message passing (Index 3) achieves competitve results. Non-linear messages improve upon linear messages when MLP attention is used while non-linear messages increase the number of tensor product operations in each block from to 2 and thus inevitably increase training time. On the other hand, MLP attention achieves similar results to DP attention. We conjecture that DP attention with linear operations expressive enough to capture common attention patterns as the numbers of nighboring nodes and atom species are much smaller than those in OC20. However, MLP attention is roughly 8% faster as it directly generates sealar features and attention weights from fúj instead of producing additional key and query irreps features for attention weights.
Result on OC20. We consider the setting of training without auxiliary task and summarize the comparison in Table 7. Compared with models in Table 15, Equiformer with dot product attention and linear message passing (Index 3) has already outperformed all previous models. Non-linear messages consistently improve upon linear messages. In contrast to the results on QM9, MLP attention achieves better performance than DP attention and is 8% faster. We surmise this because OC20 contains larger atomistic graphs with more diverse atom species and therefore requires more expressive attention mechanisms. Note that Equiformer can potentially improve upon previous equivariant Transformers (Fuchs et al., 2020; Thölke & Fabritiis, 2022; Le et al., 2022) since they use less expressive attention mechanisms similar to Index 3 in Table 7.
6 CONCLUSION
In this work, we propose Equiformer, 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 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, MD17 and OC20 demonstrate the effectiveness of Equiformer and ablation studies show the improvement of the proposed equivariant graph attention over typical attention in Transformers.
9

---

Published as a conference paper at ICLR 2023
7 ETHICS STATEMENT
Equiformer achieves more accurate approximations of quantum properties calculation. We believe there much more be gained by harnessing these abilities for productive investigation of molecules and materials relevant to application such as energy, electronics, and pharmaceuticals, 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 or harmful molecule to its large-scale deployment.
Moreover, we discuss several limitations of Equiformer and the proposed equivariant graph attention in Sec. G in appendix.
REPRODUCIBILITY STATEMENT
We include details on architectures, hyper-parameters and training time Sec. D.I (QM9), Sec. E,1 (MD17) and Sec. F.2 (OC20).
The code for reproducing the results of Equiformer on QM9, MD17 and OC20 datasets is available at https://githhe com/ http://accccccctttttte/eeouooooa
ACKNOWLEDGEMENT
We thank Simon Batzner, Albert Musaclian, Mario Geiger, Johannes Brandstetter, and Rob Hesselink for helpful discussions including help with the OC20 dataset. We also thank the e3nn (Geiger et al., 2022) developers and community for the library and detailed documentation. We acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Center (Reuther et al., 2018) for providing high performance computing and consultation resources that have contributed to the research results reported within this paper.
Yi-Lun Liao and Tess Smidt were supported by DOE ICDI grant DE-SCC0222115.
REFERENCES
Brandon Anderson, Truong-Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. In Conference on Neural Information Processing (NeurlPS), 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization. arxiv preprint 160.v:660.06644 2016.
Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A framework for self-supervised learning of speech representations. Conference on Neural Information Processing (NeurIPS). 2020.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Ko- rnbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13 doi: 10.10388441467-022-299933 s41467- 022 -29-29939-222--
(1), May 2022.
URL https: http:/oo.ooo...11/ 1038//
Johannes Brandstetter, Rob Hesselink, Elise van der Pol, Erik Bekkers, and Max Welling. Geometric and physical quantities improve e(3) equivariant message passing. In International Conference on Learning Representations (ICLR), 2022. URL https::///// http:::tt////eeet /ttp::::///peeeeieew... net/ wtt.::/weeeeeeeeeeeeeeemee
xwr8g0BeV1.
Shaked Brody, Uri Alon. and Eran Yahav. How attentive are graph attention networks? In Interna- tional Conference Learning Representations (ICLR), 2022. URL httpp::/ net forum? net/forum? 1d=772x msx7c1.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel
10

---

Published as a conference paper at ICLR 2023
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 Berner, Sam McCandlish, Alec Radford, llya Sutskever, and Dario Amodei. Language models are few-shot learners. In Conference on Neural Information Processing (NeurIPS), 2020.
Dan Busbridge, Dane Sherburn, Pietro Cavallo, and Nils Hammerla. Relational graph attention networks. arxiv preprint 190.008804.08// 2019.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zogoruyko. End-to-end object detection with transformers. In European Conference on Computer Vision (ECCV), 2020.
Lowik Chanussot*, Abhishek 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. doi: 10.102 /atcccaaaall.0004555..
Zhengdao Chen, Lisha Li, and Joan Bruna. Supervised community detection with line graph neural networks. In International Conference Learning Representations (ICLR), 2019. URL battps://openzeeee/
net forum? berum?id-tie -
Stefan Chmiela, Alexandre Tkatchenko, Huziel E. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller. Machine learning accurate energy-conserving molecular force fields. Science Advances, 3(5):e 2017. doi: 10. 10.1116/sciadv.16000 1603011. URL https:/ /www. sci sccience.org do /abs/ / sc adv. 603015.
Stefan Chmiela, Huziel E. Sauceda, Klaus-Rotee Robert Müller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-learee force fields. Nature Communications, 9(1),
sep 2018. doi: (1138//41467-018-061
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International Conference on Machine Learning (ICML),2015
Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling. Spherical CNNs. In International
Conference on Learning Representations (ICLR), 2018. URL httpts//ppeeeeeee..... forum? id=#kbd5x2Rbb
George V. Cybenko. Approximation by superpositions of sigmoidal funetion. Mathematics of Control, Signals and Systems, 2:303-3144 1989.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arxiv preprint arxiv:1810.04805, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representations (ICLR), 2021. URL https: // //openreview.ro oot
Mildred S Dresselhaus, Gene Dresselhaus, and Ado Jorio. Group theory. Springer, Berlin, Germany, 2008 edition, March 2007.
Vijay Prakash Dwivedi and Xavier Bresson. generalization of transformer networks to graphs. 2020.
Stefan Elfwing, Eiji Uchibe, and Kenji Doya. Siggoid- linear units for neural network function approximation in reinforcement learning. arXiv preprint 17Xiv:1702.03) 18, 2017.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. arXiv preprint arXiv:2002. / /2880, 2020.
11

---

Published as a conference paper at ICLR 2023
Nathan Frey, Ryan Soklaski, Simon Axelrod, Siddharth Samsi, Rafael Gomez-Bombarelle Connor Coley and Vijay Gadepally, Neural scaling of deep chemical models. ChemRxiv, 2022. doi: 200.64444cchmmmxi222222--23222
Fabian Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. Seaa)-transoormmrs: 3d roto- translation equivariant attention networks. In Conference on Neural Information Processing (NeurIPS), 2020.
Hongyang Gao and Shuiwang Ji. Graph representation learning via hard and channel-wise attention networks. Proceedings of the 25th ACM SIGKDD International Conference Knowledge Discovery & Data Mining, 2019.
Johannes Gasteiger, Shankari Giri, Johannes 'T. Margraf, and Stephan Günnemann. Fast and uncertainty-aware directional message passing for non-equilibrim molecules. In Machine Learn- ing for Molecules Workshop, NeurIPS, 2020a.
Johannes Gasteiger, Janek GroB, and Stephan Günnemann. Directional message passing for molecular graphs. International Conference on Learning Representations (ICLR), 2020b.
Johannes Gasteiger, Muhammed Shuaibi, Anuroop Sriram, Stephan Günnemann, Zachary Ulissi, C. Lawrence Zitnick, and Abhishek Das. How do graph networks generalize to large and diverse molecular systems? arxiv preprint 22x4:2200..02782 2022.
Mario Geiger, Tess Smidt, Alby M., Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice, Kos- tiantyn Lapchevskyi, Maurice Weiler, Michat Tyszkiewicz, Simon Batzner, Dylan Madisetti, Martin Uhrin, Jes Frellsen, Nuri Jung, Sophia Sanborn, Mingjian Wen, Josh Rackers, Marcel Rød, and Michael Bailey. e3nn/e3nn: 2022-04-13, 03, April 2022. URL https: /ttii//oiiio/////// orgg 5281/ zenodo. 6459381.
Justin Gilmer, Samuel Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. In International Conference Machine Learning (ICML), 2017.
Jonathan Godwin, Michael Schaarschmidt, Alexander L 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 Repre- sentations (ICLR), 2022. URL https::///// httplp:::///ppereeeeee /wwp/pppeieeeeeee/o/////pe net wtt//oorrmmimo d=1
William L. Hamilton, Rex Ying, and Jure Leskovec. Inductive representation learning on large graphs. In Conference Neural Information Processing (NeurIPS), 2017.
Kurt Hornik.
Approximation capabilities of multilayer feedforward networks. Neural Net-
works, 4(2):251-257, 1991. ISSN 0893-6080. doi: http:///oi.orr//000011//99900000 90009-T. URL https:/ /tops///www.. www.enncaiieeett.c////eeeeeee//e
0893608091900091.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neural Networks, 2(5):359-366, 1989. ISSN 0893-6080. doi: https:
URL rttt://eeeccceeeeeeeeee.cccc com/
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/1704.04861, 2017.
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 arXiv:2005.00633 2020.
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-Isc: A large-scale challenge for machine learning graphs. arXiv preprint arXiv: 10o.022333999300 2021.
12

---

Published as a conference paper at ICLR 2023
Weile Jia, Han Wang, Mohan Chen, Denghui Lu, Lin Lin, Roberto Car, Weinan E. and Linfeng Zhang. Pushing the limit molecular 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.
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. URL https: https////pppeeveeeeeeeeeeeeeeeee net http://tt/eo/tittte///// forum?id-
1YLJDVSx6J4.
Salman Khan, Muzammal Naseer, Munawar Hayat, Syed Waqas Zamir, Fahad Shahbaz Khan, and Mubarak Shah. Transformers in vision: A survey, 2021.
Dongkwan Kim and Alice Oh. How to find your friendly neighborhood: Graph attention design with selffsupervision. In International Conference Learning Representations (ICLR), 2021. URL
.nt/ forum?id=w// 2orumm21d-1115KUN1gg$33
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional networks. In International Conference on Learning Representations (ICLR), 2017.
Johannes Klicpera, Florian Becker, and Stephan Günnemann. Gemnet: Universal directional graph neural networks for molecules. In Conference Neural Information Processing (NeurIPS), 2021.
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch Clebsch-gordan nets: a fully fourier space spherical convolutional neural network. In Advances in Neural Information Processing Systems 32, pp. 10117-10126, 2018.
Devin Kreuzer, Dominique Beaini, William L. Hamilton, Vincent Létourneau, and Prudencio Tossou. Rethinking graph transformers with spectral attention. In Conference on Neural Information Pro- cessing (NeurIPS), 2021. URL ht https: wpw.revieeeeeeeeeeeeeeeeeeeeeeee .nt hu Tj4yG.
Tuan Le, Frank Noé, and Djork-Arné Clevert. Equivariant graph attention networks for molecular property prediction. arXiv preprint arXiv:2202.0989) 2022.
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. URL https:://///s.... openreview net. babrumiid=ge VsRXsOt 9r.
Denghui Lu, Han Wang, Mohan Chen, Lin Lin, Roberto Car, Weinan E, Weile Jia, and Linfeng Zhang, 86 pflops deep potential molecular dynamics simulation of 100 million atoms with ab initio accuracy. Computer Physics Communications, 159:14 2021. ISSN 0010-4655. doi: https://doi. /tg//o0..0166f....220....0
URL https: /www. sciencedi com/ science/ atticcle/ps1/00004555500000
Benjamin Kurt Miller, Mario Geiger, Tess E. Smidt, and Frank Noé. Relevance of rotationally equivariant convolutions for predicting molecular properties. arxiv preprint arxiv:2008.0844 2020.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron J. Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. arxiv preprint arxiv:2204.05249, 2022.
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, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett (eds.). Advances in Neural Information Processing Systems 32, pp. 0024-8035. Curran Associates, Inc., 2019. URL http: httptt::ppppp.....ppeee www.appers..eeeiiiee neurips cc/ hacp/pap////////// 1015-pytorch-- an- pdf.
13

---

Published as a conference paper at ICLR 2023
Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and Thomas F. Miller. OrbNet: Deep learning for quantum chemistry using coametty-adaapted atomic-orbital features. The Journal of Chemical Physics, 153(12):124111, sep 2020. doi: 10. N0.1063/5.0021955 URL https::////.......t.m / org / 10. 063 &2F5. 0021 955.
Joshua A Rackers, Lucas Tecot, Mario Geiger, and Tess E Smidt. A recipe for cracking the quantum scaling limit with machine learned electron densities. Machine Learning: Science and Technology, 502), feb 2023. doi: 10.
URL https:: : /ttp//xxx////oo doi .ww.doo/ooooooooooooo
1088/2632 - acb314.
Prajit Ramachandran, Barret Zoph, and Quoe V. Le. Searching for activation functions, arXiv preprint arXiv:1710.05941, 2017.
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1, 2014.
Bharath Ramsundar, Peter Eastman, Patrick Walters, Vijay Pande, Karl Leswing, and Zhenqin Wu.
Deep Learning for the Life Sciences. O' O'Reilly Media, 2019, https:/ //tos:://wwww /wwww ama zon com/ Leep--eerning Li fe See-SSciencese croscopy
/ttcoooyy/ppll/ /p/ 492039837.
Albert Reuther. Jeremy Kepner, Chansup Byun, Siddharth Samsi, William Arcand, David Bestor, Bill Bergeron, Vijay Gadepally, Michael Houle, Matthew Hubbell, Michael Jones, Anna Klein, Lauren Milechin, Julia Mullen, Andrew Prout, Antonio Rosa, Charles Yee, and Peter Michaleas. Interactive supercomputing 40,000 cores for machine learning and data analysis. In 2018 IEEE High Performance extreme Computing Conference (HPEC), pp. 1-6. IEEE, 2018.
Yu Rong, Yatao Bian, Tingyang Xu, Weiyang Xie, Ying WEI, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. In Conference on Neural Information Processing (NeurIPS), 2020.
Lars Ruddigkeit, Ruud van Deursen, Lorenz 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, (2:12))2 2875, 2012. doi: 10.1021/ci300415d URL fttps://dof: org/ 10 1021 /110//133001551 PMID: 23088335.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115 (3):211-2522 2015. doi: 10. 1263-015-0816- 16-y.
Alvaro Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying. Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning (ICMI., 2020.
Víctor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In International Conference on Machine Learning (ICML). 2021.
K. T. Schütt, P.-. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, and K.-R. Müller. Schnet: A continuous-filter convolutional neural network for modeling quantum interactions. In Conference on Neural Information Processing (NeurIPS). 2017.
Kristof T. Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R. Müller, and Alexandre Tkatchenko. Quantum-chemical insights from deep tensor neural networks. Nature Communications, 8(1), jan 2017. doi: 10. 10.1008/ncomms 3890..
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.
Yu Shi, Shuxin Zheng, Guolin Ke, Yifei Shen, Jiacheng You, Jiyan He, Shengjie Luo, Chang Liu, Di He, and Tie- Yan Liu. Benchmarking graphormer on large-scale molecular modeling datasets. arxiv preprint arxiv:220..04811 2022.
14

---

Published as a conference paper at ICLR 2023
Yunsheng Shi, Zhengjie Huang, Shikun Feng, Hui Zhong, Wenjin Wang. and Yu Sun. Masked label prediction: Unified message passing model for semi-supervised classification. arxiv preprint arxiv:2009.0333 03500, 2020.
Muhammed Shuaibi, Adeesh Kolluru, Abhishek Das, Aditya Grover, Anuroop Sriram, Zachary Ulissi, and C. Lawrence Zitmick. Rotation invariant graph neural networks using spin convolutions. arxiv preprint 000r:::2006005755 2021.
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. doi: 10.1103/ phystevresearch..11110.00 URL httph:tt////.m http://oooooo/o/////////// org/10. 11 3.1012002.
Anuroop Sriram, Abhishek Das, Brandon M Wood, and C. Lawrence Zitmick. Towards training billion parameter graph neural networks for atomic simulations. In International Conference Learning Representations, 2022. URL https: / /ttp:////ppeeeeeeeee///oeee wttp:eeee/oouummdd0e0eeee.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: A simple way to prevent neural networks from overfitting. Journal of Machine Learning Research, 15(56):1929-1999 1958, 2014. URL http:// httr.//mrr../// wwwpl..org//pppeere.e httg:/apprrs//vl//
srivastavee 4a html.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers neural network based molecular potentials. In International Conference on Learning Representations (ICLR), 2022. URL https: / //apenreview. net /fooomfidddiiiiii2222 2iinniddd22iiig22e223
Nathaniel Thomas, Tess E. Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-equiiariaat neural networks for 3d point clouds. arxiv preprint arXiv: 0829,, 2018.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention. arXiv preprint 20t.:20012.12877. 2020.
Raphael L. Townshend, Brent Townshend, Stephan Eismann, and Ron O. Dror. Geometric prediction: Moving beyond scalars. arXiv preprint arXiv:2006. /4163, 2020.
Raphael John Lamarre Townshend, Martin Vögele, Patricia Adriana Suriana, Alexander Derry, Alexander Powers, Yianni Laloudakis, Sidhika Balachandar, Bowen Jing, Brandon M. Anderson, Stephan Eismann, Risi Kondor, Russ Altman, and Ron O. Dror. ATOM3d: Tasks on molecules in three dimensions. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021. URL https: http////peeeeeeeeeeeeeeeeeeeeeee net forum?? iid= FM
Oliver T. Unke and Markus Meuwly. PhysNet: A neural network for predicting energies, forces, dipole moments, and partial charges. Journal of Chemical Theory and Computation, 15(6): 3678-3693, may 2019.
Oliver Thorsten Unke, Mihail Bogojeski, Michael Gastegger, Mario Geiger, Tess Smidt, and Klaus Robert Muller. SE(3)-equivariant prediction of molecular wavefunctions and electronic densities, In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Conference on Neural Information Processing (NeurIPS), 2021. URL http:: htttps::/openrvieeeeeee net forrmm? id=auGY2UQEhSu.
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.
Petar Veličković, Guillem Cucurull, Arantxa Casanova. Adriana Romero, Pietro Liò, and Yoshua Bengio. Graph attention networks. In International Conference on Learning Representations (ICLR), 2018. 8. URL https: /RL.ttts:://ppnreeer http//pppeeeieeeeeeee ew. net http://frummmmmmmdeeeeee kCZ.
15

---

Published as a conference paper at ICLR 2023
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. pp. 10402-10413, 2018.
Daniel E. Worrall, Stephan Garbin, Daniyar Turmukhambetov, and Gabriel J. Brostow. Harmonic networks: Deep translation and rotation equivariance. arxiv preprint crxiv:16/2.04642 2016.
Tian Xie and Jeffrey C. Grossman. Crystal graph convolutional neural networks for an accurate and interpretable prediction of material properties. Physical Review Letters, 120 apr 2018. doi: 10. 10.1103/physrevlet. 120. 145301. URL https: /ttp:///ooiii//oo ortp:///.......//.1.. 110.1.113110000 120.145301
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? In International Conference Learning Representations (ICLR), 2019. URL https: / //openreview. net htt///oorummdrfoo/e 61555m 6m.5Km.
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. URL https: //pp///eee/////////// net forum? id=OeWooOxEwDaa
A. Zee. Group Theory in a Nutshell for Physicists. Princeton University Press, USA, 2016.
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 on Uncertainty in Artificial Intelligence, pp. 339-349, 2018a.
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E. Deep potential molecular dynamics: scalable model with the accuracy of quantum mechanics. Phys. Rev. Lett., 120: 143001, Apr 201 doi: 10.1103/PhssRev ett. 143001. URL https: //1 iink. orr/
doi/10, 1103 htt.111.111///y/eeeeeeeeee wwwpseaeesseeeeeeeeeeeeeeaaaaeee 120 143001.
16

---

Published as a conference paper at ICLR 2023
APPENDIX
A Additional background
A.1 Group theory
A.2 Equivariance
A.3 Equivariant features based on vector spaces of irreducible representations
A.4 Tensor product
B Related works
B.1 Graph neural networks for 3D atomistic graphs
B.2 Detailed comparison between equivariant Transformers
B.3 Invariant GNNs
B.4 Attention and Transformer
C Details of architecture
C.I Equivariant operation used in Equiformer
C.2 Equiformer architecture
C.3 Dot product attention
C.4 Incorporating Equivarrrvee
C.5 Discussion on computational complexity
D Details of experiments on QM9
D.1 Training details
D.2 Comparison between SE(3) and E(3) equivariance
D.3 Comparison of training time and numbers of parameters
E Details of experiments on MD17
E.1 Training details
E.2 Additional comparison of performance to TorchMD-NET
E.3 Comparison of training time and numbers of parameters
F Details of experiments OC20
F.1 Detailed description of OC20 dataset
F.2 Training details
Results on IS2RE validation set
F.4 Comparison between SE(3) and E(3) equivariance
F.5 Comparison of training time and numbers of parameters
F.6 Error distributions
G Limitations
A ADDITIONAL BACKGROUND
In this section, we provide additional mathematical background helpful for the discussion of the proposed method. Other works (Thomas et al., 2018; Weiler et al., 2018; Kondor et al., 2018; Anderson et al., 2019; Fuchs et al., 2020: Brandstetter et al., 2022) also provide similar background. We encourage interested readers to see these works (Zee, 2016; Dresselhaus et al., 2007) for more in-depth and pedagogical presentations.
A.1 GROUP THEORY
Definition of Groups. A group is an algebraic structure that consists of set G and binary operator o : x G + and is typically denoted as G. Groups satisfy the following four axioms:
1. Closure: g ook G for all g, G.
2. Identity: There exists identity element G such that g o eog = for all g G.
17

---

Published as a conference paper at ICLR 2023
3. Inverse: For each 9 G, there exists inverse element g-! E G such that g o g- g oge
4. Associativity: fog h (f (f o h) for all f.g,heG. G.
In this work, we focus on 3D rotation, translation and inversion. Relevant groups include:
1. The Euclidean group in three dimensions 333:: 3D rotation, translation and inversion.
2. The special Euclidean group in three dimensions SE(3): 3D rotation and translation.
3. The orthogonal group in three dimensions (3): 3D rotation and inversion.
4. The special orthogonal group three dimensions SO(3): 3D rotation.
Group Representations. The actions of groups define transformations. Formally, a transformation acting vector space X parametrized by group element g G' is injective function T, X - X. A powerful result of group representation theory is that these transformations can be expressed as matrices which act vector spaces via matrix multiplication. These matrices called the group representations, Formally, a group representation D G + GI.(N) is a mapping between a group G and set of N x invertible matrices. The group representation D(g) X + X maps an N-dimensional vector space X onto itself and satisfies (h) = Dgg o h) for all g. h.he E
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 = (g) for all G, then we say the two group representations are equivalent.
(g)g) is block diagonal, which means that g acts on independent subspaces of the vector space, the representation D(y) is reducible. A particular class of representations that are convenient for composable functions are irreducible representations or "inreps", which cannot be further reduced. We can express any group representation of 80(3) direct sum (concatentation) of irreps (Zee, 2016; Dresselhaus et al., 2007: Geiger et al., 2022):
P p-P-1 P-P-1 (
+ (DL
Dio (g)
)
D(g) p-1
(g)())
(6)
where (i,(g) Wigner-D matrices with degree metnioned in Sec. 3.2.
A.2 EQUIVARIANCE
Definition of Equivariance and Invariance. Equivariance is a property of a function f: X - mapping between vector spaces X and Y. Given group G and group representations Dx (g) and (g) in input and output spaces and Y, f is equivariant to G if (g) = =(Dxxg)) for
all E X and g E G. Invariance corresponds to the case where Dy (g) is the identity / for all g G.
Equivariance in Neural Networks. Group equivariant neural networks are guaranteed to make equivariant predictions on data transformed by group. Additionally, they are found to be data- efficient and generalize better than non- con--yymetry--aee aware and invariant methods (Batzner et al., 2022; Rackers et al., 2023; Frey et al., 2022). For 3D atomistic graphs, we consider equivariance to the Euclidean group E3), which consists of 3D rotation, translation and inversion. For translation, operate relative positions and therefore our networks invariant to 3D translation. We achieve equivariance to rotation and inversion by representing our input data, intermediate features and outputs in vector spaces of 0(3) irreps and acting them with only equivariant operations.
A.3 EQUIVARIANT FEATURES BASED ON VECTOR SPACES OF IRREDUCIBLE REPRESENTATIONS
Irreducible Representations of Inversion. The group of inversion Ee 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 do. We create irreps of O(3) by simply multiplying those of (3) and Z, and introduce parity p to type-L vectors to denote how they transform under inversion. Thus, type- I. vectors in SO(3) become (ype-(T,,p) vectors in O(3). where p is or o.
18

---

Published as a conference paper at ICLR 2023
Irreps Features. As discussed in Sec. 3.2 in the main text, equivariant irreps features' and type-(L,p) p) vectors for equiivviiar irreps features. Parity denotes whether vectors change sign under inversion and can be either (even) (odd). Vectors with = o change sign under inversion while those with p = e do not. Scalar features correspond type-O vectors in the case of -euiivequivarriace and correspond type-(0, e) in the case of E(3) -equivarrree whereas type-(0, o) vectors correspond to pseudo-scalars. Euclidean vectors in R correspond to type-1 vectors and type-(1, o) vectors whereas type-(1, e) vectors correspond to pseudo-vectors. Note that type-(L, e) vectors and (ype-(L, vectors are considered vectors of different types in equivariant linear layers and layer normalizations.
type-L vectors for SE(3)-
Spherical Harmonics. Euclidean vectors F in R can be projected into type-L vectors ((4) by using spherical harmonics Y(L); f(b) = y f
[[T (Smidt et al., 2021). This is equivalent to the Fourier transform of the angular degree of freedom [[1] which can be optionally weighted by [/7]] f(L) transforms in the same manner as type-I. vectors. For
In the case of 5 FF(3)--quivvrian) Euuivequuer f(L) behaves as type-(I,, L, p) vectors, where = if L is even and p = o is odd. Visualization of spherical harmonics can be found in this website.
Vectors of Higher and Other Parities. Although previously we have restricted concrete ex- amples of vector spaces of ((3) irreps to commonly encountered scalars (type-(0, e) vectors) and Euclidean vectors (type- (1, o) vectors), vector of higher L and other parities are equally physical For example, the moment of inertia (how object rotates under torque) transforms components behaving as type-(2, e) vectors. Elasticity (how
3 x symmetric
matrix, which has symmetric-traceld object deforms under loading) transforms rank-4 or x 3 x 3 symmetric tensor, which includes components acting as type-(4, e) vectors.
A.4 TENSOR PRODUCT
Tensor Product for 0(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-Gordal coefficients to combine (ppe((L,,)) vector f(LI.P)) and type-(T.2. P2) vector g gloz)2)) and produces (ype-(T.,,)) vector h (1a.))))) as follows:
L
L2
(La.Ps) (f(L).21) p @ ((L2)P)) mg
Σ
Σ
clla.ma)
(.g.p)))lllpp (12m)) gmn
p(L..p.) (L2.p2)
(7)
my
L1 mg
Ps 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 that we additionally keep track of the output parity pa as in Eq. and use the following multiplication rules: x e, o e, and x of a type-(1, o) vector and a type-(1, (1, e) vector can result in one type-(0, o) vector, one type-(1, vector, and one t2pe-(2, vector.
For example, the tensor product
Clebsch-Gordan Coefficients. The Clebsch-Gordan coefficients for SO(3) are computed from 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.
ma) (1a.m1)llamma) = [i L2m2) (Lymsl = dsLy (h) (()))))(()) (())))
(9)
For many combinations of L1, Le, and L3, the Clebsch-Gordas coefficients are zero. The gives rise to the following selection rule for non-trivial coefficients: -|L, + L2| S L3 < 1 + L2).
In SEGNN (Brandstetter al., 2022), they are also referred to steerable features. We use the term "irreps features" remain consistent with e3nn (Geiger et al., 2022) lihrary.
19

---

Published as a conference paper at ICLR 2023
Examples of Tensor Products. Tensor products generally define the interaction between different (ype-(,p)) vectors in a symmetry-preserving manner and consist common operations as follows:
Scalar-scalar multiplication: sealar (I. 0, e) 69 scalar (I. = 0.P = e) + scalar
(T. 0,y e).
2. Scalar-vector multiplication: scalar (T. = 0, e) vector (T. = =1,7 0) + vector
(T. 1, o).
3. Vector dot product: vector (T. 1, = o) & vector (I. 1,p 0) - scalar (I. = 0. p=
e).
4. Vector cross product: vector (L 1,p o) @ vector (L 1,p o) pseudo-vector
(L = 1,p= = e).
B RELATED WORKS
B.I GRAPI NEURAL NETWORKS FOR 3D ATOMISTIC GRAPHS
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 represent atomistic graphs (Townshend et al., 2021), which are chemical bond graphs, sometimes denoted as 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 (Hamilton al., 2017; Gilmer et al., 2017; Kipf & Welling, 2017: Xu et al., 2019; Veličković et al., 20:8; Brody et al., 2022) can be directly applied to predict their properties (Ruddigkeit et al., 2012; Ramakrishnan et al., 2014; Ramsundar al., 2019; Hu et al., 2020; 2021). On the other hand, 3D spatial graphs consider positions of atoms in 3D spaces and therefore 3D geometry. Although 3D graphs can faithfully represent atomistic systems, one challenge of moving from chemical bond graphs to 3D spatial graphs is to remain invariant or 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-& vectors.
B.2 DETAILED COMPARISON BETWEEN EQUIVARIANT TRANSFORMERS
First, we compare the impact of previous equivariant Transformers (Fuchs al., 2020; Thölke Fabritiis, 2022; Le et al., 2022) in the following four aspects:
1. Previous equivariant Transformers do not perform well across datasets. For instance, SE(3)- Transformer (Fuchs al., 2020) is not as performant as other equivariant networks on QM9 as shown in Table 1. TorchMD-NEL (Thölke & Fabritiis, 2022) does not achieve comparable results to NequlP (Batzner al., 2022) on MD17 as shown in Table although it is competitive on QM9 in Table 1.
2. Equiformer simultaneously achieves the best results for MD17, QM9 and OC20 datasets, in- dicating that the Transformer architecture generally effective the literature equivariant neural networks and 3D atomistic graphs.
3. Extensive ablation studies have been conducted to justify better attention mechanism in this literature.
4. To the best of our knowledge, we are the first to apply equivariant Transformers to large and complicated datasets like OC20 and demonstrate that equivariant Transformers achieve competitive results to large models like GNS (Godwin et al., 2022) and Graphormer (Shi et al., 2022) while saving 2.3x to 15.5x training time summarized in Table 5.
Second, we compare the technical differences of architectures of equivariant Transformers. The proposed Equiformer consists of equivariant graph attention and an equivariant Transformer archi- tecture. The latter is obtained by simply replacing orginal operations in Transformers with their equivariant counterparts and including tensor product operations and corresponds to "Equiformer with dot product attention and linear message passing" indicated by Index in Table 6 and 7.
20

---

Published as a conference paper at ICLR 2023
Although with minimal modifications to original Transformers, note that this architecture has not been explored in previous equivariant Transformers and achieves competitive results on QM9 and OC20 datasets. Below discuss the advantages of the architecture, Equiformer with dot product attention and linear message passing, over other equivariant Transformers:
1. Simpler. We simply replace original operations with their equivariant counterparts and include tensor products without making further modifications. In contrast, SE(3)- Transformer (Fuchs et al., 2020) merges normalization and activation to form norm nonlin- earities, which does not exist in original Transformers. We empirically find that the norm nonlinearities (Fuchs et al., 2020) leads to higher errors compared to the equivariant layer norm used by our work.
2. More general. SEa())-Tranffrmerr (Fuchs et al., 2020) and our proposed architecture can support vectors of any degree L while other equivariant Transformers (Thölke & Fabritiis, 2022; Le et al., 2022) are limited to L and It has been shown that higher L (e.g., L up to 2 and 3) can improve the performance of networks on QM9 (Brandstetter et al., 2022) and MD17 (Batzner al., 2022). Thus, the incapability to use higher than 1 can limit their performance.
3. More efficient tensor products. Compared to (Sa))-Tranformmrr (Fuchs et al., 2020), we use more efficient depth-wise tensor products instead of fully connected tensor products. Since the proposed architecture and Srans-Tranrfrr (Fuchs et al., 2020) relative distances to parametrize the weights of tensor products, depth-wise tensor products enalbe using more channels without incurring out-of-memory errors. Specifically, for QM9, SE(3)- Transformer (Fuchs et al., 2020) only uses 16 channels for vectors of each degree while the proposed architecture can use 128, 64, and 32 channels for vectors of degree 0, 1, and 2. Using a very small number of channels can potentially lead to insufficient model capacity.
We note that the novelty of the proposed architecture, Equiformer with dot product attention and linear message passing, lies in how we choose the right operations as well as internal representations (i.e.. vectors of any degree L) and combine them in an effective manner that achieves the three advantages mentioned above. We further improve this simple architecture with our proposed equivariant graph attention, which consists of MLP attention and non- non-linear message passing.
B.3 INVARIANT GNNs
Previous works (Schütt et al., 2017; Xie & Grossman, 2018; Unke & Meuwly, 2019: Gasteiger et al., 2020b;a; Qiao et al., 2020; Liu et al., 2022; Shuaibi et al., 2021; 1; Klicpera et al., 2021) 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 (Schütt et al., 2017) uses relative distances and proposes continuous-fill ffltte convolutional layers to learn local interaction between atom pairs. DimeNet series (Gasteiger et al., 2020b;a) incorporate bond angles by using triplet representations of atoms. SphereNet (Liu et al., 2022) and GemNet (Klicpera et al., 2021; Gasteiger et al., 2022) further extend to consider dihedral angles for better performance. In order to consider directional information contained in angles, they rely on triplet or quadruplet representations of atoms. In addition to being memory- memory-intensive intensive (Sriram al., 2022), they also change graph structures by introducing higher-order interaction terms (Chen et al., 2019), which would require non-th non-trivial modifications to generic GNNs in order to apply them 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 design of generic GNNs.
B.4 ATTENTION AND TRANSFORMER
Graph Attention. Graph attention networks (GAT) (Veličković et al., 2018; Brody et al., 2022) use multi-layer perceptrons (MLP) to calculate attention weights in similar manner to message passing networks. Subsequent works using graph attention mechanisms follow either GAT-like MLP attention (Busbridge et al., 2019; Kim & Oh, 2021) or Transformerliike dot product attention (Zhang et al., 2018a; Gao & Ji, 2019; Shi et al., 2020; Dwivedi & Bresson, 2020; Kim & Oh, 2021; Kreuzer et al., 2021). In particular, Kim et al. (Kim & Oh, 2021) compares these two types of attention
21

---

Published as a conference paper at ICLR 2023
loger tmepe foruns
Iepur impe fentura
Irpar imps Iratono
Irpat imaos fomeos
Inpat wrrppe fooness
Inger imape forums 2
FOOOEGE
famur
(0) Paths contriburing to type-l vectors.
(b) Paths contriburting to type-/ vectors.
(c) Paths contributing to type-2 vecfors.
Figure 3: An alternative visualization of the depth-wise tensor product. We follow the visual- ization tensor products in e3nn (Geiger et al., 2022) 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 on one vector in each input irreps feature.
mechanisms empirically under a self-supervised setting, Brody et al. (Brody et al., 2022) analyzes their theoretical differences and compares their performance in general settings.
Graph Transformer. different line research focuses on adapting standard Transformer net- works to graph problems (Dwivedi & Bresson, 2020; Rong et al., 2020; Kreuzer et al., 2021; Ying et al., 2021; Shi et al., 2022). They adopt dot product attention in Transformers (Vaswani et al., 2017) and propose different approaches to incorporate graph- inductive biases into their networks. GROVE (Rong et al., 2020) includes additional message passing layers or graph convolutional layers incorporate local graph structures when calculating attention weights. SAN (Kreuzer al., 2021) proposes to learn position embeddings of nodes with full Laplacian spectrum. Graphormer (Ying al., 2021) proposes encode degree information in centrality embeddings and encode distances and edge features in attention biases. The proposed Equiformer belongs to one of these attempts to generalize standard Transformers to graphs and 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.
C DETAILS OF ARCHITECTURE
C.I EQUIVARIANT OPERATION USED IN EQUIFORMER
We illustrate the equivariant operations used in Equiformer in Fig. and provide an alternative visualization of desth-wise tensor products in Fig. 3.
Besides, analyze how each equivariant operation remains equivariant and satisfies that (g)x) = Dy (z), where f is function mapping between vector spaces X and Y, and (x(g) and Dy (g) are transformation matrices parametrized by in X and Y.
1. Linear. Since for each degree L, output type-L vector is linear combination of other input type-L vectors, which are transformed by the same matrix (x(g), the output type-L vector is transformed by the same matrix, meaning that Dx Dy
2. Layer normalization. For scalar parts (L = 0). they are always the same regardless of E33) transformations, and thus we can apply any function to them. For non-scalar parts (T > 0), the 1.2-norm of any type-T. vector is invariant to F(3) transformations. Therefore. the scaling of dividing by the root
square value (RMS) of I.2-norm and multiplying by a learnable parameter 7 remains the same under E(3) transformations. Multiplying an equivariant feature with an invariant number results in equivariant feature, and therefore the operation is equivariant.
3. Gate. Similar to layer normalization, We apply SiLU to the first Co channels the scalar part and sigmoid to other channels obtain non-linear weights. For the non-scalar part (T. 0), multiply each type-
apply any function to the scalar part (L 0).
22

---

Published as a conference paper at ICLR 2023
Linear
Linear
Linear
DIP
fu
Restaspe
Scall
Softisax
V2
Reshappe
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 circle denotes summation over all neighbors. Gray cells indicate intermediate irreps features.
highlight the difference of dot product attention from multi-layer perceptron attention in red. Note that key kiy and value Vay irreps features and therefore fiy in dot product attention typically has more channels than that in multi-layer perceptron attention.
vector with its corresponding non-line non-linear weight. Since the non-linear weights are invariant, multiplying equivariant features with those non-linear weights results in equivariant features.
4. Depth- wise tensor product. This operation based equivariant tensor product opera- tions and restricts that one channel in output irreps feature depends on one channel in imput irreps features. The one-to-one dependence of channels does change the interaction of different type-L vectors, and therefore, the operation is equivariant.
C.2 EQUIFORMER ARCHITECTURE
For simplicity and because most works we compare with do not include equivariance to inversion, we adopt (E irreps features in Equiformer for experiments in the main text and note that E(3) eqqivvraaa irreps features can be easily incorporated into Equiformer.
We define architectural hyper-parameters parameters like the number channels in some layers in Equiformer, which are used to specify the detailed architectures in Sec. D.I, Sec. E.1 and Sec. F.2.
use dembed to denote embedding dimension, which defines the dimension of most irreps features. Specifically, all irreps features x;, ys in Fig. have dimension dembed unless otherwise stated. Besides, we use deh to represent the dimension of spherical harmonics embeddings of relative positions in all depth-wise tensor products.
For equivariant graph attention in Fig. 1(b), the first two linear layers have the same output dimension dembed- The output dimension of depth-wise tensor products (DTP) determined by that of input irreps features. Equivariant graph attention consists of h parallel attention functions, and the value vector in each attention function has dimension dhead- We refer and dhead the number of heads and head dimension, respectively. By default, we set the number of channels in scalar feature f
(U) be the same the number of channels of type-0 type-(0, e) vectors in tij- When non-linear
23

---

Published as a conference paper at ICLR 2023
messages are adopted in Dij. set the dimension of output irreps features in gate activation to be h x dhead- Therefore, we can use two hyper-parameters and dhead to specify the detailed architecture of equivariant graph attention.
As for feed forward networks (FFNs), we denote the dimension of output irreps features in gate activation dnfn- The FFN in the last Transformer block has output dimension feature. and set fn n of the last FFN, which is followed by output head, to be d feature as well. Thus, two hyper- parameters fn. and d feature are used specify architectures of FFNS and the output dimension after Transformer blocks.
Irreps features contain channels of vectors with degrees up to Imaz- We denote Ct. type- I. vectors as (CI, L) and C(L.p) type-(L.p) p) vectors as (C(t..)) L, p) and use brackets to represent concatenations of vectors. For example, the dimension of irreps features containing 256 type-0 vectors and 128 type-1 vectors can be represented as [(256, 0), (128, 1)].
C.3 DOT PRODUCT ATTENTION
We illustrate the dot product attention without non-linear message passing used in ablation study in Fig. 4. The architecture is adapted from from multi-layer perceptron attention lies in how we obtain attention weights 0zy from fij. We split fiy into two irreps features, key kij and value Vij. and obtain query % with linear layer. Then, perform scaled dot product (Vaswani et al., 2017) between qi and kij for attention weights.
(Fuchs et al., 2020). The difference
C.4 INCORPORATING E(3) EQUIVARAIEE
To incorporate (qq)-aqquvar Equiformer, we can directly use the same architecture described in Fig. with the following two modifications:
1. We change internal representations from type-T. vectors type-(I.,p) vectors as mentioned in Sec. A.3. Note that type-(L,, vectors and type-(L, o) vectors are considered different types by equivariant operations and that scalars correspond only type-(0, e) vectors and do not include type-(0, o) vectors.
2. The behaviors of equivariant operations are changed accordingly since parities are included in internal representations. The operations of linear and layer normalization will treat type- (I, e) vectors and type-(L, o) vectors as different types. This means that we linearly combine or normalize these two types of vectors in a separate manner. Sec. A.4, we discuss how tensor products behave when E(3) equivariance is considered. For the operation of gate, we apply activation functions to type-(0, e) vectors (scalars) and treat (0, o) vectors (pseudo-scalars) in the same manner as vectors of higher L. Specifically, given input containing non-scalar C(t.p) type-(L,p) p) vectors with O < L Lmaa and p E {e,o), C(o.o) type-(0, o) vectors and (C(o.e) + C(o.o) + Ehmaz Epe(c..) C(L.P)) type-(0,e) e) vectors, we apply SiLU to the first C(o.e) (0,- vectors and sigmoid function to the other (C(o.o) + Eimas Epe(e.)) C(L.p)) type-(0, e) vectors to obtain non-linear weights and multiply each pseudo-scalar or type-(L, p) vector with corresponding non- non-linear weights. After the gate activation, the number of channels for type-(0. e) vectors is reduced to C(o.e).
C.5 DISCUSSION ON COMPUTATIONAL COMPLEXITY
We discuss the computational complexity of the proposed equivariant graph attention here.
First, we compare dot product attention with MLP attention when linear messages are used for value Way- Dot product attention requires taking the dot product of two irreps features, query q and key
kij. for attention weights, and both gi and kij have the same dimension value Fig. In contrast,
MLP attention uses only scalar features f for attention weights. The dimension of scalar features f is the same as that of the scalar part of Utj. Therefore, MLP attention generates less and smaller intermediate features for attention weights and is faster than dot product attention.
24

---

Published as a conference paper at ICLR 2023
Hlyper-prraeeters
Value description
Optimizer
Adam
Learning scheduling
Cosine learning with linear warmup
Warmup epochs
Maximum learning rate
1.5 10 10
Batch size
64, 128
Number epochs
300, 600
Weight decay
0, 5 x 10-3
Dropout rate
0.0. 0.1, 0.2
Cutoff' radius (A)
Number of radial bases
128 for Gaussian radial basis, 8 for radial bessel basis
Hidden sizes of radial functions
64
Number hidden layers radial functions
Equiformer
Number Transformer blocks
Embedding dimension dembed
(128, 0), (64, 1), (32,2) (1,0), (1,
Spherical harmonics embedding dimension d4h
Number of attention heads h Atlention head dimension dhead Hidden dimension in food forward networks defm
(32, 0). (16, 1), (8, 2)
(384, 0), (192, 1), (96. 2)
Output feature dimension dycature
(512.0)
E(3) Eqquforrre
Number of Transformer blocks Embedding dimension dembed Spherical harmonics embedding dimension da Number of attention heads h
(3) (128,0,e), (2.)
(1,
Attention head dimension dneed Ilidden dimension in feed forward networks dyz=
(32, ()). (384, 0, (96, (10, (18, (48
Output feature dimension deaturs
(512.
Table 8: type-(L,p) p) vectors as (C(r.p)) p) and use brackets to represent concatenations of vectors.
for QM9 dataset. We denote Cr type-L vectors as (Cz, L) and C(L.p)
Second, compared to linear messages, using non-linear messages increases the number of tensor product operations from to 2. Since tensor products are intensive, this inevitably increases training and inference time.
Please refer to Sec. D.1 and Sec. F.2 for the exact numbers of training time on QM9 and OC20.
D DETAILS OF EXPERIMENTS ON QM9
D.1 TRAINING DETAILS
We use the same data partition TorchMD-NET (Thölke & Fabritiis, 2022). For the task of U, Un. G. and H. where single-atom reference values are available, we subtract those reference values from ground truth. For other tasks, we normalize ground truth by subtracting mean and dividing by standard deviation.
train Equiformer with blocks with Lmaz = 2. We choose Gaussian radial basis (Schütt et al., 2017; Shuaibi et al., 2021; Klicpera et al., 2021; Shi et al., 2022) for the first six tasks in Table 1 and radial Bessel basis (Gasteiger et al., 2020b;a) for the others. We apply dropout (Srivastava et al., 2014) to attention weights agy. The dropout rate is 0.0 for the tasks of G, H. U and Uo. is 0.1 for the task of R² and is 0.2 for others. Since the tasks of G, H, U, and Uo require longer training, use slightly different hyper-parameters. The number of epochs is 600 for the tasks of G, H, U, and Un and is 300 for others. The learning rate is 1.5 x 10 for the tasks of G. H. U, and Uo and is x 10 for others. The batch size 64 for the tasks of G, H, U. and Un and is 128 for others. The weight decay is 0 for the tasks of G. H, U, and Uo and is 5 10 for others. Table 8 summarizes the hyper-pare paremeeersssss for the QM9 dataset. The detailed description of architectural parammeessssees can be found in Sec. C.2.
25

---

Published as a conference paper at ICLR 2023
We one A6000 GPU with 48GB to train each model and summarize the computational cost of training for one epoch as follows. Training [33--fqufrrrmrr in Table 9 for one epoch takes about 16.3 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 6) for epoch is 12.1 minutes, 7.2 minutes and 7.8 minutes, respectively.
D.2 COMPARISON BETWEEN SE(3) AND E(3) EQUIVARIANCE
We train two versions of Equiformers, one with (E(3) equivariant features denoted as "Equiformer" and the other with -(3uivauivarrant features denoted as " EEquiforrr"r and we compare them in Table 9. As for Table 1, we compare "Equiformer" with other works since most of them do not include equivariance to inversion.
Task
Ar
FLUMD
C2 cal/mol
Training time (minutes/oocch)
Number of parametens
Methods
Units
Equifermer
016
30
15
14
.011
023
12.1
3,53M
Equiimar
045
14
.012
023
163
3.28M
Table 9: Ablation study of SE(3)E(3) equivariance on QM9 testing set. "Equiformer" operates on (E) -eequivarrii features while "E(3) "88uffquifrmrer" uses E(3) eqqiivaraa features. Including inversion achieves similar performance.
D.3 COMPARISON OF TRAINING TIME AND NUMBERS OF PARAMETERS
We compare training time and numbers of parameters between SEGNN (Brandstetter et al., 2022), TorchMD-NET (Thölke & Fabritiis, 2022) and Equiformer and summarize the results in Table 10. Training Equiformer for 300 epochs and for 600 epochs takes 61 and 122 GPU-hours, respectively.
Compared to SEGNN, which is written with the same e 3nn library (Geiger et al., 2022), Equiformer with MLP attention and non-linear message is faster. Although Equiformer has more channels and more parameters, the training time is comparable. The reasons are as follows. Equiformer uses more efficient depth-wise tensor products (DTP), where output channel depends only input channel. SEGNN uses more compute-intensive fully connected tensor products (FCTP), where one output channel depends all input channels. Besides, SEGNN uses 4 FCTPs in each message passing block while Equiformer uses only 2 DTPs in each block.
Compared TorchMD-NET, which trained for 3000 epochs, Equiformer achieves competitve results after trained for 300 or 600 epochs. Equiformer takes more time for each epoch since Equiformer uses more expressive non-linear messages, which compared to linear messages used in other equivariant Transformers, doubles the number of tensor products and therefore almost doubles the training time. Moreover. Equiformer incorporates tensors of higher degrees (e.g., Linax = 2). which improves performance but slows down the training.
E DETAILS OF EXPERIMENTS ON MD17
E. TRAINING DETAILS
We use the same data partition as TorchMD-NET (Thölke & Fabritiis, 2022). For energy prediction, normalize ground truth by subtracting mean and dividing by standard deviation. For force prediction, we normalize ground truth by dividing by standard deviation of ground truth energy.
We train Equiformer with 6 blocks with Lmax 2 and 3. We choose the radial basis function used by PhysNet (Unke & Meuwly, 2019). We do not apply dropout to attention weights a;j. For Equiformer with Lmaz 2, the learning rate is 1 x 10-4 for benzene and is 5 x 10-4 for others. The batch size is 8, and the number of epochs is 1500. The model has about 3.50M parameters. For Equiformer with Lmaaz = 3, the learning rate is x 10 for benzene and is 2 x 10 for others. The batch size 5, and the number of epochs is 2000. The model has about 5.50M parameters. Table 11 summarizes the hyper-pare arper-parameters for the MD17 dataset. The detailed description of architectural Soperr-arameteer can be found in Sec. C.2.
26

---

Published as a conference paper at ICLR 2023
Methods
parameters
Thinnng time (GPL--ourr)
SEGNN (Brandstette al., 2022)
1 16M
81
TorchMD-NET (Thölke Fabritiis, 2022)
6.30M
Equiformer
Table 10: Comparison of training time and numbers of parameters for QM9 dataset.
Hyper-parameters
Value or description
Optimizer
Adam
Learning rate scheduling
Cosine learning rate with lincar wannup
Warmup epochs
10
Maximum learning
1 10 10-9, 10
Batch size
8,
Number of epochs
1500, 2000
Weight decay
1 10-6
Dropout rate
0.0
Weight for energy loss Weight for force loss
1
80
Cutoff radius (A)
5
Number of radial bases
32
Hidden sizes of rudial functions
64
Number of hidden layers in radial functions
2
Equiformer (Lmmx
Number of Transformer blocks Embedding dimension dembed Spherical harmonics embedding dimension dah
(128. (64, 16, (32, (1,0), (1,
Number of atlention heads h Attention hcad dimension daead Hidden dimension in feed forward networks difn
[32, 0), (16, (8, (192, (96, 2)] (512.0)
Output feature dimension d jutsure
Equiformer (Lmex 3)
Number Transformer blocks Embedklig dimension demteed Spherical harmonics embedding dimension d.a.
(128, 0), (64, (64. (32.3) (1, (1),,
Number of attention heads h
Attention head dimension dhead Hidden dimension in feed forward networks dysm.
(32, 0), (1,,1), (16, 2), (8, (384, (192, 1). (192, 2), (96.
Output feature dimension desture
(512, 0)
Table 11: Mlyee--prameele for MD17 dataset. We denote C₁ type-L vectors as (Oz, L) and
C(L.)) type-(L.p) vectors as (C(t.p), L,p) and use brackets to represent concatenations of vectors.
We use one A5000 GPU with 24GB to train different models for each molecule. Training Equiformer with Lina= 2 takes about 15.4 hours, and training Equiformer with Limaz takes about 54.4 hours.
E.2 ADDITIONAL COMPARISON TO TORCHMD-NET
Since TorchMD-NE NET (Thölke Fabritiis, 2022) also an equivariant Transformer but uses dot prod- uet attention instead of the proposed equivariant graph attention, we provide additional comparisons in Table 12. For each molecule, we adjust the ratio of the weight for force loss the weight for energy loss so that Equiformer with T.maz = 2 can achieve lower MAE for both energy and forces.
E.3 COMPARISON OF TRAINING TIME AND NUMBERS OF PARAMETERS
We compare training time and number of parameters between NequlP (Batzner et al., 2022) and Equiformer and summarize the results in Table 13. Since NequlP does not report the number of epochs, we compare the time spent for each epoch and note that NequlP is trained for more than 1000 epochs. Equiformer with Lmaz 2 is faster than NequlP with Lmar since smaller Lmaz used. Moreover, Equiformer with Lmax = 2 achieves overall better results as we use equivariant graph attention instead of linear convolution used by NequlP. When increasing Tanaz from 2 to
27

---

Published as a conference paper at ICLR 2023
Aspiria
Benzene
Malonaldehydde
Naphrhalene
Salicylic acid
Toluene
Methods
eoergy
Corces
energy
forces
energy
forves
energy
forces
energy
forces
energy
lorces
eneegy
forces
energy
forces
TorchMD-NET
2.5
2.3
3.3
2.6
5.6
2.9
Equiformer (Tonar
2.2
5.3
Table 12: Additional comparison to TorchMD-Net (Thölke & Fabritiis, 2022) on MD17 dataset. Energy and force are in units of meV and m/A/
Methods
parameters
Training (secs/epuch))
NequlP (lesos 3) (Hatzner et al., 2022)
2.97M
48.7
Equiformer (Prax
3.50M
36,9
Equiformer (Pross
5.50M
Table 13: Comparison of training time and numbers of parameters for MD17 dataset.
Equiformer achieves better results than NequIP for most molecules. However, the training time is longer than NequlP since we use 2 tensor product operations in each equivariant graph attention instead of tensor product in each linear convolution.
F DETAILS OF EXPERIMENTS ON OC20
F.1 DETAILED DESCRIPTION OF OC20 DATASET
The dataset consists of larger atomic systems, each composed of molecule called adsorbate placed on a slab called catalyst. Each input contains more atoms and more diverse atom types than QM9 and MD17. The average number of atoms in system is more than 70, and there are over atom species. The goal is to understand interaction between adsorbates and catalysts through relaxation. An adsorbate is first placed on top catalyst to form initial structure (IS). The positions atoms are updated with forces calculated by density function theory until the system is stable and becomes relaxed structure (RS). The energy of RS, or relaxed energy (RE), correlated with catalyst activity and therefore metric for understanding their interaction. We focus 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.
F.2 TRAINING DETAILS
IS2RE without Node-Level Auxiliary Task. We use hyper-parameters similar those for QM9 dataset and summarize in Table 14. For ablation study in Sec. 5.4, we use a smaller learning rate 1.5x x 10-4 for DP attention this improves the performance. The detailed description of architectural hyper-parameters can be found in Sec. C.2.
IS2RE with IS2RS Node-Level Auxiliary Task. We increase the number Transformer blocks to 18 deeper networks can benefit more from IS2RS node-level auxiliary task (Godwin et al., 2022). We follow the same hyper-parameters in Table 14 except that we increase maximum learning rate to 5 x 10-4 and set d'feature to [5122, 0), (256, 1)]. Inspired by Graphormer (Shi al., 2022), 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 1. For Noisy Nodes (Godwin et al., 2022) data augmentation, we first interpolate between initial structure and relaxed structure and then add Gaussian noise as described by Noisy Nodes (Godwin et al., 2022). When Noisy Nodes data augmentation is used, we increase the number of epochs to 40.
We use two A6000 GPUs, each with 48GB, to train models when IS2RS is not included during training. Training Equiformer and Eqqufoorrorr in Table 16 takes about 43.6 and 58.3 hours. Training Equiformer with linear messages (indicated by Index 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 data augmentation takes about days and training with Noisy Nodes takes 6 days. We note that the proposed Equiformer in Table 5 achieves competitive results even with much less computation. Specifically, training
28

---

Published as a conference paper at ICLR 2023
paper-parameters
Value or description
Optimizer
Adam
Learning scheduling
Cosine leurning with linear warmup
Warmup epochs
2
Maximum learning rate
2 10-4
Ratch size
32
Number of epochs
20
Weight decay
1 X 10
Dropout rate
0.2
Cutoff radius (Á)
5
Number of radial basis
128
Hidden size of rudial function
64
Number hidden layers radial function
2
Equiformer
Number Transformer blocks Embedding dimension dembed Spherical harmonics embedding dimension dai
(256. 03. (128)12))))) 1) (1,
Number of attention heads h Attention head dimension daend Hidden dimension food forward networks dysn
(32, 0). (16, 1)] (768. 0), (384, 1)
Output feature dimension d feature
(512,0)
E(3) Eq3ifooooo
Number of Transformer blocks Embedding dimension dembed Spherical harmonics embedding dimension dan.
(256.0, 0,8.0,ee. (64.0.o), (,), (,), (1.1,)) 8
Number of attention heads h
Attention head dimension dhead liideen dimension Output feature dimension funture
0,), (). (8, o,]]
forward networks dry
(768. 0) (19 (192))...11)))))))) 1,0)]
(512, 0.e]
Table 14: Ayper-parameters for OC20 dataset under the setting of training without IS2RS aux- iliary task. We denote CL type-L vectors (CL, L) and C(r.p) (ppe(LL.p) vectors as (C(t..)) L,P) and use brackets to represent concatenations of vectors.
Energy (eV
(5)
Methods
ID
OOD Ads
OOD Cat
OOD Both
Average
ID
OOD Ads
OOD Cat
00D Both
Average
SchNet (Schütt et 2017)
0.6465
0.7074
0.6475
0.6626
0.6660)
2.96
2.23
3.03
2.38
DimeNett+ (Gasteiger al.,,, 2020ag
0,5636
0.7127
0.5612
0.6492
0.6217
4.25
2.48
4.40
2.56
3.42
GemNet-T (Klicpem
2021)1
4.37
SphereNet (Lau et al., 2022)
0.5632
0.6682
0.3590
B.6199
0.6022
4.36
2.70
4.59
2,70
3.64
(SJEGNN (Brandstetier et 2022)
0,5497
0.6851
0.5519
0.6102
0.5992
4.99
2.50
4.71
2.85
3.77
(Brandstetter al., 2022)
Equiformer
0.5088
0.6271
0.5051
0,5545
0,5489
4.88
2.93
4.92
2.98
3.95
Table 15: Results on OC20 IS2RE validation set. f denotes results reported by Liu et al. (2022).
"Equiformer Noisy Nodes" takes about 24 GPU-days when A6000 GPUs used. The training time of "GNS Noisy Nodes" (Godwin et al., 2022) is 56 TPU-days. "Graphormer" (Shi et al., 2022) uses ensemble of 31 models and requires 372 GPU-days train all models when A100 GPUs are used.
F.3 RESULTS ON IS2RE VALIDATION SET
For completeness, we report the results of Equiformer on the validation set Table 15.
F.4 COMPARISON BETWEEN SE(3) AND EI() EQUIVARIANCE
train two versions of Equiformers, one with (E) (qqu)arrna features denoted as "Equiformer" and the other with [(3)-qquivariant features denoted as in Table 16. Including inversion improves the MAE results on ID and OOD Cat sub-splits but degrades the performance the other sub-splits. Overall, using eqqivvaria features results in slightly inferior performance. We surmise the reasons are as follows. First, inversion might not be the key bottleneck. Second, including inversion would break type-1 vectors into parts, type-(1,
and compare them
29

---

Published as a conference paper at ICLR 2023
Enegy (eV)
EarT (16
Training (minutev/epoch)
Number paramesers
Methods
ID
OOD Ads
OOD Cat
OOD Both
Average
ID
0OD Ads
OOD Cat
OOD Both
Average
Equiformer
5088
0.6271
0.5051
0,5545
0.548)
4.88
2.93
4.92
3.93
130.8
9.12M
B(3)-Bquioomer
0.5034
0,5658
0.5528
5.10
174.9
8.77M
Table 16: Ablation study of (3(3))(3) "Equiformer" operates SE(3) eqquivarari features while (qui)-qquformmrr" uses E(3- equivariant features.
equivariance on OC20 IS2RE validation set.
Methods
Number of parameters
Training (GPU--hous)
SEGNN (Brandstetter et al., 2012)
4.21M
Equiformer
9.12M
Table 17: Comparison of training time and numbers of parameters for OC20 dataset.
and type-(1, o) vectors. They are regarded as different types in equivariant linear layers and layer normalizations, and therefore, the directional information captured in these two types of vectors can only exchange in depth-wise tensor products. Third, we mainly tune Syper-parameters for Equiformer with (E(3) (eq))-vavva features, and it is possible that using E(3) (equivrrraa features would favor different parameerass
For Table 15, 3, 4, and 5, we compare "Equiformer" with other works since most of them do not include equivariance to inversion.
F.5 COMPARISON OF TRAINING TIME AND NUMBERS OF PARAMETERS
We compare training time and numbers of parameters between SEGNN (Brandstetter et al., 2022) and Equiformer when IS2RS auxiliary task is not adopted during training and summarize the results in Table 17. Equiformer achieves better results with comparable training time. Please refer to Sec. D.3 for a detailed discussion.
The comparison of training time when IS2RS auxiliary task is adopted can be found in Table 5,
F.6 ERROR DISTRIBUTIONS
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 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 uniform among different sub-splits. For example, using MLP attention and non-lincar messages improves the errors 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 calculate the metric of energy within threshold (EwT). For instance, on the OOD Both sub-split, using non-linear messages, which corresponds 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 the threshold of 0.02 eV. This explains why using non-linear messages in Table improves MAE from 0.5657 to 0.5545 but results in almost the same EWT.
G LIMITATIONS
We discuss several limitations of the proposed Equiformer and equivariant graph attention below.
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 (Geiger et al., 2022). For example, using higher degrees L can result in larger features and using tensor products can be compute-intent inttesive Part of the reasons that tensor products are that the kernels have not been heavily optimized and customized as other operations in common libraries like PyTorch (Paszke al., 2019). But this is the issue related software, not the design of
be computationally expensive
30

---

Published as a conference paper at ICLR 2023
10
10'
Absokte Error
Threshard
Absokie 10"
0.52 Threshold
10
10
Examples
Examples by Error
Equiformer with IS2RS and Noisy Nodes
Equioree with IS2RS and Noley Nodes
Equiformer with IS2RS
Equforer with IS2RS
Equifermer
Equifomer
Equiformer with linear messages
Equifommr with Inear messages
Equiformer with linear massages and dot product attention
Equitommr with Inear messages and dot product attention
aa) ID sub-split.
(b) OOD Ads sub-split.
10
10'
10
10"
Absolute Error
0.02 eV Threshold
0.02 eV Threshold
Absolute
10
10
10
5000
10000
15300
200ce
25000
5300
10000
15020
20000
29305
Examples Sorted Absolute Emar
Examples Sorted Absolute Emar
Equiformer with IS2RS and Noisy Nodes Equiformer
Equifomer with IS2RS and Naiey Nodes Equifomar Equitomer
Equiformer
Equiformer with linear messages
Equifommer with inear messages
Equiformer with linear messages and dot product attention
Equifomer with Inear messages and dat 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.
31

---

Published as a conference paper at ICLR 2023
networks. While tensor products of irreps naively do not scale well, if all possible interactions and paths are considered, some paths tensor products can also be pruned for computational efficiency. We leave these potential efficiency gains to future work and in this work focus general equivariant attention all possible paths up to Lmmax in tensor products are allowed.
Second, the improvement of the proposed equivariant graph attention depend tasks and datasets. For QM9, MLP attention improves not significantly upon dot product attention as shown in Table 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 Table Non-linear messages improve upon linear ones for the datasets.
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-linear message passing, it increases the number of tensor products from one two and requires more computation. We that if there is 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 can slow down training.
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 the proposed attention in other domains.
32


Reviews:
[{"review_id": "R5r5sBrVTw", "reviewer": "Reviewer_8jGX", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "An equivariant Transformer/GNN model is proposed to deal with the 3D symmetries (rotation, transliteration, and inversion). The main targets are atomic graphs such as molecules. \nThe main technical advancement of this paper is the equivariant attention operation with multi-layer perception and nonlinear activations. \nExperimentally, the proposed model records better figures consistently compared to the existing equivariant-aware Transformer/GNNs. ", "summary_of_the_review": "A Transformer with 3D equivariant awareness is presented. The main contribution is the MLP-based attention that is capable of equivariant outcomes. \nExperimental results are strong. \nThe manuscript is not easy to read for readers unfamiliar with equivariant topics. More explanations through feedback is expected to more correctly evaluate the paper. ", "strength_and_weaknesses": "(+) Developed a generic MLP attention operator with a non-linear activation function. \n\n\n(+) Experimental results show the efficacy of the proposed method against the existing equivariant models. \n\n(-) difficult to read especially for readers who are unfamiliar with equivariance or group theory.\n\n(-) As a result, the technical novelty of the proposed method is unclear (to me, at least)\n\n(-) The difficulty of the tackled technical challenge is not discussed. This makes it difficult for me to understand the significance of the proposed solution in the equivariant GNN literature. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "K6qHnlIu_-", "reviewer": "Reviewer_sfAh", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a new equivariant neural network to process molecular graphs.\nThe authors successfully combine two already effective equivariant architecture designs - i.e. non-linear message passing and transformers - and achieve improved performance and computational gains on multiple datasets and tasks.\nMoreover, a large ablation study is included to validate the importance of each contribution.\n", "summary_of_the_review": "SUMMARY\n\nThe work combines existing methods in a new, successful way and is supported by strong empirical results.\nThe authors could improve the work by proving the statistical significance of some results (e.g. reporting mean and std over multiple runs of their experiments).\n\n\n\nOTHER QUESTION\n\n\nYou argue that (1) the MLP attention is more computationally efficient and more expressive than (2) Dot-Product attention.\nIn C.4, you mention that the higher cost in (2) is partially because the query and the key need to have the same dimension of the value.\nWhy can't smaller keys and queries be used?\nIf only frequency-0 keys and queries were used in (2), as is done in (1), would its computational cost be lower than (1) ?\n\n\nThe description of radial basis and radial functions at page 6 is a bit unclear. What does it mean that the function consists of *two* MLPs? Do you mean a 2-layers MLP?\nAlso, I am not sure I understood the first part of this sentence: \n\"[We transform radial basis with a learnable radial function] to generate weights for those DTP layers.\"\n\nSince this work builds mostly on top of SEGNN and SE(3)-Transformer, isn't it better to use the train/val/test split used by those works in the QM9 experiments?\nDo you think the different split could make a significative difference in the results?\n", "strength_and_weaknesses": "\n\nWhile the main theoretical contribution is limited to combining existing methods, I think the paper has high practical relevance.\nIndeed, besides the state-of-the-art results, the extensive experimental section provides insights about benefits and costs of different sub-modules.\n\nEven if many related works often don't do it, I encourage the authors to report mean and std of their results, aggregated over multiple runs (the difference in performance in some experiments is relatively small).\nI think this is especially important in the ablation study: a more statistically sound comparison would increase the impact of this work.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Jxa3M3DdV9", "reviewer": "Reviewer_1yMo", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a novel transformer framework that involve equivariance for the inputs. It shows high performance on OC20 dataset and comparable results on MD17 and QM9.", "summary_of_the_review": "Overall, the novelty, experimental evaluations, and the motivation are enough in the current form. However, the paper is hard to follow because the novel part of the proposed method and the existing modules are mixed in the paper. The writing should be improved", "strength_and_weaknesses": "Strength:\n1. The proposed method considers various equivariances in transformer frameworks.\n2. The experimental results of IS2RE for OC20 dataset outperforms the other baselines.\n\nWeakness:\n1. In Sec. 4, all the components of equiformer are existing modules. The authors should emphasis their contributions or modifications to these existing modules.\n2. In Sec. 4, the analyses of equivariance are missing. Why the proposed components still keep equivariance?\n3. It seems that the performances on MD17 and QM9 are marginal. Why the baselines evaluated on MD17, QM9, and OC20 are different?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Borl1l8meQ", "reviewer": "Reviewer_d6yC", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes an E(3)/SE(3) equivariant transformer on 3D molecular graphs. The central point is to apply MLP attention + non-linear message for better capturing the interaction between atoms.  The attention weights are computed by 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, MD17 and OC20, which supports the benefit of the proposed method.", "summary_of_the_review": "Overall, this paper has made some variable efforts on improving current equivariant GNN models, by considering the techiniques from Transformers.  However, for its current presentation, it seems the technical novelty is marginal and the experimental evaluations are still insufficient to support the contributions. ", "strength_and_weaknesses": "Strengths:\n\n1. This paper did a good job on explaining and comparing the difference between different equivariant models and transformers: TFN, SE3-Transformer, NeuIP, SEGNN.\n\n2. Besides QM9 and MD17 (which are actually well-explored datasets and could be overfitted by many recent methods), the authors evaluate their method on OC20 (the direct setting), 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\nWeaknesses:\n\n1. I totally understand that MLP attention + non-linear message is the most desirable combination as supported by this paper. Extending current methods to the form proposed in this paper is not that surprising. The authors are suggested to futher highlight the novelty of the proposed method. The current writing makes people think that it is just a simple combination of serveral components (Sec. 4). Readers are interested in knowing which part is novel and which part is previously proposed. \n\n2.  The ablation studies are insufficient. From Table 6-7, it seems the contribution of this paper only lies in the form of the attention it used. If this is the case, the contribution is weak. In other words,  the authors are suggested to perform more experiments to show what other novelty (other than the attention) is and why it works.\n\n3. For the OC experiments, did you try the relaxation setting?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%