Equivariant Transformers for Neural Network based Molecular Potentials
Authors: Philipp Thölke
Gianni De Fabritiis
Venue: ICLR 2022 Spotlight (2022)
Keywords: Molecular Modeling
Quantum Chemistry
Attention
Transformers

Abstract: The prediction of quantum mechanical properties is historically plagued by a trade-off between accuracy and speed. Machine learning potentials have previously shown great success in this domain, reaching increasingly better accuracy while maintaining computational efficiency comparable with classical force fields. In this work we propose TorchMD-NET, a novel equivariant Transformer (ET) architecture, outperforming state-of-the-art on MD17, ANI-1, and many QM9 targets in both accuracy and computational efficiency. Through an extensive attention weight analysis, we gain valuable insights into the black box predictor and show differences in the learned representation of conformers versus conformations sampled from molecular dynamics or normal modes. Furthermore, we highlight the importance of datasets including off-equilibrium conformations for the evaluation of molecular potentials.

Published as a conference paper at ICLR 2022
TORCHMD-NET: EQUIVARIANT A TRANSFORMERS FOR NEURAL NETWORK BASED MOLECULAR POTENTIALS
Philipp Thölke
Computational Science Laboratory. Pompeu Fabra University, PRBB, C/ Doctor Aiguader 88, 08003 Barcelona, Spain and Institute of Cognitive Science, Osnabrück University, Neuer Graben 29/ / Schloss, 49074 Osnabrüick, Germany philipp thoe a ke ke@posteo.cccmm de
Gianni De Fabritiis
Computational Science Laboratory, Pompeu Fabra University, C/ Doctor Aiguader 88, 08003 Barcelona, Spain and ICREA, Passeig Lluis Companys 23, 08010 Barcelona, Spain and Acellera Labs, C/ Doctor Trueta 183, 08005 Barcelona, Spain giann de
ABSTRACT
The prediction of quantum mechanical properties is historically plagued by a trade-off between accuracy and speed. Machine learning potentials have previ- ously shown great success in this domain, reaching increasingly better accuracy while maintaining computational efficiency comparable with classical force fields. In this work we propose TorchMD-NET, a novel equivariant Transformer (ET) architecture, outperforming state-of-the on MD17, ANI-1, and many QM9 tar- gets in both accuracy and computational efficiency. Through an extensive atten- tion weight analysis, we gain valuable insights into the black box predictor and show differences in the learned representation of conformers versus conforma- tions sampled from molecular dynamics or normal modes. Furthermore, we high- light the importance of datasets including off-equilibrium conformations for the evaluation molecular potentials.
INTRODUCTION
Quantum mechanics are essential for the computational analysis and design of molecules and materi- als. However, the complete solution of the Schrödinger equation is analytically and computationally not practical, which initiated the study of approximations in the past decades (Szabo & Ostlund, 1996). A common quantum mechanics approximation method is to model atomie systems accord- ing to density functional theory (DFT), which can provide energy estimates with sufficiently high accuracy for different application cases in biology, physics, chemistry, and materials science. Even more accurate techniques like coupled-cluster exist but both still lack the computational efficiency be applied on larger scale, although recent advances are promising in the case of quantum Monte Carlo (Pfau et al., 2020 Hermann et al. 2020). Other methods include force-field and semi- empirical quantum mechanical theories, which provide very efficient estimates but lack accuracy.
The field of machine learning molecular potentials is relatively novel. The first important contri- (BP) representation (Behler & Parrinello, 2007) and the
butions are rooted in the seminal work from Rupp et al. (2012). One of the best transferable machine learning potentials for biomolecules, called ANI (Smith et al. (a). is based on BP. second class of methods, mainly developed in the field of materials science and quantum chemistry, uses more modern graph convo- lutions (Schütt et al. 2018 Unke & Meuwly, 2019 Qiao et al., 2020 Schütt et al., 2021). SchNet (Schütt et al.] 2017b 2018), for example, uses continuous filter convolutions in graph network architecture to predict the energy of a system and computes forces by direct differentiation of the neural network against atomic coordinates. Outside of its original use case, this approach has been

---

Published as a conference paper at ICLR 2022
extended to coupled-cluster solvers (Hermann al. 2020) and protein folding using coarse- -aaanned systems (Wang et al. 2019 Husic et al. 2020; Doerr et al., 2021]. Recently, other work has shown that shift towards rotationally equivariant networks (Anderson et al.] 2019 Fuchs al., 2020) Schütt et al. 2021), particularly useful when the predicted quantities are vectors and tensors, can
also improve the accuracy scalars (e.g. energy).
Next to the parametric group of neural network based methods, a nonparametric class of approaches exists. These are usually based on kernel methods, particularly used in materials science. In this work, we will focus on parametric neural network potentials (NNPs) because they have scaling advantage large amounts of data, while kernel methods usually work best scarce data regime.
Previous deep learning based work in the domain of quantum chemistry focused largely on graph neural network architectures (GNNs) with different levels of handcrafted and learned features (Schütt et al. 2017b; Qiao et al. 2020 Klicpera et al. 2020b) [Unke & Meuwly 2019 Liu et al.) A020|Schübl et al. 2021). For example, Qiao et al. (020)) first perform a low-cost mean-
mean-field field elec- tronic structure calculation, from which different quantities are used as input to their neural network. Recently proposed neural network architectures in this context usually include some form of atten- tion (Luong et al. 2015) inside the GNN's message passing step (Qiao et al. 2020 Unke & Meuwly. 2019 Liu et al., 2020).
In this work, we introduce TorchMD -EET, equivariant Transformer (ET) architecture for the prediction of quantum mechanical properties. By building on top of the Transformer (Vaswani et al. 2017) architecture, state-of-the-art accuracy on multiple benchmarks while relying solely on learned featurization of atomic types and coordinates. Furthermore, we gain insights into the black box prediction of neural networks by analyzing the Transformer's attention weights and comparing latent representa- tions between different types of data such as energy- meegyy-iiimmzzd (QM9 (Ramakrishnan al. 2014). ),
centering the design around the attention mechanism, achieving
molecular dynamics (MD17 (Chmicla et al. 2017) and normal mode sampled data (ANI-1 (Smith
al. 2017b)).
METHODS
The traditional Transformer architecture as proposed by Vaswani et al. (2017) operates sequence of tokens. In the context of chemistry, however, the natural data structure for the representation of molecules is graph. To work on graphs, one interpret self-attention as constructing fully con- nected graph over input tokens and computing interactions between nodes. We leverage this concept and extend include information stored in the graph's edges, corresponding to interatomic dis- tances in the context of molecular data. This requires the use of modified attention mechanism, which we introduce in the following sections, along with the overall architecture of our equivariant Transformer.
The equivariant Transformer is made up of three main blocks. An embedding layer encodes atom types Z and the atomic neighborhood of each atom into a dense feature vector r,, Then, a series of update layers compute interactions between pairs of atoms through modified multi-head attention mechanism, with which the latent atomic representations are updated. Finally, a layer normalization (Ba et al. 2016) followed by an output network computes scalar atomwise predictions using gated equivariant blocks (Weiler et al. 2018 Schütt et al. 2021), which get aggregated into single molec- ular prediction. This can be matched with a scalar target variable or differentiated against atomic coordinates, providing force predictions. An illustration of the architecture is given in Figure [
2.1 NOTATION
To differentiate between the concepts of scalar and vector features, this work follows a certain no- tation. Scalar features are written as R while refer to vector features ग 33FF", The vector norm ||||| and scalar product (..)) of vector features are applied to the spatial dimension, while all other operations act the feature dimension. Upper case letters denote matrices A E N
2

---

Published as a conference paper at ICLR 2022
a
(b)
Z
R
X,
Tij
Fi
-
Tij
z
V = X embed
attention 344 By
RBF
U1
U2
U3
update
Q
K
D
D
update
a
a(-)
V
update
output network
Σj
O
(EE))
Σ
update
attention
A
A
3/1
ij
Figure I: Overview the equivariant Transformer architecture. Thin lines: scalar features in R.. thick lines: vector features in R F dashed lines: multiple feature vectors. (a) Transformer consist- ing of an embedding layer, update layers and an output network. (b) Residual update layer including attention based interatomic interactions and information exchange between scalar and vector fea- tures. (c) Modified dot-product attention mechanism, scaling values (blue) by the attention weights (red).
2.2 EMBEDDING LAYER
The embedding layer assigns two learned vectors each atom type 2,. One is used to encode infor- mation specific to an atom, the other takes the role of a neighborhood embedding. The neighborhood embedding, which an embedding of the types of neighboring atoms, is multiplied by a distance filter. This operation resembles a continuous-filce convolution (Schütt et al.) (017b) but, as it is used in the first layer, allows the model store atomic information in two separate weight matrices. These can be thought of as containing information that is intrinsic to an atom versus information about the interaction of two atoms.
The distance filter is generated from expanded interatomic distances using linear transformation F.F. First, the distance diy between two atoms i and is expanded via a set of exponential normal radial basis functions RRF defined as
cl
RBF (BFF(dd) = o(dg)
-x)))
(1)
where Se and μa are fixed parameters specifying the center and width of radial basis function k. The y vector initialized with values equally spaced between exp(-dent) and 1, B is initialized as (2K-1(1 exp(- de for all h as proposed by Unke & Meuwly (2019). The cutoff dis-
tance deut was set to 5Á. The cosine cutoff o(d;)) ) is used to ensure a smooth transition to 0 as d: approaches deut in order to avoid jumps in the regression landscape. It is given by
1
cos
Td2y
+1),
if dij S deut.
((deg)
deut
(2)
0.
if dij > deat.
The neighborhood embedding n, for atom then defined
Σ
embed (zz) @ W RBF (diz)
with ubb being the neighborhood embedding function and N the number of atoms in the graph. The final atomic embedding x; is calculated as a linear projection of the concatenated intrin-
sic embedding and neighborhood embedding [embed (zz), nz] resulting in
C [embod (2+). n²| +
(4)
with inb being the intrinsic embedding function. The vector features i are initially set to 0.
3

---

Published as a conference paper at ICLR 2022
2.3 MODIFIED ATTENTION MECHANISM
We use a modified multi-head attention mechanism (F) extending dot-product attention, in order include edge data into the calculation of attention weights. First, the feature vectors are passed through layer normalization. Then, edge data, i.e. interatomic distances raj, are projected into two multidimensional filters D and DV, according to
D = o (W RBF (riy) + BD*)
DV o (WD' RBF (riz) + 600)
The attention weights are computed via an extended dot product, i.e. an elementwise multiplication and subsequent sum over the feature dimension, of the three input vectors: query Q, key K and distance projection D*;
Q = W and K = W x4
(6)
(o,
Σ
Qs Q KR o DE
The resulting matrix is passed through nonlinear activation function and is weighted by cosine cutoff & (see equation? 2). ensuring that atoms with a distance larger than deut do not interact.
A = (iL K, D))) ((di))
(8)
Traditionally, the resulting attention matrix A is passed through a softmax activation, however, replace this step with a SiLU function to preserve the distance cutoff. The softmax scaling factor of Vd which normally rescales small gradients from the softmax function, is left out. Work by Choromanski et al. (202I) suggests that replacing the softmax activation function in Transformers with ReLU-like functions might even improve accuracy, supporting the idea of switching to SiLU this case.
We place continuous filter graph convolution (Schütt et al., 2017b) in the attention mechanism's ' value pathway. This enables the model to not only consider interatomic distances in the attention weights but also incorporate this information into the feature vectors directly. The resulting repre- sentation is split into three equally sized vectors s/y- & by E FR.. The vector s is scaled by the attention matrix A and aggregated the value-dimension, leading to updated list of feature vectors. The linear transformation is used to combine the attention heads' outputs into a single feature vector yi R384
3
$ij S sis split(V, @ DV m))
y/ =
Σ
Aij
The attention mechanism's output, therefore, corresponds to the updated scalar feature vectors y/ and scalar filters sty and siy. which are used to weight the directional information inside the update layer.
2.4 UPDATE LAYER
The update layer (Figure [[)) is used to compute interactions between atoms (attention block) and exchange information between scalar and vector features. The updated scalar features yi from the attention block are split up into three feature vectors q/ 9f qf E IRF. The first feature vector, g!' takes the role of residual around the scaled vector features. The resulting scalar feature update Ax; of this update layer is then defined as
Ar, = qt + 2 (U10, U27))
(10)
where (U, úg U2ū1) denotes the scalar product of vector features v, transformed by linear projec- tions U, and Uz.
4

---

Published as a conference paper at ICLR 2022
On the side of the vector features, scalar information is introduced through multiplication between q. and a linear projection of the vector features Uz The representation is updated with equivariant features using the directional vector between two atoms. The edge-wise directional information is multiplied with scalar filter si and added to the rescaled vector features sig Ty. The result is aggregated inside each atom, forming wi The final vector feature update At, for the current update layer is then produced by adding the weighted scalar features to the equivariant features w;
N mi
rj F. - Fy
Σ
(11)
3
A Cे + UU
2.5 TRAINING
Models are trained using mean squared error loss and the Adam optimizer (Kingma & Ba 2017)
with parameters 3 = 0.9, 3g = 0.999 and e = 10-8 Linear learning rate warm-up is applied as suggested by Vaswani et al. 2017) by scaling the learning rate with & Patego period, we systematically decrease the learning rate by scaling with decay factor upon reaching plateau in validation loss. The learning rate is decreased down to a minimum of 10- We found that weight decay and dropout do not improve generalization and forces, we apply exponential smoothing to the energy's train and validation loss. New losses are discounted with factor of a = 0.05. See Appendix A for a more comprehensive summary of hyperparameters.
step After the warm-up
this context. When training on energies
EXPERIMENTS AND RESULTS
We evaluate the equivariant Transformer on the QM9 (Ramakrishnan al. 2014), MD17 (Chmiela iet al. (017)) and ANI-1 (Smith et al., 2017b) benchmark datasets. QM9 comprises 133,885 small orgamic molecules with up to nine heavy atoms of type O, N. and F. It reports computed geometric, thermodynamic, energetic, and electronic properties for locally optimized geometries. As suggested by the authors, we used a revised version of the dataset, which excludes 3,054 molecules due to failed geometric consistency checks. The remaining molecules were split into a training set with 110,000 and a validation set with 10,000 samples, leaving 10,831 samples for testing.
Table compares the equivariant Transformer's results QM9 with the invariant architectures SchNet (Schitt et al. (2188). PhysNet (Unke & Meuwly 2019) and DimeNet++ (Klicpera et al.) 2020a), the covariant Cormorant (Anderson al.) 2019) architecture and equivariant EGNN (Sator- ras et al. 2021), Lic Tirrnssffrerr (we compare to their best variant, LicTransformer T3+S03 Aug) (Hutchinson al. 2020) and PaiNN (Schütt et al.) 2021). We use specialized output models for two of the QM9 targets, which add certain features directly to the prediction. For the molecular dipole moment μr, both scalar and vector features are used in the final calculation. The output MLP consists of two gated equivariant blocks (Weiler et al. (2018) Schütt et al. (021)) with the same layer sizes as in the otherwise used output network. The updated scalar features v; and vector features i, are then used to compute
Σ Th + + (T r)
(12)
where F is the center of mass of the molecule. For the prediction of the electronic spatial extent ( (R?). scalar features are transformed using gated equivariant blocks as described above, yielding scalar atomic predictions z4, and multiplied by the squared norm of atomic positions
N
RR2)
Σ
x4 1
(13)
The MD17 dataset consists of molecular dynamics (MD) trajectories of small organic molecules, including both energies and forces. In order to guarantee conservation of energy, forces are predicted using the negative gradient of the energy with respect atomic coordinates F, = - To
5

---

Published as a conference paper at ICLR 2022
Table 1: Results on all QM9 targets and comparison to previous work. Scores are reported as mean absolute errors (MAE). LieTF refers to the best performing variant of Lie'rrassformet (Hutchinson
al. 2020), i.e. TieTrassforeer- T3+S33 Aug.
Target
Unit
SchNet
EGNN
PhysNet
LieTF
DimeNet++
Cormorant
PaiNN
ET
D
0.033
0.029
0.0529
0.041
0.0297
0.038
0.012
0.011
0.235
0.071
0.0615
0.082
0.0435
0.085
0.045
0.059
ннооо LLMM
meV
41
29
32.9
33
24.6
34
27.6
20.3
meV
34
25
24.7
27
19.5
38
20.4
17.5
Ae
meV
63
48
42.5
51
32.6
61
45.7
36.1
(R²)
ag
0.073
0.106
0.765
0.448
0.331
0.961
0.066
0.033
ZPVE
meV
1.7
1.55
1.39
2.10
1.21
2.027
1.28
1.84
Un
meV
14
11
8.15
17
6.32
22
5.85
6.15
U
meV
19
12
8.34
16
6.28
21
5.83
6.38
H
meV
14
12
8.42
17
6,53
21
5.98
6.16
G
meV
14
12
9.4
19
7.56
20
7.35
7.62
C
cal
0.033
0.031
0.028
0.035
0.023
0.026
0.024
0.026
molK
Table 2: Results MD trajectories from the MD17 dataset. Scores given by the MAE of energy predictions (kcal/mol) and forces (kcal/mol/A). NequlP does not provide errors on energy, for PaiNN we include the results with lower force error of training only on forces versus on forces and energy. Benzene corresponds to the dataset originally released in Chmiela et al. (2017), which is sometimes left out from the literature. ET results are averaged over three random splits.
Molecule
SchNet
PhysNet
DimeNet
PaiNN
NequlP
ET
energy 0.37
0.230
0.204
0.167
0.123
Aspirin
forces
1.35
0.605
0,499
0.338
0.348
0.253
energy
0.08
0.078
0.058
Benzene
forces
0.31
0.187
0.187
0.196
Ethanol
energy
0.08
0.059
0.064
0.064
0.052
forces
0.39
0.160
0.230
0.224
0.208
0.109
Malondialdehyde
energy
0.13
0.094
0.104
0.091
0.077
forces
0.66
0.319
0.383
0.319
0.337
0.169
Naphthalene
energy
0.16
0.142
0.122
0.116
0.085
forces
0.58
0.310
0.215
0.077
0.097
0.061
Salicylic Acid
energy
0.20
0.126
0.134
0.116
0.093
forces
0.85
0.337
0.374
0.195
0.238
0.129
Toluene
energy
0.12
0.100
0.102
0.095
0.074
forces
0.57
0.191
0.216
0.094
0.101
0.067
Uracil
energy
0.14
0.108
0.115
0.106
0.095
forces
0.56
0.218
0.301
0.139
0.173
0.095
evaluate the architecture's performance in a limited data setting, the model is trained on only 1000 samples from which 50 used for validation. The remaining data is used for evaluation and is the basis for comparison with other work. Separate models are trained for each molecule using combined loss function for energies and forces where the energy loss is multiplied with factor of 0.2 and the force loss with 0.8. An overview of the results and comparison to the invariant models SchNet (Schütt et al. 2017b), PhysNet (Unke Meuwly 2019) and DimeNet (Klicpera et al.) 2020b), well as the equivariant architectures PaiNN (Schitt et al. 2021 and NequlP (Batzner ctt al. 202I) can be found in Table
To evaluate the architecture's capabilities a large collection of off-equilibrium conformations. we train and evaluate the equivariant Transformer on the ANI-1 dataset. It contains 22,057.374 configurations of 57,462 small organic molecules with up to 8 heavy atoms and atomic species H, C, N, and O. The off-equilibrium data points are generated via exhaustive normal mode sampling of the energy minimized molecules. The model is fitted DFT energies from 80% of the dataset,

---

Published as a conference paper at ICLR 2022
while 5% used as validation and the remaining 15% of the data make up the test set. Figure compares the equivariant Transformer's 's performance to previous methods DTNN (Schütt et al.
2017a), SchNet (Schütt al. 2017b), MGCN (Lu et al. 2019) and ANI (Smith et al. 2017a).
Comparison of testing MAE on the ANI-I dataset in eV. Results for DTNN, SchNet and MGCN are provided by Lu et al. 2019). The ANI method refers to the ANAKIN-ME (Smith et al.) 2017a) model used for constructing the ANI-1 dataset.
Figure 2:
0.10
0.113
0.108
MAE (eV)
0.08
0.078
0.06
0.057
0.04
0.02
0.012 ET
0,00
DTNN
SchNet
MGCN
ANI
3.1
ATTENTION WEIGHT ANALYSIS
Neural network predictions are notoriously difficult to interpret due to the complex nature of the learned transformations. To shed light into the black box predictor, we extract and analyze the equivariant Transformer's attention weights. We run inference the ANI-I, QM9, and MD17 7 test sets for all molecules and extract each sample's attention matrix from all attention heads in layers. Attention rollout (Abnar& & Zuidema 2020) under the single head assumption is applied during the extraction, resulting in a single attention matrix per sample. Figure visualizes these attention weights for random QM9 molecules (see Appendix F for further examples).
To analyze patterns in the interaction of different chemical elements, we average the attention weights each unique combination of interacting atom types (hydrogen, carbon, oxygen, ni- trogen, fluorine). This generates two attention scores for each pair of atom types, one from the perspective of atom type zı attending 22 and vice versa. The attention scores compared the probability of this bond occurring in the respective dataset, making sure the network's 's attention is simply proportional to the relative frequency of the interaction. Figure presentss summary of these bond probabilities and attention scores for QM9, ANI-1, and the average attention scores for all MD17 models. For further details, see Appendix H
Since the equivariant Transformer is trained to predict the energy of a certain molecular conforma- tion, we expect to pay attention to atoms that are displaced from the equilibrium conformation, i.e. the energy- minimmzed structure, of the molecule. We test this hypothesis by comparing the attention weights of displaced atoms to those the equilibrium conformation. Using the QM9 test set as source of equilibrium conformations, we displace single atoms by 0.4À in a random direction and compare the absolute magnitude of attention weights involving the displaced atom to the remaining attention weights. We find increased attention for displaced carbon and oxygen atoms in all models, however, only the model trained on ANI-1 attends more to displaced hydrogen atoms than to hydro- gen in its equilibrium position. Attention for displaced hydrogen atoms decreases in the model trained on QM9, which suggests that the energy labels in QM9 do not depend strongly on the exact location of hydrogen atoms. For detailed overview of the results, see Appendix C
is interesting to see that the training dataset influences attention. For static structures, like in QM9, attention analysis shows that very little importance is attributed hydrogens, while structural atoms like carbons are very important. For datasets which have dynamical data like ANI-1 and MD17, we see that hydrogen attended strongly. This is consistent with the fact that hydrogens are important for hydrogen bond-type interactions and therefore important for dynamics. This suggests that the network is not only learning meaningful chemical representation but also that training on dynamical datasets is important.
3.2 ABLATION STUDIES
We quantify the effectiveness of the neighbor embedding layer by the change in accuracy when ablating this architectural component. The neighbor embedding layer is replaced by regular atom type embedding, which evaluate by comparing the testing MAE on Up from QM9. The ablation causes a drop in accuracy of roughly 6% (from 6.24 to 6.60 V). We also try replacing the neighbor embedding by an additional update layer as the neighhor embedding resembles graph convolution

---

Published as a conference paper at ICLR 2022
ANI-1 C N
MD17
QM9
H
O
HCNOP
N O
H
Bond probabilities
2i N
F
H
Attention scores
c
C
2j N
N
F
H C
O F
HCNOF
H C ACNOE Zj
2j
2j
Figure 3: Depiction of bond probabilities and attention scores extracted from the ET model of TorchMD-NET using QM9 (total energy Uo), MD17 (average over 8 discussed molecules) and ANI- testing data. Attention scores are given 2; attending 2j. bond probabilities follow the same idea, showing the conditional probability of a bond between 2; and 2j. given 2, Darker colors correspond to larger values, element pairs without data grayed out. See Appendix G for overview of elemental composition in the respective datasets.
operation, which is the equivalent operation an update layer in graph convolutional networks such as SchNet. Here, the drop in accuracy is even more pronounced with decrease of around 10% (from 6.24 to 6.85 V). However, this may also the result overfitting as each update layer has about 4.6 times more parameters than neighbor embedding layer.
The hyperparameter set used for MD17 and A results in a model size of 1.34M trainable param- eters, which is significantly more than recent similar architectures such as PaiNN (600k) NequlP (290k). To rule out the possibility that the ET results are simply caused by an increased number of parameters, train smaller versions of the ET, which comparable to the size of PaiNN and NequIP. find that smaller versions the ET are still competitive and outperform previous state-of-the-ar results MD17. See Appendix D for details the results, hyperparameters and computational efficiency.
3.3 COMPUTATIONAL EFFICIENCY
To assess the computational efficiency of the equivariant Transformer, we measure the inference time of random QM9 batches comprising 50 molecules (including computing pairwise distances) on an NVIDIA V100 GPU (see Table We report times for different sizes the model, differing in the number of update layers, the feature dimension, and the size of the RBF distance expansion. ET-large uses the QM9 hyperparameter set, while ET-small is constructed using MDI7/ANI-I hy- perparameters (see Table 4p. The measurements were conducted using just-in-time (JIT) compiled models. The Jompllee versions ET-small and ET-large are 1.5 and 1.8 times respectively more efficient than the raw implementation. We only measure the duration of the forward pass, excluding the backward pass required to predict forces. The force prediction decreases inference speed by approximately 75.9%, resulting in 39.0mm per batch using ET-small and 48.5m using ET-large.
8

---

Published as a conference paper at ICLR 2022
[GDB 105807]
[GDB 11383]
[GDB 120698]
[GDB 5605]
[GDB 96582)
ANI-1
trained on
MD17
9M0
Figure 4: Visualization of five molecules from the QM9 dataset with attention scores corresponding to models trained on ANI-1, MD17 (uracil) and QM9. Blue and red lines represent negative and positive attention scores respectively.
Table 3: Comparison of computational efficiency between PaiNN, DimeNet++ and different sizes of TorchMD-NET ET. The time is measured at inference using random batches of 50 molecules from QM9. Speed of ET models ToT reported mean + standard deviation over 1000 calls. Values for PaiNN and DimeNet++ are taken from Schütt et al. (2021) so differences in efficiency may to some degree originate from different implementations.
PaiNN
DimeNet++
ET-small
ET-large
time per batch ms parameters
45 ms
9.4ms 3.4ms 11.7ms 4.0ms
600k
1.8M
1.34M
6.87M
DISCUSSION
In this work, we introduce a novel attention-based architecture for the prediction of quantum me- chanical properties, leveraging the use rotationally cquivariant features. We show high degree of accuracy the QM9 benchmark dataset, however, the architecture's effectiveness is particu- larly clear when looking at the prediction of energies and atomic forces in the context of molec- ular dynamics. We set molecule Benzene) and demonstrate the architecture's ability to work in low data regime. As de- scribed in previous work (Schütt al. 2021 Fuchs al.) 2020), the model's vector features and equivariance can be utilized in the prediction of variables beyond scalars. Here, only the dipole moment is predicted in this fashion, however, the architecture capable of predicting tensorial properties. By extracting and analyzing the model's attention weights, we gain insights into the molecular representation, which characterized by the nature the corresponding training data. We show that the model does not pay much attention to the location of hydrogen when trained only on energy-minit minimized molecules, while a model trained on data including off-equilibruum conforma- tions focuses to a large degree hydrogen. Furthermore, analyzing attention weights involving atoms displaced from their equilibrium location. We demon- strate that offieuuiiir conformations in the training data play important role in the accurate prediction of molecule's 's energy. This highlights the importance of configurational diversity in the evaluation of neural network potentials.
state-of-the-are
all MD17 targets (except force prediction of the
validate the learned representation by

---

Published as a conference paper at ICLR 2022
After the final review of this paper, NequlP's preprint (Batzner et al. (021)) updated the results with better accuracy for MD17. However, it requires using high order spherical harmonics which are likely substantially slower than Transformer models.
SOFTWARE AND DATA
The equivariant Transformer is implemented in PyTorch (Paszke et al. 2019). using PyTorch Geo-
metric (Fey & Lenssen| 2019) as the underlying framework for geometric deep learning. Training is done using pytorch-lightning (Falcon & The Py PyTorch I. Lightning team 2019). high-level interface for training Py PyTorch models. The datasets QM MD17 and ANI are publicly available and all source code for training, running and analyzing the models presented in this work is available at github.com/orrbmm//orcmntt
ACKNOWLEDGMENTS
GDF thanks the project HID2020-116564G11 funded by MCIN/ AEI /10. 13.33303//00 10001 1033, Unidad de Excelencia María de Maeztu" funded by the MCIN and the AEI (DOI: 00.13039/51110001100 Ref: LEX2018-000792--0 and the European Union's Horizon 2020 re- search and innovation programme under grant agreement No. 823712. Research reported in this publication was supported by the National Institute of General Medical Sciences (NIGMS) of the National Institutes of Health under award number GM140090.
REFERENCES
Samira Abnar and Willem Zuidema.
Quantifying Attention Flow in Transformers. arXiv:2005. 00928 [cs], May 2020. URL http:///axii..orraa/////0000000 00928 arXiv: 2005.00928.
Brandon Anderson. Truong-Son Hy, and Risi Kondor. Cormorant: Covariant Molecular Neural
Networks. arXiv: 19Xx:1996600015 /physics, stat], November 2019. URL http://arxiv.
abs/ 04015
arXiv: 1906.04015.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization. arXiv:1607.0643: 06450 arXiv:
[cs, stat], July 2016. TRLp:htp://aaxxt www.:o/gg/aaaaaaaeooo
Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Ko- mbllth, Nicola Molinari, Tess Smidt, and Boris Kozinsky. SEqiiarquivariant Graph Neu- ral Networks for Data-Efficient and Accurate Interatomic Potentials. arXiv:2101.033 164 [cond- mat, physics:physics). July 2021. URL http:://axxiv www.:/xwvvvoggobb//b/aaae abtp///./.111.000 03164 2101.03164.
arXiv:
Jörg Behler and Michele Parrinello. Generalized Neural-Network Representation of High- Dimensional Potential-Energy Surfaces. Physical Review Letters, 98(1))144401, April 2007. doi: 10. 10.1 146401. URL https: httt:::/A..kkkppp...oo///////// org. 1103/
aps
PhysRevLett. 98 146401 Publisher: American Physical Society.
Stefan Chmiela, Alexandre Tkatchenko, Huziel F. Sauceda, Igor Poltavsky, Kristof T. Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Sci- ence A https::/advaacces http://wwww.s.....cccccc..c sciencemag. org/ wontent/t////.n 11030155 Publisher: Amer-
May 2017. ISSN 2375-2548. doi: 10. 1e66/ccidvv.60330000 URL
ican Association for the Advancement of Science Section: Research Article.
Stefan Chmiela, Huziel E. Sauceda, Klaus-Robert Müller, and Alexandre Tkatchenko. Towards exact molecular dynamics simulations with machine-learned force fields. Nature Communi- cations, 91):3887, September 2018. ISSN 2041-1723 doi: 10. 06169-2
URL hattps://wwet wttp:aattuee..o/aat..oooooo ottp:///.....e.e....///////1100. 0118-00166661111-6111-------
Number:
Publisher: Nature Publishing Group
hutps://do..orr/10. htt.://..oorr/100000/////// 6084/09 Jieshare...978904.55
http://iishhaa.ccoommmmm///////s Dataset Rclease/100470411
Dataset
10

---

Published as a conference paper at ICLR 2022
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking Attention with Performers. arXiv:2009. 0e0::2009.44777 14794 [cs, stat), 1 14794
March 2021. URL http::://ppp....rrtrttp//
arXiv: 2009.147944
Stefan Doemr Maciej Majewski, Adrià Pérez, Andreas Krämer, Cecilia Clementi, Frank Noe, Toni Giorgino, and Gianni De Fabritüi.. TorchMD: A Deep Learning Framework for Molecular Simu- lations. Journal of Chemical Theory and Computation, 17(4):2355-2363, April 2021. ISSN 1549- 1012//acs.jct..c01111
9618. doi:
URL https: http::////oo.../1.1100000 10t jote. Publisher: American Chemical Society.
0c01343
William Falcon and The Py PyToch Lightning team. PyTorch Lightning. March 2019. URL https:
Matthias Fey and Jan Eric Lenssen. Fast Graph Representation Learning with PyTorch Geometric, May 2019. URL http:p//gggaaaa.cccoc/////////pp
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. SE(3)-Transfrmmen 3D Roto- Translation Equivariant Attention Networks. arXiv:2006. 10503 [cs, stat], November 2020.
URL httt:ht/tt/xii// www.aiii...g.aaaeaoooooa wwtp://g//...oo.o//////
10503
arXiv: 10503.
Jan Hermann, Zeno Schätzle, and Frank Noé. Deep neural network solution of the electronic Schrödinger equation. Nature Chemistry, 12(10))891-897, October 2020. ISSN 1755-4330, 1755-4349. doi: 10. 00.1038//415570020-05555
URL
org. abs/ abs/1909.. 1909.
08423 arXiv: 1909.088 .0943..
Brooke E. Husic, Nicholas E. Charron, Dominik Lemm, Jiang Wang, Adrià Pérez, Maciej Ma- jewski, Andreas Krämer, Yaoyi Chen, Simon Olsson, Gianni de Fabritiis, Frank Noé, and Ce- cilia Clementi. Coarse Graining Molecular Dynamics with Graph Neural Networks. The Jour- Novembee URL http:://axxi...oo//// 2000 1122 arXiv: 2007.11412.
nal of Chemical Physics, 153(19):144101, November 2020. ISSN 0021-9606, 1089-7690. doi:
Michael Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyun- jik Kim. Lierransformer: Equivariant self-attention for Lie Groups. December 2020. URL
https: /tas:///rxiiz org.oogg////... 2212//202222 10885......18
Diederik Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. arXiv:1412.6980 [cs], January 2017. URL /ttttp://rxxi...o org/ abs/ 1412. 6980 arXiv: 1412.6980.
Johannes Klicpera, Shankari Giri, Johannes Margraf, and Stephan Günnemann. Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules. arXiv:2011. 1.1111 [physics], December 2020a. URL http://aaxxi..orr/bbb///000
14115 arXiv: 2011.14115
Johannes Klicpera, Janek GroB, and Stephan Günnemann. Directional Message Passing for Molec- ular Graphs. arXiv:2003. 03133 [physics. stat), March 2020b. URL arXiv: 2003.03123.
0bs/2203.0330
Ziteng Liu, Liqiang Lin, Qingqing Jia, Zheng Cheng, Yanyan Jiang, Yanwen Guo, and Jing Ma. Transferable Multi-level Attention Neural Network for Accurate Predic- tion of Quantum Chemistry Properties via Multi-task Learning. ChemRriv, July 2020. URL https: http///ii..rr//aattt....... preprint/Transferrttt www.bbleeeelttieeeeee Att._tteettttt inn Neural Network._ For wwwccuratee_redicttttee
doi: 10.26434/cheemrxiv... 12588110..1
http:paottccciooouuouuuuuuuuu of
Chemlsstry_Propertieee via
Publisher: ChemRxiv.
Chengqiang Lu, Qi Liu, Chao Wang, Zhenya Huang, Peize Lin, and Lixin He. Molecular Prop- erty Prediction: A Multilevel Quantum Interactions Modeling Perspective. Proceedings of the Artificial Intelligence, 1050-100022-0000 July 2019. ISSN 2374-3468. doi: 10.1609/aaii 333101.330111552 URL httpt:t//////t.. //tj/ aaai. org/ htt///ndxx...m php/ AAAI/ Number: 01.
AAAI Conference
article/ve ew 3896

---

Published as a conference paper at ICLR 2022
Minh-Thang Luong, Hieu Pham, and Christopher D. Manning. Effective Approaches to Attention- based Neural Machine Translation. 1508v [es], September 2015. URL http://
arxiv.org/ abs 1508. 04025
arXiv: 1508.04025.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach 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. (IXiv:1912.01703 Jes, stat], December 2019. URL. http: org/ attp://..///11/....././ 01703
arXiv: 1912.01703.
David Pfau, James S. Spencer, Alexander G. D. G. Matthews, and W. M. C. Foulkes, Ab initio solution of the many-electron Schrödinger equation with deep neural networks. Physical Re- view Research, 2(3):033429, September 2020. doi: 10.11000hyyyReesseccc2...333.33
URL 0333429 Publisher:
https:://ink: aps http:///c.///..111/////a/eeeeeee
American Physical Society.
Zhuoran Qiao, Matthew Welborn, Animashree Anandkumar, Frederick R. Manby, and Thomas F. Miller III. OrbNet: Deep Learning for Quantum Chemistry Using Symmetry-Adaple Atomic- Orbital Features. 0rXivv2007.000026 [physics], September 2020. doi: 10. 10.106355.00219555 URL
http:://arxiv.or//abs...
08026
arXiv: 2007.08026
Raghunathan Ramakrishnan, Pavlo O. Dral, Matthias Rupp, and O. Anatole von Lilienfeld. Quan- tum chemistry structures and properties of 134 kilo molecules. Scientific Data, 1(1):140022,
August 2014. ISSN 2052-4463. doi: 10. (t98//sdata2001422 htthttts::/www //ww::::/www..../. nature.
Com/ coom/articlesssdie wtt.:/leessstataa111aa
Number: Publisher: Nature Publishing Group.
Matthias Rupp, Alexandre Tkatchenko, Klaus-Robert Müller, and O. Anatole Lilienfeld. Fast
and Accurate Modeling of Molecular Atomization Energies with Machine Learning. Phys-
ical Review Letters, 108(5)008330, January 2012. doi: 10.1103/Phys 108.058800.00 URL http://wwww...pp.ooooooooo////// www3:/1yse///eeettt.oomommm 108 058301 lisher: American Physical Society.
aps wwt.s//////////////////////
Pub-
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) Equivariant Graph Neural Net- works. arXiv:2102.0988 [cs, stat], May 2021. URL http:://axxv..rr//44 09844
arXiv: 2102.09844.
K. T. Schütt, H. E. Sauceda, P.-. Kindermans, A. Tkatchenko, and K.-R. Müller. SchNet A deep learning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):24122,,, March 2018. ISSN 0021-9606. doi: 111063/1.5019771 URL https: 7/aip. www.stiiiiiii.oo org./onooorroooooooooo wwwp///a.a/..ooooooo 1tt.b//.....11111...1....0. 5019779 tute of Physics.
Publisher: American Insti-
Kristof Schütt, Farhad Arbabzadah, Stefan Chmiela, Klaus R. Müller, and Alexandre Tkatchenko. Quantum-chemiin insights from deep tensor neural networks. Nature Communications, 8(1): 13890, January 2017a. ISSN 2041-1723. doi: 10. 10.1038//oomms/ 3890. URL /ttpps://www
nature. com/ arti /tco///commss 3890
Publisher: Nature Publishing Group.
Kristof T. Schütt, Pieter-Jan Kindermans, Huziel E. Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. SchNet: continuous-filter convolutional neural net- work for modeling quantum interactions. arXiv: 1706. 776666866 [physics, stat], December 2017b. URL http::/taaxiiii wwt.
08566
arXiv: 1706.08566.
Kristof T. Schütt, Oliver T. Unke, and Michael Gastegger. Equivariant message passing for the [physics), June 2021.
prediction of tensorial properties and molecular spectra. URL httthttpp//axxiv org/ abs/2102. 03150
arXiv: 2102.03150.
J. S. Smith, O. Isayev, and A. E. Roitberg. ANI-1: extensible neural network potential with DFT accuracy at force field computational cost. Chemical Science, 8(4):319233233 March 2017a. ISSN 2041-6539. doi: 10.1039/C6.005220A. URL https:/ /too:://pubss rsc org/
len/ http:///een/rtcieaaaaaaaaa////// of Chemistry.
Publish The Royal Society
Publisher:
12

---

Published as a conference paper at ICLR 2022
Justin S. Smith, Olexandr Isayev, and Adrian E. Roitberg. ANI-I, A data set of 20 million calculated off-equilibrium conformations for organic molecules. Scientific Data. 4(1):170193, December 2017b. ISSN 2052-4463. doi: 10. 10.1038/sdata.2017..
(Rthhtt::/// /www. nature. com/ Publisher: Nature Publishing Group-
Seticles/sdataa2111
Attila Szabo and Neil S. Ostlund. Modern Quantum Chemistry:
Introduction to Advanced Electronic Structure Theory. Dover Books on Chemistry. Dover Publications, July 1996. ISBN 978-0-486- 69186-2. Google-Books-ID: 6mV9gYzEkgIC.
Oliver T. Unke and Markus Meuwly. PhysNet: A Neural Network for Predicting Energies, Forces, Dipole Moments and Partial Charges. arXiv: 1902. 08408 [physics], March 2019. doi: 10.1021/ acs.jetc.bb00181.
URLhttp. /ttp/aaaavvv//.a///ii//aa./a..ia org. abs 902. 08408 arXiv: 1902.08408.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. arXiv: 03766 [cs], De- cember 2017. URL http:::/axxi...rggg/ppp/////
1066. 03762
arXiv: 1706.03762.
Jiang Wang, Simon Olsson, Christoph Wehmeyer, Adrià Pérez, Nicholas E. Charron, Gianni de Fab- ritiis, Frank Noé, and Cecilia Clementi. Machine Learning of Coarse- Coarse-Grained Molecular Dynam- ics Force Fields. ACS Central Science, 5(5):755-767, May 2019. ISSN 2374-7943. doi: (221/acscentcci.8000911
httphtttp///......
10t.:/000000////cttteecooo 8b00913
Publisher: American Chemical Society.
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3D Steerable CNNs: Learning Rotationally Equivariant Features in Volumetric Data. arXiv:I807.000 [cs, stat], October 8.. URL htttp::/aaxxv..or//abbbba
02547
arXiv: 1807 807.00554.
13

---

Published as a conference paper at ICLR 2022
A ARCHITECTURAL DETAILS AND HYPERPARAMETERS
All models in this work were trained using distributed training across two NVIDIA RTX 2080 Ti GPUs, using the DDP training protocol. This leads to training times of 16h for QM9, 10h for MD17 and 83h for ANI-I.
Table 4: Comparison of various hyperparameters used for QM9, MD17 and ANI-1.
Parameter
QM9
MD17
ANI-1 1
initial learning rate
4
10
1
10-3
7 10
Ir patience (epochs)
15
30
5
Ir decay factor
0.8
0.8
0.5
Ir warmup steps batch size
10,000
1,000
10,000
128
2048
no. layers
8
6
no. RBFs
64
32
feature dimension no. parameters
256
128
128
6.87M
1,34M
1.34M
While the ET model follows a similar idea as the (B(3)-Traasformer introduced by Fuchs et al. 2020), there are significant architectural differences. The (r(n))Trannfrmmer relies heavily on pensive features, such as Clebsch-Gordan coefficients and spherical harmonics while the ET model only requires interatomic distances. Additionally, split scalar and equivariant features into two pathways, which exchange information inside the update layer while the TraassTranso com- putes message passing updates for each type of feature vector (scalar equívariant). Finally, modified attention mechanism and update step differ significantly from the SE(3)-Transformer's message passing layer, which, for example, does not handle self-interactions in the attention mech- anism and applies only linear transformations to distance features.
B IMPORTANCE OF HYDROGEN
While the ET model trained on QM9 mostly attends to carbon atoms, models trained on molecular dynamics trajectories show strong focus carbon- -hybroen interactions. This highlights the different nature of datasets containing only energy minimized conformations in contrast to datasets containing MD data. We further support this hypothesis by comparing the reduction in accuracy for models, which are trained without hydrogen, on the datasets QM9 (total energy Un) and MD17 (aspirin). We show that the loss in accuracy when predicting the energy using the model trained on MD17 is one order of magnitude larger than when training only on molecules in the ground state. Furthermore, the accuracy of force predictions drops by another 1.5x compared to MD17 energy predictions when excluding hydrogen. A summary of the results is given in Table
Table 5: Test MAE of the TorchMD-NET ET on QM9 and MD17 (aspirin), trained with and without hydrogen.
Dataset QM9
with hydrogen
without hydrogen
relative change
6.37 meV
20.83 meV
227.0%
MD17
energy
5.33 meV
137.59 meV
2481.4%
forces
10.67 meV/Á
433.99 me
3966.5%
C ATOM DISPLACEMENT
Figure shows averaged absolute attention scores for molecules where a single atom has been displaced from the equilibrium structure. This reiterates the idea that the model trained only equilibrium structures (QM9) focuses on carbon and neglects hydrogen. The models trained on off-equilibrium conformations show higher degree of attention for displaced hydrogen atoms. We
14

---

Published as a conference paper at ICLR 2022
restrict the molecules to only contain hydrogen, carbon, and oxygen in order to compare results between models trained on QM9, MD17 (aspirin), and ANI-1.
H
C
O
QM9
0.2
model trained on
0.0
MD17
0.2
0.0
ANI1
0.2
0.0
equilibrium
displaced
equilibrium
displaced
equilibrium
displaced
Figure 5: Averaged attention weights extracted from the ET on the QM9 test set (molecules sisting of H, C, and O only) with a displacement 0.4À in single atoms. Blue bars show attention towards atoms in equilibrium locations, orange bars correspond to attention weights involving the displaced atom. Attention scores are normalized inside each molecule. The black bars show the attention weights' standard deviation.
D EQUIVARIANT TRANSFORMERS WITH REDUCED PARAMETER COUNT
compare the ET model with reduced number of parameters, matching those of PaiNN (600k) and NequlP (290k), state-of-the-art models on the MD17 benchmark. This ensures that our results are not solely consequence of larger model size but correspond to an improved architecture. The smaller ET models still competitive and outperform state-of-the the-art most MD17 targets. T provides an overview of the adjusted architectural hyperparameters and Table 7 summarizes the results MD17.
Table 6: Hyperparameter of the full ET model compared PaiNN- and NequiP-sized variants. This table only includes hyperparameters that were changed.
Hyperparameter full ET
PaiNN-sized ET
NequIP-sizer sized ET
no. layers
no. RBFs
32
16
feature dimension
128
120
80
no. parameters
1.34M
593k
273k
15

---

Published as a conference paper at ICLR 2022
Table Energy (kcal/mol) and force (kcal/mol/A) MAE of the PaiNN- and NequlP-sized ET mod- els. The "full ET" column is equal the results in Table and is meant for comparison with the smaller ET variants. Values in bold indicate the best result out of the two models in direct compari- son.
PaiNN-sized
NequlP-sized
Molecule
full ET
ET (594k)
PaiNN (600k)
ET (273k)
NequlP (290k)
Aspirin
energy
0.124
0.138
0.167
0.143
forces
0.255
0.334
0.338
0.337
0.348
energy
0.056
0.057
0.063
Benzene
forces
0.201
0.197
0.189
0.187
Ethanol
energy
0.054
0.053
0.064
0.053
forces
0.116
0.112
0.224
0.123
0.208
Malondialdehyde
energy
0.079
0.080
0.091
0.080
forces
0.176
0.209
0.319
0.218
0.337
Naphthalene
energy
0.085
0.084
0.116
0.085
forces
0.060
0.080
0.077
0.080
0.097
Salicylic Acid
energy
0.094
0.095
0.116
0.097
forces
0.135
0.175
0.195
0.184
0.238
Toluene
energy
0,074
0.075
0.095
0.075
forces
0.066
0.088
0.094
0.091
0.101
energy 0.096
0.094
0.106
0.096
Uracil
forces
0.094
0.122
0.139
0.128
0.173
E ABLATION OF EQUIVARIANT FEATURES
We perform ablation of the TorchMD-NET ET's equivariance and compare the performance of the resulting rotationally invariant model to that of the equivariant Transformer. Without equivari- ance, the MAE in total energy Uo in QM9 increases from 6.24 6.64 meV (6%). On aspirin inside the MD17 benchmark, removing the equivariance causes the energy MAE to rise from 5.37 to 13.23 meV (146%), while force errors go up from 11.05 to 30.27 V/À (174%). As, without equivari- ance, the error increases much more drastically on dynamical data, we hypothesize that equivariant features are particularly useful when dealing with non-zero forces.
F MOLECULAR REPRESENTATION BY DATASET
Figures 6a 6bb and 60 show three dimensional visualization of three random molecules from the datasets QM9, MDT7 (aspirin) and ANI-1 respectively. We extract the attention weights from the best performing equivariant Transformer on the test sets of the three datasets respectively to make sure that no model has seen visualized conformation during training. The red and blue lines be- tween atoms depict the 10 largest absolute attention weights where the line width and alpha value represent the absolute attention weight. Red lines show positive attention weights, blue lines cor- respond to negative attention weights, which occur due to using SiLU activations inside the atten- tion mechanism. While the model trained on QM9 focuses largely carbon-carbon interactions, is clear that models trained on MD17 and ANI-1 have a strong focus on hydrogen-card hydrogen-oxye interactions. This corresponds to the results found in the attention weights aver- aged over pairwise interactions between elements. The attention weights are normalized inside each molecule.
and
16

---

Published as a conference paper at ICLR 2022
Molecules from QM9 molecule
molecule
molecule
ANI-1
model trained on MD17 (uracil)
QM9
(a) Random molecules from the QM9 test set.
Molecules from MD17 (aspirin)
molecule 1
molecule 2
molecule
ANI-1
model trained on MD17 (uracil)
QM9
Random conformations from the MD17 test set of aspirin trajectories.
17

---

Published as a conference paper at ICLR 2022
Molecules from ANI-1
molecule
molecule
molecule
ANI-1
model trained on
MD17 (uracil)
ONB
(c) Random molecules from the ANI-1 test set.
Figure 6: Visualization of 10 largest attention weights by absolute value on random molecules from QM9 (a), MDI7-aspirin (b) and ANI-1 (c). Each column shows the same molecule, rows correspond the same ET model trained on QM9, MDI7-uracil and ANI-1 respectively.
G FREQUENCY OF ELEMENTS
Figure shows the distribution of elements in the datasets QM9, MD17 and ANI-1 where MD17 corresponds to the combination of all target molecules. It aims to assist with the interpretation of attention scores where certain pairs of elements are largely underrepresented in the model's attention. This corresponds predominantly to nitrogen and fluorine as MD17 contains only single molecule with nitrogen (uracil) and fluorine is only found in QM9.
ANI-1
MD17
QM9
Atom occurance (%))
40
20
0
C
N
O
F
H
C
N
O
F H
OF
Figure 7: Distribution clements in the datasets QM9, MD17 (combination of all target molecules)
and ANI-1.
18

---

Published as a conference paper at ICLR 2022
H MD17 ATTENTION WEIGHTS
Figures and 8b contain the extracted rolled out attention weights for the MD17 target molecules aspirin, benzene, ethanol, malondialdehyde, naphthalene, salicylic acid, toluene and uracil. While the molecules show different attention patterns, a high degree attention for hydrogen atoms is common for most of the molecules. However, models trained aspirin, malondialdehyde and uracil additionally exhibit a strong focus on oxygen.
The sum each equals one, meaning that the probabilities represent the conditional proba- bility
Nbouded (Zms kn)
P2. anoondmmzilll 21) Σnee EZ Nbouded (Zmi 2k)
(14)
where Noonded (Z₁, 2z) is the total number of bonds between atom types and 2j found in the collection of molccules.
MD17 Aspirin
MD17 Benzene
MD17 Ethanol
MD17 Malondialdehyde
Bond probabilities
H
x
N
N
H
Attention scores
2
N
N
0
2j
Attention scores and bond probabilitics from aspirin, benzene, ethanol and malondialdehyde.
19

---

Published as a conference paper at ICLR 2022
MD17 Naphthalene
MD17 Salicyllic Acid
MD17 Toluene
MD17 Uracil
H
Bond probabilities
N
Attention soores
2
N
2
2
(b) Attention scores and bond probabilities from naphthalene, salicylic acid, toluene and uracil.
Figure 8: Bond probabilities and attention scores extracted from the ET using testing data from all individual molecules in the MD17 dataset. Attention scores are given as 2/ attending to 2j. bond probabilities follow the same principle, showing the conditional probability of bond between 2; and Zj. given The rightmost subfigure displays the total number of atoms for each type in the data.
PERFORMANCE ON REVISED MD17 TRAJECTORIES
On top of the MD17 dataset, trajectories with higher numerical accuracy were published for some MD17 molecules. This includes aspirin at CCSD) and benzene, malondialdehyde, toluene and ethanol at CCSD(T) level of accuracy (Chmiela et al., 2018). Here, we present ET results on these trajectories with the same training protocol used for the original MD17 dataset (950 training samples, 50 validation samples).
Table ET results on CCSD/CCSD(T) trajectories from Chmiela et al. 2018 Scores are given by the MAE of energy predictions (kcal/mol) and forces (kcal/mol/A). Results averaged over two random splits.
Aspirin
Benzene
Ethanol
Malondialdehydee
Toluene
energy
0.068
0.002
0.016
0.024
0.011
forces
0.268
0.008
0.103
0.168
0.062
20


Reviews:
[{"review_id": "UPawlAhAAYq", "reviewer": "Reviewer_RnuV", "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": "# Strengths\n\n1. The idea to put attention at the very center of an ML force field is novel.\n2.  The performance of the method is state-of-the-art.\n\n# Weaknesses\n\n1. The expose of the attention mechanism and update layer lacks detail.\n2. The value of inspecting attention weights is not clear to me.", "summary_of_the_paper": "The authors introduce a novel architecture for ML force fields, the Equivariant transformer (ET). It is based on the Transformer approach and can be used to predict energies (and forces) and other molecular properties (e.g., QM targets). The performance on standard benchmarks such as QM9 and MD17 is impressive. The authors inspect the attention weights.", "summary_of_the_review": "1. The introduction is well-written and relates the work to previous works.\n2. Sections 2.1 to 2.4 are not detailed enough. Considering that these Sections is at the very center of this work, they deserve more details:\n   1) is the function $a_n$ defined anywhere? \n   2) where are the matrices (?) Q and K coming from?\n   3) how is this attention matrix A computed?\n   4) how did the authors arrive at the update rules in equations (8) and below?\n3. Specification of the training details is welcome.\n4. Did you consider using the updated version of MD17, rMD17, for your experiments? Why did you decide against using the updated dataset?\n5. There is a typo in Table 3 (malonaldehyde).\n6. In the abstract, the authors claim to have gained \"valuable insights into the black box predictor\". In the manuscript, the value of their insights is not properly laid out. What did the authors learn about the model that would help them improve it, for instance?\n\n# Update\n\nThe authors updated the manuscript and addressed my questions. For this reason, I am willing to increase my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "TjdbneJCYa4", "reviewer": "Reviewer_AQ9R", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is well structured and the proposed method is described clearly. The use of the modified attention mechanism could be better motivated. Often, attention is used to capture non-local interactions for which convolutions are not suitable (e.g. Unke et al 2021). Here, the attention is forced to be local by a cutoff function and more reminiscent of factorized tensor layers.\n\nET improves on the QM9 dataset for 5 of 12 chemical properties. It does not improve however on the energies-related properties (U0, U, G, H), while improvng on energies on the MD17 dataset. This might indicate overfitting. ET uses more than 10x the number of parameters as NequIP or PaiNN. Therefore, it does not become clear whether the improved performance stems from the ET architecture or the size of the model. This should be demonstrated in the ablation studies with reduced number of layers, feature dimensions and/or RBFs.\n\nThe analysis of attention scores is interesting, however, it does not become clear from the paper how to interpret the results, e.g. in Figs 3 and 4. In the current form, the shown results are not too useful to \"shed light into the black box predictor\". The analysis of attention to displaced hydrogen is aligned with what one would expect from equilibrium vs out-of-equilibrium data, but could probably also have been achieved by a sensitivity analysis. The paper could be improved by extracting patterns from the attention scores that, ideally, go beyond pairwise interactions.\n\nThe demonstrated computational efficiency is impressive, in particular given the model size. Is there an explanation why the computation time is comparable to the much smaller model except implementation issues? Do the benchmarks include the computation of interatomic distances with scalable neighbor lists, or is the implementation considering all pairwise interactions? A fast evaluation is particularly important for molecular dynamics simulations. Therefore, having a comparison for MD17 molecules with PaiNN and and NequIP would be very relevant.\n\nFurther comments:\n- p. 1: Pfau et al and Hermann et al have proposed neural network approaches to VMC, not coupled cluster\n- p. 1: the original publication of Behler-Parinello should be cited (PRL, 2007) when referring to first work in this field\n- p. 2: when using gated equivariant blocks as in PaiNN, Weiler et al 2018 should also be cited who introduced equivariant gated non-linearities", "summary_of_the_paper": "The paper describes an equivariant neural network (ET) with attention mechanisms that is applied to the prediction of molecular properties. The ET performs competitive with previous approaches on commonly used benchmark data.", "summary_of_the_review": "The presented approach shows iterative improvements over previous work. Unfortunately, at the current stage it is not clear whether those are due to the ET architecture or the increased model size. The analysis of attention scores is promising, but still lacks the necessary interpretation.\n\nUpdate:\nBased on the response of the authors, I have raised my score.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Dwa3Jbqmshg", "reviewer": "Reviewer_ZUVc", "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": "Strengths: The experiment results of this paper looks very promising. In QM9 dataset, the proposed model outperform other models in many targets. Especially for dipole moment and electronic spatial extent, the proposed model is much better. For MD17 dataset, the proposed model outperform all the baselines on all the targets. For ANI-1 dataset, the result is also much better. \n\nWeakness: It will be better to include more discussion about the components of the model, especially the components not in regular transformers. Readers might be curious about \n(1) What are the motivations of these components. \n(2) Will the performance drop if we replace the components by others? Which components are the most important ones?\n\nAlso, there are some other famous baselines such as EGNN [1], DeepPotential [2]. It will be better to include the results in the table.\n\n[1] Satorras, Victor Garcia, Emiel Hoogeboom, and Max Welling. \"E (n) equivariant graph neural networks.\" arXiv preprint arXiv:2102.09844 (2021).\n[2] End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems\nL Zhang, J Han, H Wang, WA Saidi, R Car - arXiv preprint arXiv:1805.09003, 2018", "summary_of_the_paper": "This paper proposed equivariant transformers --- a neural network based algorithm to predict properties of molecules. The architecture is built upon the traditional transformer architecture, combined with some modifications specific to molecular property prediction tasks, such as exponential normal radial basis functions, SiLU activation function and design of update layers. It is shown that the proposed model have good performance on QM9, MD17 and ANI-1 dataset. Ablation studies analyze the attention weight and give some insight about how the model works.", "summary_of_the_review": "Overall, this paper is a good paper with strong experiment results on various datasets. The components of the model look very interesting. Possible improvement include more discussion on model components and comparison with more baseline algorithms.", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "18PH99BmXaW", "reviewer": "Reviewer_9RJX", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "Strengths\n1. The presented transformer model is novel and is applicable to atomic graphs because it is equivariant to rotations and includes edge features in the attention computation. These edge features include an envelope term that prevents discontinuities in the loss landscape. This model is also well motivated by the underlying physics of the problem.\n2. The model obtains excellent results on a variety of datasets showing its generality.\n3. The paper includes detailed analysis of the learned attention weights.\n4. The presented model has good runtime performance compared to models like Dimenet++ that require higher order interactions, while achieving good performance. In practical applications like  catalyst discovery, these models are often used to search through millions of examples making efficient inference important.\n\nWeakness\n1. Other equivariant transformers have recently been proposed in the literature (for e.g. SE(3)-Transformers: https://arxiv.org/abs/2006.10503). Since these models are comparable to the method proposed in this paper, it would be good to compare these models. Also, clarifying the differences between these models would make the contributions of this paper clearer.\n2. The analysis of attention seems difficult to follow for somebody who is unfamiliar with chemistry terminology. Since ICLR is a general DL conference, I would urge the authors to edit the text to make it easier to follow.\n", "summary_of_the_paper": "The paper presents an equivariant transformer model for predicting quantum mechanical properties from an atomic graph. The model obtains SOTA or near-SOTA results on three popular datasets while maintaining good computational efficiency. The primary novelty in their method is a new way to compute the attention score using edge features. The paper also presents a detailed analysis of the attention weights that give insights into what the model is attending over. This is interesting from a chemistry perspective.\n", "summary_of_the_review": "The paper presents a new rotation equivariant transformer model for atomic prediction. The resulting method obtains good results on 3 popular datasets, while being computationally efficient at inference time. The paper also provides a nice analysis of the attention weights that shed light on the inner workings of their model. While I would recommend the authors to clarify this analysis for the final version, I am leaning towards accepting the paper in its current form.\n", "strength_and_weaknesses": "", "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": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a rotationally equivariant transformer architecture for predicting molecular properties. The proposed architecture demonstrates good computational efficiency and good results on three benchmarks.

All four reviewers recommend acceptance (two weak, two strong), citing the novelty of the architecture, the good computational efficiency of the model and the good empirical results as the main strengths of the paper. The reviewers expressed minor criticisms and recommendations for improvement, some of which were addressed by the authors during the reviewing process, which led to an increase in scores.

Overall, this is a nice contribution of machine learning to science, and I'm happy to recommend acceptance to ICLR.
Decision: Accept (Spotlight)