Title: EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
Venue: NeurIPS 2023
Accepted: no
Keywords: equivariant neural networks

graph neural networks

computational physics

transformer networks
Abstract: Equivariant Transformers such as Equiformer have demonstrated the efficacy of applying Transformers to the domain of 3D atomistic systems. However, they are still limited to small degrees of equivariant representations due to their computational complexity. In this paper, we investigate whether these architectures can scale well to higher degrees. Starting from Equiformer, we first replace $SO(3)$ convolutions with eSCN convolutions to efficiently incorporate higher-degree tensors. Then, to better leverage the power of higher degrees, we propose three architectural improvements – attention re-normalization, separable $S^2$ activation and separable layer normalization. Putting this all together, we propose EquiformerV2, which outperforms previous state-of-the-art methods on the large-scale OC20 dataset by up to 15% on forces, 5% on energies, offers better speed-accuracy trade-offs, and 2$\times$ reduction in DFT calculations needed for computing adsorption energies.

Abstract

Equivariant Transformers such as Equiformer have demonstrated the efficacy of applying Transformers to the domain of 3D atomistic systems. However, they are still limited to small degrees of equivariant representations due to their computational complexity. In this paper, we investigate whether these architectures can scale well to higher degrees. Starting from Equiformer, we first replace SO p 3 q convolutions with eSCN convolutions to efficiently incorporate higher-degree tensors. Then, to better leverage the power of higher degrees, we propose three architectural improvements - attention re-normalization, separable S 2 activation and separable layer normalization. Putting this all together, we propose EquiformerV2, which outperforms previous state-of-the-art methods on the large-scale OC20 dataset by up to 15% on forces, 5% on energies, offers better speed-accuracy trade-offs, and 2 ˆ reduction in DFT calculations needed for computing adsorption energies.

1 Introduction

In recent years, machine learning (ML) models have shown promising results in accelerating and scaling high-accuracy but compute-intensive quantum mechanical calculations by effectively accounting for key features of atomic systems, such as the discrete nature of atoms, and Euclidean and permutation symmetries [1-10]. By bringing down computational costs from hours or days to fractions of seconds, these methods enable new insights in many applications such as molecular simulations, material design and drug discovery. A promising class of ML models that have enabled this progress is equivariant graph neural networks (GNNs) [5,11-18].

Equivariant GNNs treat 3D atomistic systems as graphs, and incorporate inductive biases such that their internal representations and predictions are equivariant to 3D translations, rotations and optionally inversions. Specifically, they build up equivariant features of each node as vector spaces of irreducible representations (or irreps) and have interactions or message passing between nodes based on equivariant operations such as tensor products. Recent works on equivariant Transformers, specifically Equiformer [17], have shown the efficacy of applying Transformers [19, 20], which have previously enjoyed widespread success in computer vision [21-23], language [24, 25], and graphs [26-29], to this domain of 3D atomistic systems.

Abottleneck in scaling Equiformer as well as other equivariant GNNs is the computational complexity of tensor products, especially when we increase the maximum degree of irreps L max . This limits these models to use small values of L max (e.g., L max ď 3), which consequently limits their performance. Higher degrees can better capture angular resolution and directional information, which is critical to accurate prediction of atomic energies and forces. To this end, eSCN [18] recently proposes efficient convolutions to reduce SO p 3 q tensor products to SO p 2 q linear operations, bringing down the computational cost from O p L 6 max q to O p L 3 max q and enabling scaling to larger values of L max (e.g., L max up to 8 ). However, except using efficient convolutions for higher L max , eSCN still follows SEGNN [15]-like message passing network design, and Equiformer has been shown to improve upon SEGNN. Additionally, this ability to use higher L max challenges whether the previous design of equivariant Transformers can scale well to higher-degree representations.

Tij

Linear

Linear

Concat

Type z;

Atom

Constant

One

Type z;

Atom

Constant

One

T'i

000

o

Linear

DTP

Convolution

eSCN

SO(2) Linear)

fi

Linear

Linear)

DTP

Convolution

eSCN

Linear)

Linear

SO(2)

Embedding

fj

Reshape

Linear

D

Reshape

Separable

Irreps Feature f

Scale

Layer Norm

Irreps Feature f

Separable S

Layer Norm)

SeparableS

Scale

GraphAttention

Feed Forward

Layer Norm

Equivariant

Separable

Softmax

Linear)

ajj

Linear

DTP

Gate

Vij

Re-normalization

Convolution

Activation

Attention

eSCN

Softmax

Linear

aj

SO(2) Linear)

Activation

Vij

(c) Embedding (Atom +Edge-Degree)

X;

X;

X;

X;

Network

×N

Reshape

?

Reshape

D

Linear)

Gate

Separable S

Activation

SeparableS

Activation

Linear)

Layer Norm

Separable

Linear

Linear

Linear

Linear

Output Head

yi

yi

yi

(a)EquiformerV2

(b) Equivariant Graph Attention

(d)Feed Forward Network

Architecture

In this paper, we are interested in adapting eSCN convolutions for higher-degree representations 40 to equivariant Transformers. We start with Equiformer [17] and replace SO p 3 q convolutions with 41 eSCN convolutions. We find that naively incorporating eSCN convolutions does not result in better 42 performance than the original eSCN model. Therefore, to better leverage the power of higher 43 degrees, we propose three architectural improvements - attention re-normalization, separable S 2 44 activation and separable layer normalization. Putting this all together, we propose EquiformerV2, 45 which is developed on the large and diverse OC20 dataset [30]. Experiments on OC20 show that 46 EquiformerV2 outperforms previous state-of-the-art methods with improvements of up to 15% on 47 forces and 5% on energies, and offers better speed-accuracy trade-offs compared to existing invariant 48 and equivariant GNNs. Additionally, when used in the AdsorbML algorithm [10] for performing 49 adsorption energy calculations, EquiformerV2 achieves the highest success rate and 2 ˆ reduction in 50 DFT calculations to achieve comparable adsorption energy accuracies as previous methods. 51

52

2 Related Works

SE(3)/E(3) -Equivariant GNNs. Equivariant neural networks [5,7,11-18,31-38] use equivariant 53 irreps features built from vector spaces of irreducible representations (irreps) to achieve equivariance 54 to 3D rotation [11-13]. They operate on irreps features with equivariant operations like tensor 55 products. Previous works differ in equivariant operations used in their networks and how they 56 combine those operations. TFN [11] and NequIP [5] use equivariant graph convolution with linear 57 messages built from tensor products, with the latter utilizing extra equivariant gate activation [12]. 58 SEGNN [15] introduces non-linearity to messages passing [1,39] with equivariant gate activation, and 59 the non-linear messages improve upon linear messages. SE(3)-Transformer [14] adopts equivariant 60 dot product attention [19] with linear messages. Equiformer [17] improves upon previously mentioned 61 equivariant GNNs by combining MLP attention and non-linear messages. Equiformer additionally 62 introduces equivariant layer normalization and regularizations like dropout [40] and stochastic 63 depth [41]. However, the networks mentioned above rely on compute-intensive SO p 3 q tensor 64 products to mix the information of vectors of different degrees during message passing, and therefore 65 they are limited to small values for maximum degrees L max of equivariant representations. SCN [42] 66 proposes rotating irreps features based on relative position vectors and identifies a subset of spherical 67 harmonics coefficients, on which they can apply unconstrained functions. They further propose 68 relaxing the requirement for strict equivariance and apply typical functions to rotated features during 69 message passing, which trades strict equivariance for computational efficiency and enables using 70 higher values of L max . eSCN [18] further improves upon SCN by replacing typical functions with 71 SO p 2 q linear layers for rotated features and imposing strict equivariance during message passing. 72

However, except using more efficient operations for higher L max , SCN and eSCN mainly adopt 73 the same network design as SEGNN, which is less performant than Equiformer. In this work, we 74 propose EquiformerV2 which includes all the benefits of the above networks by incorporating eSCN 75 convolutions into Equiformer and adopts three additional architectural improvements. 76

77

78

79

80

81

82

83

84

85

86

87

Invariant GNNs. Prior works [4,8,43-51] extract invariant information from 3D atomistic graphs and operate on the resulting graphs augmented with invariant features. Their differences lie in leveraging different geometric features such as distances, bond angles (3 atom features) or dihedral angles (4 atom features). SchNet [43] models interaction between atoms with only relative distances. DimeNet series [4,46] use triplet representations of atoms to incorporate bond angles. SphereNet [48] and GemNet [50, 51] further include dihedral angles by considering quadruplet representations. However, the memory complexity of triplet and quadruplet representations of atoms do not scale well with the number of atoms, and this requires additional modifications like interaction hierarchy used by GemNet-OC [51] for large datasets like OC20 [30]. Additionally, for the task of predicting DFT calculations of energies and forces on the large-scale OC20 dataset, invariant GNNs have been surpassed by equivariant GNNs recently.

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

3 Background

3.1 SE p 3 q{ E p 3 q -Equivariant Neural Networks

We discuss the relevant background of SE p 3 q{ E p 3 q -equivariant neural networks here. Please refer to Sec. A in appendix for more details of equivariance and group theory.

Including equivariance in neural networks can serve as a strong prior knowledge, which can therefore improve data efficiency and generalization. Equivariant neural networks use equivariant irreps features built from vector spaces of irreducible representations (irreps) to achieve equivariance to 3D rotation. Specifically, the vector spaces are p 2 L ` 1 q -dimensional, where degree L is a non-negative integer. L can be intuitively interpreted as the angular frequency of the vectors, i.e., how fast the vectors rotate with respect to a rotation of the coordinate system. Higher L is critical to tasks sensitive to angular information like predicting forces [5,18,42]. Vectors of degree L are referred to as typeL vectors, and they are rotated with Wigner-D matrices D p L q when rotating coordinate systems. Euclidean vectors ⃗ r in R 3 can be projected into typeL vectors by using spherical harmonics Y p L q p ⃗ r || ⃗ r || q . We use order m to index the elements of typeL vectors, where ´ L ď m ď L . We concatenate multiple typeL vectors to form an equivariant irreps feature f . Concretely, f has C L typeL vectors, where 0 ď L ď L max and C L is the number of channels for typeL vectors. In this work, we mainly consider C L ' C , and the size of f is p L max ` 1 q 2 ˆ C . We index f by channel i , degree L , and order m and denote as f p L q m,i .

Equivariant GNNs update irreps features by passing messages of transformed irreps features between nodes. To interact different typeL vectors during message passing, we use tensor products, which generalize multiplication to equivariant irreps features. Denoted as b L 3 L 1 ,L 2 , the tensor product uses Clebsch-Gordan coefficients to combine typeL 1 vector f p L 1 q and typeL 2 vector g p L 2 q and produces typeL 3 vector h p L 3 q :

where m 1 denotes order and refers to the m 1 -th element of f p L 1 q . Clebsch-Gordan coefficients C p L 3 ,m 3 q p L 1 ,m 1 qp L 2 ,m 2 q are non-zero only when | L 1 ´ L 2 | ď L 3 ď | L 1 ` L 2 | and thus restrict output vectors to be of certain degrees. We typically discard vectors with L ą L max , where L max is a hyper-parameter, to prevent vectors of increasingly higher dimensions. In many works, message passing is implemented as equivariant convolutions, which perform tensor products between input irreps features x p L 1 q and spherical harmonics of relative position vectors Y p L 2 q p ⃗ r || ⃗ r || q .

3.2 Equiformer

Equiformer [17] is an SE p 3 q / E p 3 q -equivariant GNN that combines the inductive biases of equivariance with the strength of Transformers [19,22]. First, Equiformer replaces scalar node features with equivariant irreps features to incorporate equivariance. Next, it performs equivariant operations on these irreps features and equivariant graph attention for message passing. These operations include

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

tensor products and equivariant linear operations, equivariant layer normalization [52] and gate activation [12, 34]. For stronger expressivity in the attention compared to typical Transformers, Equiformer uses non-linear functions for both attention weights and message passing. Additionally, Equiformer incorporates regularization techniques common in Transformers applied to other domains, e.g., dropout [40] to attention weights [53] and stochastic depth [54] to the outputs of equivariant graph attention and feed forward networks. Please refer to the Equiformer paper [17] for more details.

3.3 eSCN Convolution

While tensor products are necessary to interact vectors of different degrees, they are compute-intensive. To reduce the complexity, eSCN convolutions [18] are proposed to use SO p 2 q linear operations for efficient tensor products. We provide an outline and intuition for their method here, and please refer to Sec. A and their work [18] for mathematical details.

A traditional SO p 3 q convolution interacts input irreps features x p L i q m i and spherical harmonic projections of relative positions Y p L f q m f p r ⃗ ij q with an SO p 3 q tensor product with Clebsch-Gordan coefficients C p L o ,m o q p L i ,m i q , p L f ,m f q . The projection Y p L f q m f p r ⃗ ij q becomes sparse if we rotate the relative position vector ⃗ r ij with a rotation matrix D ij to align with the direction of L ' 0 and m ' 0 , which corresponds to the z axis traditionally but the y axis in the conventions of e3nn [55]. Concretely, given D ij ⃗ r ij aligned with the y axis, Y p L f q m f p D ij ⃗ r ij q ‰ 0 only for m f ' 0 . If we consider only m f ' 0 , C p L o ,m o q p L i ,m i q , p L f ,m f q can be simplified, and C p L o ,m o q p L i ,m i q , p L f , 0 q ‰ 0 only when m i ' ˘ m o . Therefore, the original expression depending on m i , m f , and m o is now reduced to only depend on m o . This means we are no longer mixing all integer values of m i and m f , and outputs of order m o are linear combinations of inputs of order ˘ m o . eSCN convolutions go one step further and replace the remaining non-trivial paths of the SO p 3 q tensor product with an SO p 2 q linear operation to allow for additional parameters of interaction between ˘ m o without breaking equivariance. To summarize, eSCN convolutions achieve efficient equivariant convolutions by first rotating irreps features based on relative position vectors and then performing SO p 2 q linear operations on the rotated features. The key idea is that the rotation sparsifies tensor products and simplifies the computation.

4 EquiformerV2

Starting from Equiformer [17], we first use eSCN convolutions to scale to higher-degree representations (Sec. 4.1). Then, we propose three architectural improvements, which yield further performance gain when using higher degrees: attention re-normalization (Sec. 4.2), separable S 2 activation (Sec. 4.3) and separable layer normalization (Sec. 4.4). Figure 1 illustrates the overall architecture of EquiformerV2 and the differences from Equiformer.

4.1 Incorporating eSCN Convolutions for Efficient Tensor Products and Higher Degrees

The computational complexity of SO p 3 q tensor products used in traditional SO p 3 q convolutions during equivariant message passing scale unfavorably with L max . Because of this, it is impractical for Equiformer to use beyond L max ' 1 for large-scale datasets like OC20 [30] and beyond L max ' 3 for small-scale datasets like MD17 [56-58]. Since higher L max can better capture angular information and are correlated with model expressivity [5], low values of L max can lead to limited performance on certain tasks such as predicting forces. Therefore, we replace original tensor products with eSCN convolutions [18] for efficient tensor products, enabling Equiformer to scale up L max to 6 or 8 on the large-scale OC20 dataset.

Equiformer uses equivariant graph attention for message passing. The attention consists of depthwise tensor products, which mix information across different degrees, and linear layers, which mix information between channels of the same degree. Since eSCN convolutions mix information across both degrees and channels, we replace the SO p 3 q convolution, which involves one depth-wise tensor product layer and one linear layer, with a single eSCN convolutional layer, which consists of a rotation matrix D ij and an SO p 2 q linear layer as shown in Figure 1b.

4.2 Attention Re-normalization

Equivariant graph attention in Equiformer uses tensor products to project node embeddings x i and x j , which contain vectors of different degrees, to scalar features f p 0 q ij and applies non-linear functions to f p 0 q ij for attention weights a ij . The node embeddings x i and x j are obtained by applying equivariant

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

Input irreps

Degree = 0

Degree =1

features

SiLU

F

F

SiLU

Sigmoid

G-1

Output irreps

features

(a)GateActivation

(b)S²Activation

(c)SeparableS²Activation

Subtract

Subtract

mean μ/0)

Divide by

Divide by

Divide by

Divide by

Divide by the

std f0)

RMS

RMS o(2)

std (0)

same RMS L >0)

(a)EquivariantLayerNormalization

(b)SeparableLayerNormalization

layer normalization [17] to previous outputs. We note that vectors of different degrees in x i and x j are normalized independently, and therefore when they are projected to the same degree, the resulting f p 0 q ij can be less well-normalized. To address the issue, we propose attention re-normalization and introduce one additional layer normalization (LN) [52] before non-linear functions. Specifically, given f p 0 q ij , we first apply LN and then use one leaky ReLU layer and one linear layer to calculate z ij ' a J LeakyReLU p LN p f p 0 q ij qq and a ij ' softmax j p z ij q ' exp p z ij q ř k P N p i q exp p z ik q , where a is a learnable vector of the same dimension as f p 0 q ij .

4.3 Separable S 2 Activation

The gate activation [12] used by Equiformer applies sigmoid activation to scalar features to obtain non-linear weights and then multiply irreps features of degree ą 0 with non-linear weights to add non-linearity to equivariant features. The activation, however, only accounts for the interaction from vectors of degree 0 to those of degree ą 0 and could be sub-optimal when we scale up L max .

2

To better mix the information across degrees, SCN [42] and eSCN [18] propose to use S activation [59]. The activation first converts vectors of all degrees to point samples on a sphere for each channel, applies unconstrained functions F to those samples, and finally convert them back to vectors. Specifically, given an input irreps feature x P R p L max ` 1 q 2 ˆ C , the output is y ' G ´ 1 p F p G p x qqq , where G denotes the conversion from vectors to point samples on a sphere, F can be typical SiLU activation [60, 61] or typical MLPs, and G ´ 1 is the inverse of G .

While S 2 activation can better mix vectors of different degrees, we find that directly replacing the gate activation with S 2 activation results in training instability (row 3 in Table 1a). To address the issue, we propose separable S 2 activation, which separates activation for vectors of degree 0 and those of degree ą 0 . Similar to gate activation, we have more channels for vectors of degree 0 . As shown in Figure 2c, we apply a SiLU activation to the first part of vectors of degree 0 , and the second part of vectors of degree 0 are used for S 2 activation along with vectors of higher degrees. After S 2 activation, we concatenate the first part of vectors of degree 0 with vectors of degrees ą 0 as the final output and ignore the second part of vectors of degree 0 . Additionally, we also use separable S 2 activation in point-wise feed forward networks (FFNs). Figure 2 illustrates the differences between gate activation, S 2 activation and separable S 2 activation.

4.4 Separable Layer Normalization

As mentioned in Sec. 4.2, equivariant layer normalization used by Equiformer normalizes vectors of different degrees independently, and when those vectors are projected to the same degree, the projected vectors can be less well-normalized. Therefore, instead of performing normalization to each degree independently, we propose separable layer normalization (SLN), which separates normalization for vectors of degree 0 and those of degrees ą 0 . Mathematically, let x P R p L max ` 1 q 2 ˆ C denote an input irreps feature of maximum degree L max and C channels, and x p L q m,i denote the L -th degree, m -th order and i -th channel of x . SLN calculates the output y as follows. For L ' 0 , y p 0 q ' γ p 0 q ˝ ´ x p 0 q ´ µ p 0 q σ p 0 q ¯ ` β p 0 q , where µ p 0 q ' 1 C ř C i ' 1 x p 0 q 0 ,i and σ p 0 q ' b 1 C ř C i ' 1 p x p 0 q 0 ,i ´ µ p 0 q q 2 . For L ą 0 , y p L q ' γ p L q ˝ ´ x p L q σ p L ą 0 q ¯ , where σ p L ą 0 q ' b 1 L max ř L max L ' 1 ` σ p L q ˘ 2 and σ p L q ' c 1 C ř C i ' 1 1 2 L ` 1 ř L m '´ L ´ x p L q m,i ¯ 2

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

260

261

262

263

γ p 0 q , γ p L q , β p 0 q P R C are learnable parameters, µ p 0 q and σ p 0 q are mean and standard deviation of vectors of degree 0 , σ p L q and σ p L ą 0 q are root mean square values (RMS), and ˝ denotes elementwise product. The computation of y p 0 q corresponds to typical layer normalization. We note that the difference between equivariant layer normalization and SLN lies only in y p L q with L ą 0 and that equivariant layer normalization divides x p L q by σ p L q , which is calculated independently for each degree L , instead of σ p L ą 0 q , which considers all degrees L ą 0 . Figure 3 compares how µ p 0 q , σ p 0 q , σ p L q and σ p L ą 0 q are calculated in equivariant layer normalization and SLN.

4.5 Overall Architecture

Here, we discuss all the other modules in EquiformerV2 and focus on the differences from Equiformer.

Equivariant Graph Attention. Figure 1b illustrates equivariant graph attention after the above modifications. As described in Sec. 4.1, given node embeddings x i and x j , we first concatenate them along the channel dimension and then rotate them with rotation matrices D ij based on their relative positions or edge directions ⃗ r ij . The rotation enables reducing SO p 3 q tensor products to SO p 2 q linear operations, and we replace depth-wise tensor products and linear layers between x i , x j and f ij with a single SO p 2 q linear layer. To consider the information of relative distances || ⃗ r ij || , in the same way as eSCN [18], we transform || ⃗ r ij || with a radial function to obtain distance embeddings and then multiply distance embeddings with concatenated node embeddings before the first SO p 2 q linear layer. We split the outputs f ij of the first SO p 2 q linear layer into two parts. The first part is scalar features f p 0 q ij , which only contains vectors of degree 0 , and the second part is irreps features f p L q ij and includes vectors of all degrees up to L max . As mentioned in Sec. 4.2, we first apply an additional LN to f p 0 q ij and then follow the design of Equiformer by applying one leaky ReLU layer, one linear layer and a final softmax layer to obtain attention weights a ij . As for value v ij , we replace the gate activation with separable S 2 activation with F being a single SiLU activation and then apply the second SO p 2 q linear layer. While in Equiformer, the message m ij sent from node j to node i is m ij ' a ij ˆ v ij , here we need to rotate a ij ˆ v ij back to original coordinate frames and the message m ij becomes D ´ 1 ij p a ij ˆ v ij q . Finally, we can perform h parallel equivariant graph attention functions given f ij . The h different outputs are concatenated and projected with a linear layer to become the final output y i . Parallelizing attention functions and concatenating can be implemented with 'Reshape'.

Feed Forward Network. As illustrated in Figure 1d, we replace the gate activation with separable S 2 activation. The function F consists of a two-layer MLP, with each linear layer followed by SiLU, and a final linear layer.

Embedding. This module consists of atom embedding and edge-degree embedding. The former is the same as that in Equiformer. For the latter, as depicted in the right branch in Figure 1c, we replace original linear layers and depth-wise tensor products with a single SO p 2 q linear layer followed by a rotation matrix D ´ 1 ij . Similar to equivariant graph attention, we consider the information of relative distances by multiplying the outputs of the SO p 2 q linear layer with distance embeddings.

Radial Basis and Radial Function. We represent relative distances || ⃗ r ij || with a finite radial basis like Gaussian radial basis functions [43] to capture their subtle changes. We transform radial basis with a learnable radial function to generate distance embeddings. The function consists of a two-layer MLP, with each linear layer followed by LN and SiLU, and a final linear layer.

Output Head. To predict scalar quantities like energy, we use one feed forward network to transform irreps features on each node into a scalar and then perform sum aggregation over all nodes. As for predicting forces acting on each node, we use a block of equivariant graph attention and treat the output of degree 1 as our predictions.

5 OC20 Experiments

Our experiments focus on the large and diverse OC20 dataset [30] (Creative Commons Attribution 4.0 License), which consists of 1 . 2 MDFT relaxations for training and evaluation, computed with the revised Perdew-Burke-Ernzerhof (RPBE) functional [62]. Each structure in OC20 has an adsorbate molecule placed on a catalyst surface, and the core task is Structure-to-Energy-Forces (S2EF), which is to predict the energy of the structure and per-atom forces. Models trained for the S2EF task are evaluated on energy and force mean absolute error (MAE). These models can in turn be used for performing structure relaxations by using the model's force predictions to iteratively update the atomic positions until a relaxed structure corresponding to a local energy minimum is found. These

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

(a) Architectural improvements. Attention re-normalization improves energies, and separable S 2 activation ('Sep. S 2 ') and separable layer normalization ('SLN') improve forces.

6

6

6

8

8

Epochs

12

20

30

12

20

eSCN

forces energy

21.3

20.6

20.1

21.3

294

290

285

296

forces energy

20.46

19.78

19.42

20.46

285

280

278

279

-

-

19.95

273

(b) Training epochs. Training for more epochs consistently leads to better results.

(c) Degrees L max . Higher degrees are consistently helpful.

(d) Orders M max . Higher orders mainly improve energy predictions.

(e) Number of blocks. Adding more Transformer blocks can help both force and energy predictions.

relaxed structure and energy predictions are evaluated on the Initial Structure to Relaxed Structure (IS2RS) and Initial Structure to Relaxed Energy (IS2RE) tasks. The 'All' split of OC20 contains 134 Mtraining structures spanning 56 elements, and 'MD' split consists of 38 Mstructures. We first conduct ablation studies on EquiformerV2 trained on the smaller S2EF-2M subset (Sec. 5.1). Then, we report the results of training on S2EF-All and S2EF-All+MD splits (Sec. 5.2). Additionally, we investigate the performance of EquiformerV2 when used in the AdsorbML algorithm [10] (Sec. 5.3). Please refer to Sec. B and C for details of models and training.

5.1 Ablation Studies

Architectural Improvements . In Table 1a, we ablate the three proposed architectural changes attention re-normalization, separable S 2 activation and separable layer normalization. First, with attention re-normalization (row 1 and 2), energy errors improve by 2 . 4% , while force errors are about the same. Next, we replace the gate activation with S 2 activation used in SCN [42] and eSCN [18], but that does not converge (row 3). Instead, using the proposed separable S 2 activation (row 4), where we have separate paths for invariant and equivariant features, converges to 5% better forces albeit hurting energies. Similarly, replacing equivariant layer normalization with separable layer normalization (row 5) further improves forces by 1 . 5% . Finally, these modifications enable training for longer without overfitting (row 7), further improving forces by 3 . 6% and recovering energies to similar accuracies as Index 2. Overall, our modifications improve forces by 10% and energies by 3% . Note that simply incorporating eSCN convolutions into Equiformer (row 1) and using higher degrees does not result in improving over the original eSCN baseline (row 8), and that the proposed architectural changes are necessary.

Scaling of Parameters . In Tables 1c, 1d, 1e, we systematically vary the maximum degree L max , the maximum order M max , and the number of Transformer blocks and compare with equivalent eSCN variants. There are several key takeaways. First, across all experiments, EquiformerV2 performs better than its eSCN counterparts. Second, while one might intuitively expect higher resolution features and larger models to perform better, this is only true for EquiformerV2, not eSCN. For example, increasing L max from 6 to 8 or M max from 3 to 4 degrades the performance of eSCN on energy predictions but helps that of EquiformerV2. In Table 1b, we show that longer training regimes are crucial. Increasing the training epochs from 12 to 30 with L max ' 6 improves force and energy predictions by 5% and 2 . 5% , respectively.

Comparison of Speed-Accuracy Trade-offs . To be practically useful for atomistic simulations and material screening, models should offer flexibility in speed-accuracy tradeoffs. We compare these trade-offs for EquiformerV2 with prior works in Figure 4a. Here, the speed is reported as the number of structures processed per GPU-second during inference and measured on V100 GPUs.

L

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

318

319

320

321

Throughput

S2EF validation

S2EF test

IS2RS test

IS2RE test

Training

Samples /

Energy MAE

Force MAE

Energy MAE

Force MAE

AFbT

ADwT

Energy MAE

set

Model

Ò

(meV)

Ó

(meV/Å)

Ó

(meV)

Ó

(meV/Å)

Ó

(%)

Ò

(%)

Ò

(meV)

Ó

CGCNN [44]

-

590

74.0

608

73.3

-

-

-

SchNet [43]

-

549

56.8

540

54.7

-

14.4

764

ForceNet-large [63]

15.3

-

33.5

-

32.0

12.7

49.6

-

DimeNet++-L-F+E [4]

4.6

515

32.8

480

31.3

21.7

51.7

559

SpinConv [49]

6.0

371

41.2

336

29.7

16.7

53.6

437

GemNet-dT [50]

25.8

315

27.2

292

24.2

27.6

58.7

400

GemNet-XL [8]

1.5

-

-

270

20.5

30.8

62.7

371

GemNet-OC [51]

18.3

244

21.7

233

20.7

35.3

60.3

355

SCN L=8 K=20 [42]

-

-

-

244

17.7

40.3

67.1

330

eSCN L=6 K=20 [18]

2.9

-

-

242

17.1

48.5

65.7

341

λ

E

'

2

)

1.8

236

15.7

229

14.8

53.0

69.0

316

All+MD

GemNet-OC-L-E [51]

7.5

239

22.1

230

21.0

-

-

-

GemNet-OC-L-F [51]

3.2

252

20.0

241

19.0

40.6

60.4

-

GemNet-OC-L-F+E [51]

-

-

-

-

-

-

-

348

SCN L=6 K=16 (4-tap 2-band) [42]

-

-

-

228

17.8

43.3

64.9

328

SCN L=8 K=20 [42]

-

-

-

237

17.2

43.6

67.5

321

eSCN L=6 K=20 [18]

2.9

243

17.1

236

16.2

50.3

66.7

327

λ

E

'

2

)

1.8

230

14.6

227

13.8

55.4

69.8

311

λ

E

'

4

)

1.8

227

15.0

219

14.2

54.4

69.4

309

(a) Trade-offs between inference speed

and validation force MAE.

(b) Trade-offs between training cost and validation force MAE.

30

30

GemNet-dT(2021)

A

28

A

28

SCN (2022)

MAE(meV/)

26

ForceMAE(meV/)

26

eSCN (2023)

GemNet-OC(2022)

24

24

Force

22

8

1.6x

22

-8%

E

20

01.5x

20

18

18

5

10

20

40

60

500

1000

2000

4000

8000

Samplesper GPU-second

Training GPU hours

For the same force MAE as eSCN, EquiformerV2 is up to 1 . 6 ˆ faster, and for the same speed as eSCN, EquiformerV2 is up to 8% more accurate. Compared to GemNet-OC [51] at the same speed, EquiformerV2 is 5% more accurate. Comparing to the closest available EquifomerV2 point, GemNet-dT [50] is 1 . 25 ˆ faster but 30% worse. Overall, EquiformerV2 clearly offers a better tradeoff between speed and accuracy. In similar spirit, we also study the training cost of EquiformerV2 compared to prior works in Figure 4b, and find that it is substantially more training efficient.

5.2 Main Results

Table 2 reports results on the test splits for all the three tasks of OC20, averaged across the indistribution, out-of-distribution adsorbates, out-of-distribution catalysts, and out-of-distribution both subsplits. Models are trained on either OC20 S2EF-All and S2EF-All+MD splits. All test results are computed via the EvalAI evaluation server 1 . EquiformerV2 outperforms all previous models across all tasks, improving by 4% on S2EF energy MAE, by 15% on S2EF force MAE, by 5% absolute on IS2RS Average Forces below Threshold (AFbT), and by 4% on IS2RE energy MAE. In particular, the improvements in force predictions are significant. Going from SCN [42] to eSCN [18], S2EF test force MAE improves from 17 . 2 meV/Å to 16 . 2 meV/Å , largely due to replacing approximate equivariance in SCN with strict equivariance in eSCN during message passing and scaling to higher degrees. Similarly, by scaling up the degrees of representations in Equiformer [17], EquiformerV2 further improves force MAE to 13 . 8 meV/Å, more than doubling the gain of going from SCN to eSCN. These better force predictions also translate to higher IS2RS test AFbT, which is computed via DFT single-point calculations to check if the DFT forces on the predicted relaxed structures are close to zero. A 5% improvement on AFbT is a strong step towards replacing DFT with ML.

5.3 AdsorbML Results

Lan et al. [10] recently proposed the AdsorbML algorithm, wherein they show that recent state-of-theart GNNs (e.g. SCN [42]) can achieve more than 1000 ˆ speedup over DFT relaxations at computing

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

Table 3: AdsorbML results with EquiformerV2 ( λ E ' 4 ) trained on S2EF-All+MD from Table 2.

adsorption energies within a 0 . 1 eV margin of DFT results with an 87% success rate. This is done by using OC20-trained models to perform structure relaxations for an average 90 configurations of an adsorbate placed on a catalyst surface, followed by DFT single-point calculations for the topk structures with lowest predicted relaxed energies, as a proxy for calculating the global energy minimum or adsorption energy. We refer the reader to the AdsorbML paper [10] for more details. We benchmark AdsorbML with EquiformerV2, and Table 3 shows that it improves over SCN by a significant margin, with 8% and 5% absolute improvements at k ' 1 and k ' 2 , respectively. Moreover, EquiformerV2 at k ' 2 is more accurate at adsorption energy calculations than all the other models even at k ' 5 , thus requiring at least 2 ˆ fewer DFT calculations.

6 Conclusion

In this work, we investigate how equivariant Transformers can be scaled up to higher degrees of equivariant representations. We start by replacing SO p 3 q convolutions in Equiformer with eSCN convolutions, and propose three architectural improvements to better leverage the power of higher degrees - attention re-normalization, separable S 2 activation and separable layer normalization. With these modifications, we propose EquiformerV2, which outperforms state-of-the-art methods on the S2EF, IS2RS, and IS2RE tasks on the OC20 dataset, improves speed-accuracy trade-offs, and achieves the best success rate when used in AdsorbML.

Broader Impacts. EquiformerV2 achieves more accurate approximation of quantum mechanical calculations and demonstrates one further step toward replacing DFT force fields with machine learned ones. By demonstrating its promising results, we hope to encourage the community to make further progress in applications like material design and drug discovery than to use it for adversarial purposes. Additionally, the method only facilitates identification of molecules or materials of specific properties, and there are substantial hurdles from their large-scale deployment. Finally, we note that the proposed method is general and can be applied to different problems like protein structure prediction [64] as long as inputs can be modeled as 3D graphs.

Limitations. Although EquiformerV2 improves upon state-of-the-art methods on the large and diverse OC20 dataset, we acknolwdge that the performance gains brought by scaling to higher degrees and the proposed architectural improvements can depend on tasks and datasets. For example, the increased expressivity may lead to overfitting on smaller datasets like QM9 [65,66] and MD17 [5658]. However, the issue can be mitigated by pre-training on large datasets like OC20 [30] and PCQM4Mv2 [67] optionally via denoising [68] and then finetuning on smaller datasets.

References

J. Gilmer, S. S. Schoenholz, P. F. Riley, O. Vinyals, and G. E. Dahl, 'Neural message passing for quantum chemistry,' in International Conference on Machine Learning (ICML) , 2017. 1, 2

L. Zhang, J. Han, H. Wang, R. Car, and W. E, 'Deep potential molecular dynamics: A scalable model with the accuracy of quantum mechanics,' Phys. Rev. Lett. , vol. 120, p. 143001, Apr 2018. 1

W. Jia, H. Wang, M. Chen, D. Lu, L. Lin, R. Car, W. E, and L. Zhang, 'Pushing the limit of 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. 1

J. Gasteiger, S. Giri, J. T. Margraf, and S. Günnemann, 'Fast and uncertainty-aware directional message passing for non-equilibrium molecules,' in Machine Learning for Molecules Workshop, NeurIPS , 2020. 1, 3, 8, 9

S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. 366 Smidt, and B. Kozinsky, 'E(3)-equivariant graph neural networks for data-efficient and accurate 367 interatomic potentials,' Nature Communications , vol. 13, May 2022. 1, 2, 3, 4 368

D. Lu, H. Wang, M. Chen, L. Lin, R. Car, W. E, W. Jia, and L. Zhang, '86 pflops deep potential 369 molecular dynamics simulation of 100 million atoms with ab initio accuracy,' Computer Physics 370 Communications , vol. 259, p. 107624, 2021. 1 371

O. T. Unke, M. Bogojeski, M. Gastegger, M. Geiger, T. Smidt, and K. R. Muller, 'SE(3)372 equivariant prediction of molecular wavefunctions and electronic densities,' in Advances in 373 Neural Information Processing Systems (NeurIPS) (A. Beygelzimer, Y. Dauphin, P. Liang, and 374 J. W. Vaughan, eds.), 2021. 1, 2 375

A. Sriram, A. Das, B. M. Wood, and C. L. Zitnick, 'Towards training billion parameter graph 376 neural networks for atomic simulations,' in International Conference on Learning Representa377 tions (ICLR) , 2022. 1, 3, 8 378

J. A. Rackers, L. Tecot, M. Geiger, and T. E. Smidt, 'A recipe for cracking the quantum scaling 379 limit with machine learned electron densities,' Machine Learning: Science and Technology , 380 vol. 4, p. 015027, feb 2023. 1 381

J. Lan, A. Palizhati, M. Shuaibi, B. M. Wood, B. Wander, A. Das, M. Uyttendaele, C. L. Zitnick, 382 and Z. W. Ulissi, 'AdsorbML: Accelerating adsorption energy calculations with machine 383 learning,' arXiv preprint arXiv:2211.16486 , 2022. 1, 2, 7, 8, 9 384

N. Thomas, T. E. Smidt, S. Kearnes, L. Yang, L. Li, K. Kohlhoff, and P. Riley, 'Tensor field 385 networks: Rotation- and translation-equivariant neural networks for 3d point clouds,' arxiv 386 preprint arXiv:1802.08219 , 2018. 1, 2 387

M. Weiler, M. Geiger, M. Welling, W. Boomsma, and T. Cohen, '3D Steerable CNNs: Learning 388 Rotationally Equivariant Features in Volumetric Data,' in Advances in Neural Information 389 Processing Systems 32 , pp. 10402-10413, 2018. 1, 2, 4, 5 390

R. Kondor, Z. Lin, and S. Trivedi, 'Clebsch-gordan nets: a fully fourier space spherical 391 convolutional neural network,' in Advances in Neural Information Processing Systems 32 , 392 pp. 10117-10126, 2018. 1, 2 393

394

395

396

F. Fuchs, D. E. Worrall, V. Fischer, and M. Welling, 'Se(3)-transformers: 3d roto-translation equivariant attention networks,' in Advances in Neural Information Processing Systems (NeurIPS) , 2020. 1, 2

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

J. Brandstetter, R. Hesselink, E. van der Pol, E. J. Bekkers, and M. Welling, 'Geometric and physical quantities improve e(3) equivariant message passing,' in International Conference on Learning Representations (ICLR) , 2022. 1, 2

A. Musaelian, S. Batzner, A. Johansson, L. Sun, C. J. Owen, M. Kornbluth, and B. Kozinsky, 'Learning local equivariant representations for large-scale atomistic dynamics,' arxiv preprint arxiv:2204.05249 , 2022. 1, 2

Y.-L. Liao and T. Smidt, 'Equiformer: Equivariant graph attention transformer for 3d atomistic graphs,' in International Conference on Learning Representations (ICLR) , 2023. 1, 2, 3, 4, 5, 8

S. Passaro and C. L. Zitnick, 'Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs,' in International Conference on Machine Learning (ICML) , 2023. 1, 2, 3, 4, 5, 6, 7, 8

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, 'Attention is all you need,' in Advances in Neural Information Processing Systems (NeurIPS) , 2017. 1, 2, 3

S. Khan, M. Naseer, M. Hayat, S. W. Zamir, F. S. Khan, and M. Shah, 'Transformers in vision: A survey,' arXiv preprint arxiv:2101.01169 , 2021. 1

N. Carion, F. Massa, G. Synnaeve, N. Usunier, A. Kirillov, and S. Zagoruyko, 'End-to-end object detection with transformers,' in European Conference on Computer Vision (ECCV) , 2020. 1

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, J. Uszkoreit, and N. Houlsby, 'An image is worth 16x16 words: Transformers for image recognition at scale,' in International Conference on Learning Representations (ICLR) , 2021. 1, 3

H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou, 'Training data419 efficient image transformers & distillation through attention,' arXiv preprint arXiv:2012.12877 , 420 2020. 1 421

422

423

J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, 'BERT: Pre-training of deep bidirectional transformers for language understanding,' arxiv preprint arxiv:1810.04805 , 2019. 1

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei, 'Language models are few-shot learners,' in Advances in Neural Information Processing Systems (NeurIPS) , 2020. 1

V. P. Dwivedi and X. Bresson, 'A generalization of transformer networks to graphs,' arxiv preprint arxiv:2012.09699 , 2020. 1

D. Kreuzer, D. Beaini, W. L. Hamilton, V. Létourneau, and P. Tossou, 'Rethinking graph transformers with spectral attention,' in Advances in Neural Information Processing Systems (NeurIPS) , 2021. 1

C. Ying, T. Cai, S. Luo, S. Zheng, G. Ke, D. He, Y. Shen, and T.-Y. Liu, 'Do transformers really perform badly for graph representation?,' in Advances in Neural Information Processing Systems (NeurIPS) , 2021. 1

Y. Shi, S. Zheng, G. Ke, Y . Shen, J. You, J. He, S. Luo, C. Liu, D. He, and T.-Y . Liu, 'Benchmarking graphormer on large-scale molecular modeling datasets,' arxiv preprint arxiv:2203.04810 , 2022. 1

L. Chanussot*, A. Das*, S. Goyal*, T. Lavril*, M. Shuaibi*, M. Riviere, K. Tran, J. HerasDomingo, C. Ho, W. Hu, A. Palizhati, A. Sriram, B. Wood, J. Yoon, D. Parikh, C. L. Zitnick, and Z. Ulissi, 'Open catalyst 2020 (oc20) dataset and community challenges,' ACS Catalysis , 2021. 2, 3, 4, 6, 7, 9

B. K. Miller, M. Geiger, T. E. Smidt, and F. Noé, 'Relevance of rotationally equivariant convolutions for predicting molecular properties,' arxiv preprint arxiv:2008.08461 , 2020. 2

R. J. L. Townshend, B. Townshend, S. Eismann, and R. O. Dror, 'Geometric prediction: Moving beyond scalars,' arXiv preprint arXiv:2006.14163 , 2020. 2

B. Jing, S. Eismann, P. Suriana, R. J. L. Townshend, and R. Dror, 'Learning from protein structure with geometric vector perceptrons,' in International Conference on Learning Representations (ICLR) , 2021. 2

K. T. Schütt, O. T. Unke, and M. Gastegger, 'Equivariant message passing for the prediction of tensorial properties and molecular spectra,' in International Conference on Machine Learning (ICML) , 2021. 2, 4, 9

V. G. Satorras, E. Hoogeboom, and M. Welling, 'E(n) equivariant graph neural networks,' in International Conference on Machine Learning (ICML) , 2021. 2

P. Thölke and G. D. Fabritiis, 'Equivariant transformers for neural network based molecular potentials,' in International Conference on Learning Representations (ICLR) , 2022. 2

T. Le, F. Noé, and D.-A. Clevert, 'Equivariant graph attention networks for molecular property prediction,' arXiv preprint arXiv:2202.09891 , 2022. 2

I. Batatia, D. P. Kovacs, G. N. C. Simm, C. Ortner, and G. Csanyi, 'MACE: Higher order equivariant message passing neural networks for fast and accurate force fields,' in Advances in Neural Information Processing Systems (NeurIPS) , 2022. 2

A. Sanchez-Gonzalez, J. Godwin, T. Pfaff, R. Ying, J. Leskovec, and P. W. Battaglia, 'Learning to simulate complex physics with graph networks,' in International Conference on Machine Learning (ICML) , 2020. 2

N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov, 'Dropout: A simple way to prevent neural networks from overfitting,' Journal of Machine Learning Research , vol. 15, no. 56, pp. 1929-1958, 2014. 2, 4

G. Huang, Y. Sun, Z. Liu, D. Sedra, and K. Q. Weinberger, 'Deep networks with stochastic 470 depth,' in European Conference on Computer Vision (ECCV) , 2016. 2 471

L. Zitnick, A. Das, A. Kolluru, J. Lan, M. Shuaibi, A. Sriram, Z. Ulissi, and B. Wood, 'Spherical 472 channels for modeling atomic interactions,' in Advances in Neural Information Processing 473 Systems (NeurIPS) , 2022. 2, 3, 5, 7, 8, 9 474

K. T. Schütt, P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, and K.-R. Müller, 475 'Schnet: A continuous-filter convolutional neural network for modeling quantum interactions,' 476 in Advances in Neural Information Processing Systems (NeurIPS) , 2017. 3, 6, 8, 9 477

T. Xie and J. C. Grossman, 'Crystal graph convolutional neural networks for an accurate and 478 interpretable prediction of material properties,' Physical Review Letters , 2018. 3, 8 479

O. T. Unke and M. Meuwly, 'PhysNet: A neural network for predicting energies, forces, 480 dipole moments, and partial charges,' Journal of Chemical Theory and Computation , vol. 15, 481 pp. 3678-3693, may 2019. 3 482

J. Gasteiger, J. Groß, and S. Günnemann, 'Directional message passing for molecular graphs,' 483 in International Conference on Learning Representations (ICLR) , 2020. 3 484

Z. Qiao, M. Welborn, A. Anandkumar, F. R. Manby, and T. F. Miller, 'OrbNet: Deep learning for 485 quantum chemistry using symmetry-adapted atomic-orbital features,' The Journal of Chemical 486 Physics , 2020. 3 487

Y. Liu, L. Wang, M. Liu, Y. Lin, X. Zhang, B. Oztekin, and S. Ji, 'Spherical message passing 488 for 3d molecular graphs,' in International Conference on Learning Representations (ICLR) , 489 2022. 3 490

M. Shuaibi, A. Kolluru, A. Das, A. Grover, A. Sriram, Z. Ulissi, and C. L. Zitnick, 'Rotation 491 invariant graph neural networks using spin convolutions,' arxiv preprint arxiv:2106.09575 , 492 2021. 3, 8 493

J. Klicpera, F. Becker, and S. Günnemann, 'Gemnet: Universal directional graph neural 494 networks for molecules,' in Advances in Neural Information Processing Systems (NeurIPS) , 495 2021. 3, 8 496

J. Gasteiger, M. Shuaibi, A. Sriram, S. Günnemann, Z. Ulissi, C. L. Zitnick, and A. Das, 497 'GemNet-OC: Developing Graph Neural Networks for Large and Diverse Molecular Simulation 498 Datasets,' Transactions on Machine Learning Research (TMLR) , 2022. 3, 8, 9 499

J. L. Ba, J. R. Kiros, and G. E. Hinton, 'Layer normalization,' arxiv preprint arxiv:1607.06450 , 500 2016. 4, 5 501

P. Veliˇ ckovi´ c, G. Cucurull, A. Casanova, A. Romero, P. Liò, and Y. Bengio, 'Graph attention 502 networks,' in International Conference on Learning Representations (ICLR) , 2018. 4 503

G. Huang, Y. Sun, Z. Liu, D. Sedra, and K. Q. Weinberger, 'Deep networks with stochastic 504 depth,' in European Conference on Computer Vision (ECCV) , 2016. 4 505

506

507

508

M. Geiger, T. Smidt, A. M., B. K. Miller, W. Boomsma, B. Dice, K. Lapchevskyi, M. Weiler, M. Tyszkiewicz, S. Batzner, D. Madisetti, M. Uhrin, J. Frellsen, N. Jung, S. Sanborn, M. Wen, J. Rackers, M. Rød, and M. Bailey, 'e3nn/e3nn: 2022-04-13,' Apr. 2022. 4

509

510

511

512

513

514

515

516

S. Chmiela, A. Tkatchenko, H. E. Sauceda, I. Poltavsky, K. T. Schütt, and K.-R. Müller, 'Machine learning of accurate energy-conserving molecular force fields,' Science Advances , vol. 3, no. 5, p. e1603015, 2017. 4, 9

K. T. Schütt, F. Arbabzadah, S. Chmiela, K. R. Müller, and A. Tkatchenko, 'Quantum-chemical insights from deep tensor neural networks,' Nature Communications , vol. 8, jan 2017. 4, 9

S. Chmiela, H. E. Sauceda, K.-R. Müller, and A. Tkatchenko, 'Towards exact molecular dynamics simulations with machine-learned force fields,' Nature Communications , vol. 9, sep 2018. 4, 9

T. S. Cohen, M. Geiger, J. Köhler, and M. Welling, 'Spherical CNNs,' in International Confer517 ence on Learning Representations (ICLR) , 2018. 5 518

S. Elfwing, E. Uchibe, and K. Doya, 'Sigmoid-weighted linear units for neural network function 519 approximation in reinforcement learning,' arXiv preprint arXiv:1702.03118 , 2017. 5 520

P. Ramachandran, B. Zoph, and Q. V. Le, 'Searching for activation functions,' arXiv preprint 521 arXiv:1710.05941 , 2017. 5 522

B. Hammer, L. B. Hansen, and J. K. Nørskov, 'Improved adsorption energetics within density523 functional theory using revised perdew-burke-ernzerhof functionals,' Phys. Rev. B , 1999. 6 524

W. Hu, M. Shuaibi, A. Das, S. Goyal, A. Sriram, J. Leskovec, D. Parikh, and C. L. Zit525 nick, 'Forcenet: A graph neural network for large-scale quantum calculations,' arxiv preprint 526 arxiv:2103.01436 , 2021. 8 527

J. H. Lee, P. Yadollahpour, A. Watkins, N. C. Frey, A. Leaver-Fay, S. Ra, K. Cho, V. Gligorijevic, 528 A. Regev, and R. Bonneau, 'Equifold: Protein structure prediction with a novel coarse-grained 529 structure representation,' bioRxiv , 2022. 9 530

R. Ramakrishnan, P. O. Dral, M. Rupp, and O. A. von Lilienfeld, 'Quantum chemistry structures 531 and properties of 134 kilo molecules,' Scientific Data , vol. 1, 2014. 9 532

533

534

535

L. Ruddigkeit, R. van Deursen, L. C. Blum, and J.-L. Reymond, 'Enumeration of 166 billion organic small molecules in the chemical universe database gdb-17,' Journal of Chemical Information and Modeling , vol. 52, no. 11, pp. 2864-2875, 2012. PMID: 23088335. 9

536

537

538

539

540

541

542

M. Nakata and T. Shimazaki, 'Pubchemqc project: A large-scale first-principles electronic structure database for data-driven chemistry,' Journal of chemical information and modeling , vol. 57 6, pp. 1300-1308, 2017. 9

S. Zaidi, M. Schaarschmidt, J. Martens, H. Kim, Y. W. Teh, A. Sanchez-Gonzalez, P. Battaglia, R. Pascanu, and J. Godwin, 'Pre-training via denoising for molecular property prediction,' in International Conference on Learning Representations (ICLR) , 2023. 9

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose EquiformerV2, which is an improvement over the original Equiformer architecture. The main improvement is using a more efficient parameterization of the tensor products used in Equiformer which is computationally expensive for higher order representations. The more efficient parameterization involves using SO(2) linear layers instead of SO(3) tensor products. Moreover, they have three architectural improvements – adding an extra layer norm during attention, S2 activations instead of gates and \"separable\" layer normalization which differs from the previous layer normalization in terms of the denominator used. They perform experiments on OC20 to show improvements compared to other models in the literature.\n\nSTRENGTHS:\nOverall, the idea of using computing tensor products more efficiently using SO(2) linear layers makes good sense in its application to Equiformer in order to scale it up via higher order SO(3) representations. Empirical work around improving the architecture by appropriate use of different type of layer norms and activation functions is also valuable. The components feeding into this model aren't always original, but that shouldn't detract from the significance of putting together the full model and showing that it improves performance on a benchmark. The clarity of the writing is usually good but there are places where it can be improved (see below).\n\nWEAKNESSES:\nThe main weakness is insufficient comparison against the original Equiformer architecture. There are four differences versus the V1 architecture as far as I can tell: (1) the more efficient parameterization of tensor products via the ideas of eSCN, (2) an extra layer norm applied to scalar features in the attention module, (3) separable S2 activation instead of gates and (4) layer normalization with a different way of computing the denominator (called \"separable\" layer normalization in the paper). \n\n- There is an ablation in Table 1(a) on one of the datasets looking at the effects of changes (2), (3), and (4). It is unclear how much of the difference in performance (specially, on energies) is statistically significant. I understand training each model is expensive, but since the experiments in this table are on the smaller S2EF-2M dataset, isn't it possible to have error bars here?\n- Why isn't there a comparison in Table 1(a) of the effect of incorporating change (1) in EquiformerV2? The SO(2) linear layers replacing the tensor product are in principle as expressive as the tensor products but the optimization dynamics can be different since the parameterization of the architecture is different. It would be good to see a side-by-side comparison of Equiformer's V1 and V2 when keeping other architectural hyperparameters fixed (e.g. number of channels, maximum representation order etc.)\n- It would be good to include EquiformerV1 performance in the other sub-tables of Table 1 too, in order to see the cummulative effect of the changes to the original architecture.\n- Similarly, it would be useful to see EquiformerV1 performance in Table 2 and 3 as well (possibly with a lower $L_\\text{max}$ but with other hyperparameters tuned).\n- How does the setup for IS2RE in Table 2 differ from EquiformerV1's Table 3? Is one with relaxations after training on S2EF and the other is through direct energy prediction? A side-by-side comparison under the same setup would help a lot here. \n- How necessary are the higher orders compared to increasing the number of layers or channels? SO(2) linear layers instead of tensor products will improve the efficiency for lower orders too, so can we make up for the performance of higher orders by instead increasing the number of layers or channels?\n- Have you benchmarked EquiformerV2 on QM9 where EquiformerV1 has already been benchmarked? It's a dataset that is different from OC20 in various ways, so it would be useful to see a comparison.\n\nFor other suggestions for improvements, see the questions below. \n\nEquiformer(V1): https://arxiv.org/abs/2206.11990\n\nQUESTIONS:\n- The explanation of separable S2 activation was unclear: what does line 186 mean? \"The activation first converts vectors of all degrees to point samples on a sphere for each channel, applies unconstrained functions F to those samples, and finally convert them back to vectors.\" \n- Have you done direct IS2RE prediction as done in the EquiformerV1 paper with/without an auxiliary denoising loss?\n- Can you justify \"higher degrees can better capture angular resolution and directional information\"? Perhaps via an example?\n- Was stochastic depth used in EquiformerV2? Related work section mentions it was used in EquiformerV1?\n- Line 69: \"apply typical functions to rotated features during message passing\" was unclear to me.\n- Line 142: \"eSCN convolutions go one step further and replace the remaining non-trivial paths of the SO(3) tensor product with an SO(2) linear operation to allow for additional parameters of interaction between $\\pm m_o$ without breaking equivariance.\" Can you explain this in the context of Appendix A.3?\n- Line 175: why would $f^{(0)}_{ij}$ be less well-normalized?\n- Line 178: there is a slight overload in notation with $a$ and $a_{ij}$.\n- Line 210: why is there no centering in the layer norm formulation? Wouldn't that continue to preserve equivariance as the mean vector belongs to the same representation? There is centering in the 0th degree activations and traditional layer norm.\n- Line 210: do you use a small $\\epsilon$ in the new separable layer normalization denominator for numerical stability?\n- Line 280: \"Finally, these modifications enable training for longer without overfitting (row 7)\" don't both models improve performance when trained for longer, not just the model in row 7?\n- On comparison of speed: are all the models using the same training/inference pipeline such that the comparison is fair? It would be useful to report inference cost in FLOPs or a similar metric as well.\n\nLIMITATIONS:\nLimitations and broader impact are adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper provides a new equivariant graph neural networks named EquiformerV2 to enhance the original Equiformer performance. It uses four new modules. The first module is to use the convolution in the eSCN (https://arxiv.org/abs/2302.03655) to replace the depth-wise tensor production accelerating the speed. The second module is the separable $S^{2}$ activation using the spherical grids in SCN (https://arxiv.org/abs/2206.14331) to encourage the non-linearity. The third module is the separable layer normalization which uses the variance of all equivariant feature $\\ell \\geq 1$ to do normalization. The fourth module is the attention re-normalization using a layer normalization before the non-linear function of the attention score branch.\n\nSTRENGTHS:\n1. This paper is well organized and written. The figure clearly shows the modification on the model architecture.\n2. The experiments in Table 2 supports the proposed model architecture can achieve SOTA performance on the OC20 All training sets as well as OC20 All+MD. For example, on OC20 All, test energy MAE is improved 13meV in S2ET test and 25meV in IS2RE test. Such improvement is great. These two datasets usually take extensive training time to perform experiments on them. From the Throughput metric, the EquiformerV2 has better training efficiency compared to current baselines.\n\nWEAKNESSES:\n1. As a suggestion, it will be better if efficiency study includes both training time and inference time. Computational complexity metric such as FLOPs or MACs can help measure the inference complexity.\n2. The description of spherical grid is not very clear in the paper. Although it is introduced in the SCN paper, I think a brief introduction can help people understand why such operation can enhance the non-linearity.\n\nQUESTIONS:\n1. For the speed up of DFT calculations, could you briefly introduce the AdsorbML and explain why EquiformerV2 can be used to accelerate the DFT?\n2. Is it possible to apply such efficient equivariant architecture in directly accelerating the DFT calculations such as the SchNorb (https://www.nature.com/articles/s41467-019-12875-2), PhiSNet (https://arxiv.org/abs/2106.02347), QHNet (https://arxiv.org/abs/2306.04922) and recent QH9 dataset (https://arxiv.org/abs/2306.09549)? Since the computational complexity is also a problem for these equivariant networks.\n3. For the equivariance, since the equiformerV2 uses the spherical grids, I am curous about whether equiformerV2 is rigorous equivariant or approximate equivariant. If this architecture is approximate equivariant, I think it will be great if there are some experiments to verify the equivariance.\n\nLIMITATIONS:\n1. As an suggestion, to comprehensively study the improvement on original Equiformer, it will be better if there is a comprehensive experiments on the original datasets such as QM9.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors proposed EquiformerV2, which is an equivariant network for 3D molecular modeling. The EquiformerV2 is built on the Equiformer with several architectural modifications: 1) replace SO(3) convolutions (tensor product operations) with efficient SO(2) counterparts from eSCN; 2) Attention Re-normalization; 3) Separable S2 activation; 4) separate Layer Normalization. These changes enable EquiformerV2 to achieve good performance on the large-scale OC20 benchmark and also the new AdsorbML dataset.\n\nSTRENGTHS:\n1. The targeted problem is of great interest to the community. The EquiformerV2 provides another attempt to enlarge the maximum degree of irreducible representations and obtain performance gains on large-scale DFT benchmarks.\n\n2. Good empirical performance. On the OC20 benchmark, the EquiformerV2 achieves state-of-the-art performance on the Structure-to-Energy-Force task. The model trained on this task further serves as a force-field evaluator to achieve strong performance on IS2RS and IS2RE tasks. The EquiformerV2 outperforms the compared baselines on all these tasks, especially on force prediction. Additionally, it also improves the success rate a lot on the AdsorbML dataset.\n\nWEAKNESSES:\n1. The novelty of integrating the eSCN convolution and S^2 activation into the Equiformer is limited. Among the proposed architectural changes, the eSCN convolution is the key component to enable Equiformer to use irreducible representations of higher degrees, and the S^2 activation also replaces all non-linear activations. However, these design strategies should be mainly credited to the eSCN work.\n\n2. The motivation for the other architectural modifications should be thoroughly clarified. First, lines 174-175 in Section 4.2 suggest the \"less well-normalized\" issue, which motivated the authors to propose re-normalization and Separable Layer Normalization. It is better to provide further quantitive evidence to reveal how such an issue affects the model's performance, and why these modifications could remove or mitigate such effect. Second, the authors proposed Separable S^2 Activation because the original S^2 activation would make the training process diverge. However, it is hard to understand why such separable modifications could make the training process stable. Is there any further essential reason behind such a phenomenon? It is suggested to provide further analysis on such modifications.\n\n3. More analyses are necessary if the computational resources are acceptable. First, the authors did not provide a performance comparison between whether using the eSCN convolution. Although the eSCN convolution is equivalent to the SO(3) counterpart, the computation processes of these two parameterizations are different. Second, it is suggested to further report the number of model parameters and memory costs for the EquiformerV2 and the compared baselines in Table 2. Third, how does EquiformerV2 perform on the IS2RS and IS2RE using the direct setting (use or not the denoising setting) that is the same as EquiformerV1?\n\nOverall, the major weakness of this work lies in the novelty, unclear motivations, and incomplete analyses. If the authors could well address the above concerns, I would like to increase my scores.\n\nQUESTIONS:\n1. Towards the claim \"Higher degrees can better capture angular resolution and directional information, which is critical to accurate prediction of atomic energies and forces\". [Lines 32-33, also see Lines 97-98], could you further provide explanations and evidence from deep learning molecular models?\n\n2. How is the S^2 activation implemented? Is there any sampling process inducing randomness or incompleteness to make the module not strictly equivariant? If so, is there any measurement of such approximated equivariance?\n\n3. Could you briefly introduce how the EquiformerV2 model accelerates the DFT calculation in AdsorbML? Are there further quantitative results on the error measure (e.g., MAE) between the EquiformerV2 predictions and DFT labels?\n\nLIMITATIONS:\nThe authors carefully discuss the broader impact and limitations of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper propose EquiformerV2, a advance verison based on Equiformer and eSCN structure extend to higher degree representations, which achieve better performance in force and energy tasks.\n\nSTRENGTHS:\nThis paper is well-written and organized, presenting a clear and coherent structure throughout. The introduced EquiformerV2, an upgraded version of the original Equiformer, including three architectural improvements: attention re-normalization, separable S^2 activation, and separable layer normalization. These enhancements contribute to the SOTA performance in  OC20 dataset. \nThe proposed model achieves a high degree representation with efficiency, as outlined in the paper. The authors also provide a comprehensive ablation study to support the necessity of these modifications, effectively highlighting their respective contributions to the overall performance improvement.\n\nWEAKNESSES:\n* A few spelling errors. For instance, in Section 6, the word \"acknolwdge\" etc.  Along with any other mistakes found throughout the manuscript.\n\n* The experiments conducted in this study primarily utilize the OC20 dataset. While this dataset is relevant, it is essential to note that there are various other DFT-based datasets available that could provide a more comprehensive evaluation of the proposed architecture. Such as OC22, OQMD[1,2], SPICE[3], and PCQM4Mv2 etc.\n\n[1]  Saal, J. E., Kirklin, S., Aykol, M., Meredig, B., and Wolverton, C. \"Materials Design and Discovery with High-Throughput Density Functional Theory: The Open Quantum Materials Database (OQMD)\", JOM 65, 1501-1509 (2013). \n[2]. Kirklin, S., Saal, J.E., Meredig, B., Thompson, A., Doak, J.W., Aykol, M., Rühl, S. and Wolverton, C. \"The Open Quantum Materials Database (OQMD): assessing the accuracy of DFT formation energies\", npj Computational Materials 1, 15010 (2015).\n[3] SPICE, A Dataset of Drug-like Molecules and Peptides for Training Machine Learning Potentials\n\nQUESTIONS:\n* Would it be possible for you to include the results from the small dataset, as you mentioned? This would provide valuable insights into the performance and scalability of your proposed approach.\n* More details about structural relaxations, specifically, the relax trajectories and time efficiency.\n\nLIMITATIONS:\nsame above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-16"}