Equivariant Transformers for Neural Network based
Molecular Potentials
Philipp Thölke∗Computational Science Laboratory, Pompeu Fabra University,
PRBB, C/ Doctor Aiguader 88, 08003 Barcelona, Spain
philipp.thoelke@gmx.de
Gianni de Fabritiis
Computational Science Laboratory, Pompeu Fabra University,
PRBB, C/ Doctor Aiguader 88, 08003 Barcelona, Spain andInstitució Catalana de Recerca i Estudis Avançats (ICREA),
Passeig Lluis Companys 23, 08010 Barcelona, Spaing.defabritiis@gmail.com

Abstract
The prediction of quantum mechanical properties is historically plagued by a tradeoff 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 thiswork we propose a novel equivariant Transformer architecture, outperforming stateof-the-art on MD17 and ANI-1. Through an extensive attention weight analysis,
we gain valuable insights into the black box predictor and show differences in thelearned 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.

1

Introduction

Quantum mechanics are essential for the computational analysis and design of molecules and materials. However, the complete solution of the Schrödinger equation is analytically and computationallynot practical, which initiated the study of approximations in the past decades [Szabo and Ostlund,1996]. A common quantum mechanics approximation method is to model atomic systems according
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 moreaccurate techniques like coupled-cluster exist but both still lack the computational efficiency to be
applied on a larger scale, although recent advances are promising in the case of coupled-cluster [Pfauet 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 contributions
are rooted in the Behler-Parrinello (BP) representation [Behler, 2011] and the seminal work from
Rupp et al. [2012]. One of the best transferable machine learning potentials for biomolecules, ANI[Smith et al., 2017a], is based on BP. A second class of methods, mainly developed in the field ofmaterials science and quantum chemistry, uses more modern graph convolutions [Schütt et al., 2018,
Unke and Meuwly, 2019, Qiao et al., 2020, Schütt et al., 2021]. Recently, other work has shown that a∗

Secondary affiliation: Institute of Cognitive Science, Osnabrück University, Germany

Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021).

Figure 1: Overview of the equivariant Transformer architecture. Thin lines: scalar features in
RF , thick lines: vector features in R3×F , dashed lines: multiple feature vectors. (a) Transformer
consisting of an embedding layer, update layers and an output network. (b) Residual update layer
including attention based interatomic interactions and information exchange between scalar andvector features. (c) Modified dot-product attention mechanism, scaling values (blue) by the attention
weights (red).

shift towards rotationally equivariant networks [Anderson et al., 2019, Fuchs et al., 2020, Schütt et al.,2021], particularly useful when the predicted quantities are vectors and tensors, can also improve theaccuracy on scalars (e.g. energy).
In this work, we introduce an equivariant Transformer (ET) architecture for the prediction of quantum
mechanical properties. By building on top of the Transformer [Vaswani et al., 2017] architecture,
we are centering the design around the attention mechanism, achieving state-of-the-art accuracy on
multiple benchmarks while relying solely on a learned featurization of atomic types and coordinates.
Furthermore, we gain insights into the black box prediction of neural networks by analyzing theTransformer’s attention weights and comparing latent representations.

2

Methods

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 xi . Then, a series of
update layers compute interactions between pairs of atoms through a modified multi-head attention
mechanism, with which the latent atomic representations are updated. Finally, an output network
computes scalar atomwise predictions using gated equivariant blocks [Schütt et al., 2021], which
get aggregated into a single molecular prediction. This can be matched with a scalar target variable
or differentiated against atomic coordinates, providing force predictions. An illustration of thearchitecture is given in Figure 1. A detailed description can be found in the supplementary material.
2.1

Training

Models are trained from scratch using mean squared error loss and the Adam optimizer [Kingma
and Ba, 2017] with parameters β1 = 0.9, β2 = 0.999 and  = 10−8 . Linear learning rate warm-up isapplied as suggested by Vaswani et al. [2017] by scaling the learning rate with ξ = nstep. After the
steps
warm-up period, we systematically decrease the learning rate by scaling with a decay factor uponreaching a plateau in validation loss. The learning rate is decreased down to a minimum of 10−7 . We
found that weight decay and dropout do not improve generalization in this context. When training onenergies and forces, we apply exponential smoothing to the energy’s train and validation loss. New
losses are discounted with a factor of α = 0.05. See supplementary material for a complete list of
hyperparameters. The full model comprises 1.34 million parameters.
2

Table 1: Results on MD trajectories from the MD17 dataset. Scores are given by the MAE of energy
predictions (kcal/mol) and forces (kcal/mol/Å). NequIP does not provide errors on energy, for PaiNN
we include the results with lower force error out of training only on forces versus on forces andenergy. Benzene corresponds to the dataset originally released in Chmiela et al. [2017], which issometimes left out from the literature. ET results are averaged over three random splits ± standard
deviation.
Molecule

3

SchNetPhysNetDimeNetPaiNN

NequIPET

Aspirinenergyforces

0.37
1.35

0.230
0.605

0.204
0.499

0.167
0.338

0.348

0.124 ± 0.001
0.255 ± 0.007

Benzeneenergyforces

0.08
0.31

-

0.078
0.187

-

0.187

0.056 ± 0.003
0.201 ± 0.008

Ethanolenergyforces

0.08
0.39

0.059
0.160

0.064
0.230

0.064
0.224

0.208

0.054 ± 0.000
0.116 ± 0.001

Malondialdehyde

energyforces

0.13
0.66

0.094
0.319

0.104
0.383

0.091
0.319

0.337

0.079 ± 0.001
0.176 ± 0.003

Naphthalene

energyforces

0.16
0.58

0.142
0.310

0.122
0.215

0.116
0.077

0.097

0.085 ± 0.000
0.060 ± 0.002

Salicylic Acidenergyforces

0.20
0.85

0.126
0.337

0.134
0.374

0.116
0.195

0.238

0.094 ± 0.001
0.135 ± 0.006

Tolueneenergyforces

0.12
0.57

0.100
0.191

0.102
0.216

0.095
0.094

0.101

0.074 ± 0.000
0.066 ± 0.001

Uracilenergyforces

0.14
0.56

0.108
0.218

0.115
0.301

0.106
0.139

0.173

0.096 ± 0.000
0.094 ± 0.000

Experiments and Results

The MD17 [Chmiela et al., 2017] dataset consists of molecular dynamics (MD) trajectories of smallorganic molecules, including both energies and forces. Forces are predicted using the negative
gradient of the energy with respect to atomic coordinates F~i = −∂ Ê/∂~ri . We train on 1000 samples
from which 50 are used for validation. The remaining data is used for evaluation and is the basis forcomparison with other work. Separate models are trained for each molecule using a combined loss
function for energies and forces where the energy loss is multiplied with a factor of 0.2 and the force
loss with 0.8. An overview of the results and comparison to SchNet [Schütt et al., 2017b], PhysNet[Unke and Meuwly, 2019], DimeNet [Klicpera et al., 2020], PaiNN [Schütt et al., 2021] and NequIP[Batzner et al., 2021] can be found in Table 1.

MAE (eV)To evaluate the architecture’s capabilities on a large collection of off-equilibrium conformations, wetrain and evaluate the equivariant Transformer on the ANI-1 [Smith et al., 2017b] 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 on DFT energies from 80% of
the dataset, while 5% are used as validation and the remaining 15% of the data make up the test set.
Figure 2 compares the equivariant Transformer’s performance to previous methods DTNN [Schütt
et al., 2017a], SchNet [Schütt et al., 2017b], MGCN [Lu et al., 2019] and ANI [Smith et al., 2017a].

0.10
0.08
0.06
0.04
0.02
0.00

0.113

0.108
0.078

DTNN

SchNet

MGCN

0.057
0.010
ET

ANI

3

Figure 2: Comparison of testing
MAE on the ANI-1 dataset in eV. Results for DTNN, SchNet and MGCN
are provided by Lu et al. [2019]. The
ANI method refers to the ANAKINME [Smith et al., 2017a] model used
for constructing the ANI-1 dataset.

3.1

Attention Weight Analysis

Neural network predictions are notoriously difficult to interpret due to the complex nature of thelearned transformations. To shed light into the black box predictor, we extract and analyze theequivariant Transformer’s attention weights. We run inference on the ANI-1 [Smith et al., 2017b],
QM9 [Ramakrishnan et al., 2014], and MD17 [Chmiela et al., 2017] test sets for all molecules andextract each sample’s attention matrix from all attention heads in all layers. Attention rollout [Abnarand Zuidema, 2020] under the single head assumption is applied during the extraction, resulting ina single attention matrix per sample. We average attention weights over each unique combinationof interacting atom types, leaving two attention scores for each pair of atom types, one from theperspective of z1 attending z2 and vice versa.
The attention scores are compared to bond probabilities extracted from the same molecules to makesure the network does not simply attend interacting atoms proportional to the relative frequency inthe dataset. Figure 3 presents a summary of the distilled probabilities and attention scores for QM9,
ANI-1, and the average attention scores for all MD17 models. We normalize each row to sum to one.

Attention scores

Bond probabilities

H

C

ANI-1
N

O

F

H

C

MD17
N

O

F

H

C

QM9
N

O

F

H

H

C

C

zi N

N

O

O

F

F

H

H

C

C

zi N

N

O

O

F

F
H

C

N

zj

O

F

H

C

N

zj

O

F

H

C

N

zj

O

F

Figure 3: Depiction of bond probabilities and attention scores extracted from the ET using QM9(total energy U0 ), MD17 (average over 8 discussed molecules) and ANI-1 testing data. Attention
scores are given as zi attending zj , bond probabilities follow the same idea, showing the conditional
probability of a bond between zi and zj , given zi . Darker colors correspond to larger values, element
pairs without data are grayed out.

4

Discussion

In this work, we introduce a novel attention-based architecture for the prediction of quantum mechanical properties, leveraging the use of rotationally equivariant features. We set a new state-of-the-art
on all MD17 targets (except force prediction of the molecule Benzene) and demonstrate the architecture’s ability to work in a low data regime. By extracting and analyzing the model’s attention
weights, we gain insights into the molecular representation, which is characterized by the natureof the corresponding training data. We show that the model does not pay much attention to thelocation of hydrogen when trained only on energy-minimized molecules, while a model trained ondata including off-equilibrium conformations focuses to a large degree on hydrogen. Neural networks
and especially Transformers are known to require large amounts of training data and computational
power. It should be taken into consideration that training these kinds of models requires significant
amounts of energy and causes the emission of greenhouse gases.
4Software and Data
The equivariant Transformer is implemented in PyTorch [Paszke et al., 2019], using PyTorch Geometric [Fey and Lenssen, 2019] as the underlying framework for geometric deep learning. Training isdone using pytorch-lightning [Falcon and The PyTorch Lightning team, 2019], a high-level interface
for training PyTorch models. The datasets QM92 , MD173 and ANI-14 are publicly available andall source code for training, running and analyzing the models presented in this work is available athttps://github.com/torchmd/torchmd-net.

References
S. Abnar and W. Zuidema. Quantifying Attention Flow in Transformers. arXiv:2005.00928 [cs],May 2020. URL http://arxiv.org/abs/2005.00928. arXiv: 2005.00928.
B. Anderson, T.-S. Hy, and R. Kondor. Cormorant: Covariant Molecular Neural Networks.
arXiv:1906.04015 [physics, stat], Nov. 2019. URL http://arxiv.org/abs/1906.04015.
arXiv: 1906.04015.
S. Batzner, A. Musaelian, L. Sun, M. Geiger, J. P. Mailoa, M. Kornbluth, N. Molinari, T. E. Smidt,
and B. Kozinsky. SE(3)-Equivariant Graph Neural Networks for Data-Efficient and AccurateInteratomic Potentials. arXiv:2101.03164 [cond-mat, physics:physics], July 2021. URL http://arxiv.org/abs/2101.03164. arXiv: 2101.03164.
J. Behler. Atom-centered symmetry functions for constructing high-dimensional neural network
potentials. The Journal of Chemical Physics, 134(7):074106, Feb. 2011. ISSN 0021-9606. doi: 10.
1063/1.3553717. URL https://aip.scitation.org/doi/10.1063/1.3553717. Publisher:
American Institute of Physics.
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, 3(5):e1603015, May 2017. ISSN 2375-2548. doi: 10.1126/sciadv.1603015. URL https://advances.sciencemag.org/content/3/5/e1603015. Publisher: American Association
for the Advancement of Science Section: Research Article.
K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Davis,
A. Mohiuddin, L. Kaiser, D. Belanger, L. Colwell, and A. Weller. Rethinking Attention with
Performers. arXiv:2009.14794 [cs, stat], Mar. 2021. URL http://arxiv.org/abs/2009.
14794. arXiv: 2009.14794.
W. Falcon and The PyTorch Lightning team. PyTorch Lightning, Mar. 2019. URL https://github.
com/PyTorchLightning/pytorch-lightning.
M. Fey and J. E. Lenssen. Fast Graph Representation Learning with PyTorch Geometric, May 2019.
URL https://github.com/rusty1s/pytorch_geometric.
F. B. Fuchs, D. E. Worrall, V. Fischer, and M. Welling. SE(3)-Transformers: 3D Roto-Translation
Equivariant Attention Networks. arXiv:2006.10503 [cs, stat], Nov. 2020. URL http://arxiv.
org/abs/2006.10503. arXiv: 2006.10503.
J. Hermann, Z. Schätzle, and F. Noé. Deep neural network solution of the electronic Schrödinger
equation. Nature Chemistry, 12(10):891–897, Oct. 2020. ISSN 1755-4330, 1755-4349. doi:
10.1038/s41557-020-0544-y. URL http://arxiv.org/abs/1909.08423. arXiv: 1909.08423.
D. P. Kingma and J. Ba. Adam: A Method for Stochastic Optimization. arXiv:1412.6980 [cs], Jan.
2017. URL http://arxiv.org/abs/1412.6980. arXiv: 1412.6980.
J. Klicpera, J. Groß, and S. Günnemann. Directional Message Passing for Molecular Graphs.
arXiv:2003.03123 [physics, stat], Mar. 2020. URL http://arxiv.org/abs/2003.03123.
arXiv: 2003.03123.
2

https://doi.org/10.6084/m9.figshare.c.978904.v5
http://www.quantum-machine.org/gdml/#datasets
4
https://figshare.com/articles/dataset/ANI-1x_Dataset_Release/10047041/1
3

5

C. Lu, Q. Liu, C. Wang, Z. Huang, P. Lin, and L. He. Molecular Property Prediction: A Multilevel
Quantum Interactions Modeling Perspective. Proceedings of the AAAI Conference on Artificial
Intelligence, 33(01):1052–1060, July 2019. ISSN 2374-3468. doi: 10.1609/aaai.v33i01.33011052.
URL https://ojs.aaai.org/index.php/AAAI/article/view/3896. Number: 01.
A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein,
L. Antiga, A. Desmaison, A. Köpf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy,
B. Steiner, L. Fang, J. Bai, and S. Chintala. PyTorch: An Imperative Style, High-Performance
Deep Learning Library. arXiv:1912.01703 [cs, stat], Dec. 2019. URL http://arxiv.org/abs/
1912.01703. arXiv: 1912.01703.
D. Pfau, J. S. Spencer, A. G. D. G. Matthews, and W. M. C. Foulkes. Ab initio solution of the
many-electron Schrödinger equation with deep neural networks. Physical Review Research, 2(3):033429, Sept. 2020. doi: 10.1103/PhysRevResearch.2.033429. URL https://link.aps.org/
doi/10.1103/PhysRevResearch.2.033429. Publisher: American Physical Society.
Z. Qiao, M. Welborn, A. Anandkumar, F. R. Manby, and T. F. Miller III. OrbNet: Deep Learning
for Quantum Chemistry Using Symmetry-Adapted Atomic-Orbital Features. arXiv:2007.08026[physics], Sept. 2020. doi: 10.1063/5.0021955. URL http://arxiv.org/abs/2007.08026.
arXiv: 2007.08026.
R. Ramakrishnan, P. O. Dral, M. Rupp, and O. A. von Lilienfeld. Quantum chemistry structures andproperties of 134 kilo molecules. Scientific Data, 1(1):140022, Aug. 2014. ISSN 2052-4463. doi:
10.1038/sdata.2014.22. URL https://www.nature.com/articles/sdata201422. Number:
1 Publisher: Nature Publishing Group.
M. Rupp, A. Tkatchenko, K.-R. Müller, and O. A. von Lilienfeld. Fast and Accurate Modelingof Molecular Atomization Energies with Machine Learning. Physical Review Letters, 108(5):058301, Jan. 2012. doi: 10.1103/PhysRevLett.108.058301. URL https://link.aps.org/doi/
10.1103/PhysRevLett.108.058301. Publisher: American Physical Society.
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, 8(1):13890, Jan. 2017a.
ISSN 2041-1723. doi: 10.1038/ncomms13890. URL https://www.nature.com/articles/
ncomms13890. Publisher: Nature Publishing Group.
K. T. Schütt, P.-J. Kindermans, H. E. Sauceda, S. Chmiela, A. Tkatchenko, and K.-R. Müller.
SchNet: A continuous-filter convolutional neural network for modeling quantum interactions.
arXiv:1706.08566 [physics, stat], Dec. 2017b. URL http://arxiv.org/abs/1706.08566.
arXiv: 1706.08566.
K. T. Schütt, H. E. Sauceda, P.-J. Kindermans, A. Tkatchenko, and K.-R. Müller. SchNet – A deeplearning architecture for molecules and materials. The Journal of Chemical Physics, 148(24):241722, Mar. 2018. ISSN 0021-9606. doi: 10.1063/1.5019779. URL https://aip.scitation.
org/doi/abs/10.1063/1.5019779. Publisher: American Institute of Physics.
K. T. Schütt, O. T. Unke, and M. Gastegger. Equivariant message passing for the predictionof tensorial properties and molecular spectra. arXiv:2102.03150 [physics], June 2021. URL
http://arxiv.org/abs/2102.03150. arXiv: 2102.03150.
J. S. Smith, O. Isayev, and A. E. Roitberg. ANI-1: an extensible neural network potential with
DFT accuracy at force field computational cost. Chemical Science, 8(4):3192–3203, Mar. 2017a.
ISSN 2041-6539. doi: 10.1039/C6SC05720A. URL https://pubs.rsc.org/en/content/
articlelanding/2017/sc/c6sc05720a. Publisher: The Royal Society of Chemistry.
J. S. Smith, O. Isayev, and A. E. Roitberg. ANI-1, A data set of 20 million calculated off-equilibrium
conformations for organic molecules. Scientific Data, 4(1):170193, Dec. 2017b. ISSN 2052-4463.
doi: 10.1038/sdata.2017.193. URL https://www.nature.com/articles/sdata2017193.
Publisher: Nature Publishing Group.
A. Szabo and N. S. Ostlund. Modern Quantum Chemistry: Introduction to Advanced Electronic
Structure Theory. Dover Books on Chemistry. Dover Publications, July 1996. ISBN 978-0-48669186-2. Google-Books-ID: 6mV9gYzEkgIC.
6

O. T. Unke and M. Meuwly. PhysNet: A Neural Network for Predicting Energies, Forces, Dipole
Moments and Partial Charges. arXiv:1902.08408 [physics], Mar. 2019. doi: 10.1021/acs.jctc.
9b00181. URL http://arxiv.org/abs/1902.08408. arXiv: 1902.08408.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin.
Attention Is All You Need. arXiv:1706.03762 [cs], Dec. 2017. URL http://arxiv.org/abs/
1706.03762. arXiv: 1706.03762.

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes](c) Did you discuss any potential negative societal impacts of your work? [Yes](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A](b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes](b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes](c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [No]4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes](b) Did you mention the license of the assets? [No](c) Did you include any new assets either in the supplemental material or as a URL? [Yes](d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [Yes](e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]7

AHyperparameters

Table 2 provides an overview of hyperparameters used for training the ET on MD17 and ANI-1.
Table 2: Comparison of various hyperparameters used for MD17 and ANI-1.
Parameter
MD17
ANI-1
1 · 10−3
30
0.8
1,000
8
6
32
128

initial learning ratelr patience (epochs)
lr decay factor
lr warmup steps
batch size
no. layers
no. RBFs
feature dimension

B

7 · 10−45
0.5
10,000
2048
6
32
128

Neighbor Embedding

The embedding layer assigns two learned vectors to each atom type zi . One is used to encode information specific to an atom, the other takes the role of a neighborhood embedding. The neighborhood
embedding, which is an embedding of the types of neighboring atoms, is multiplied by a distancefilter. This operation resembles a continuous-filter convolution [Schütt et al., 2017b] but, as it is usedin the first layer, allows the model to store atomic information in two separate weight matrices. Thesecan be thought of as containing information that is intrinsic to an atom versus information about theinteraction of two atoms.
The distance filter is generated from expanded interatomic distances using a linear transformation
W F . First, the distance dij between two atoms i and j is expanded via a set of exponential normal
radial basis functions eRBF , defined aseRBF(dij ) = φ(dij ) exp(−βk (exp(−dij ) − µk )2 )(1)kwhere βk and µk are fixed parameters specifying the center and width of radial basis function k.
The µ vector is initialized with values equally spaced between exp(−dcut ) and 1, β is initializedas (2K −1 (1 − exp(−dcut )))−2 for all k as proposed by Unke and Meuwly [2019]. The cutoff
distance dcut was set to 5Å. The cosine cutoff φ(dij ) is used to ensure a smooth transition to 0 as dij
approaches dcut in order to avoid jumps in the regression landscape. It is given by



 
 1 cos πdij + 1 , if d ≤ d
ij
cut
dcut
φ(dij ) = 2(2)0,
if dij > dcut .
The neighborhood embedding ni for atom i is then defined as
ni =

N
X

W F eRBF (dij )an (zj )(3)

j=1

with an being the neighborhood embedding function and N the number of atoms in the graph. The
final atomic embedding xi is calculated as a linear projection of the concatenated intrinsic embedding
and neighborhood embedding [ai (zi ), ni ], resulting in
xi = W C [ai (zi ), ni ] + bC
with ai being the intrinsic embedding function. The vector features ~vi are initially set to 0.

C

Equivariant Transformer Architecture

C.1

Modified Attention Mechanism(4)

We use a modified multi-head attention mechanism (Figure 1c), extending dot-product attention, inorder to include edge data into the calculation of attention weights. The edge data, i.e. interatomic
8

distances rij , are projected into two multidimensional filters DK and DV , according to
K

K

V

V

DK = σ(W D eRBF (rij ) + bD )DV = σ(W D eRBF (rij ) + bD )(5)

The attention weights are computed via an extended dot product, i.e. an elementwise multiplicationand subsequent sum over the feature dimension, of the three input vectors: query Q, key K anddistance projection DK :K

dot(Q, K, D ) =

F
X

Qk

Kk

DkK

(6)

k

The resulting matrix is passed through a nonlinear activation function and is weighted by a cosine
cutoff (see equation 2), ensuring that atoms with a distance larger than dcut do not interact. Traditionally, the resulting attention matrix A is passed through a softmax activation, however, we replace√ −1this step with a SiLU function to preserve the distance cutoff. The softmax scaling factor of dk ,which normally rescales small gradients from the softmax function, is left out. Work by Choromanski
et al. [2021] suggests that replacing the softmax activation function in Transformers with ReLU-likefunctions might even improve accuracy, supporting the idea of switching to SiLU in this case.
We place a 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 representationis split into three equally sized vectors s1ij , s2ij , s3ij ∈ RF . The vector s3ij is scaled by the attention
matrix A and aggregated over the value-dimension, leading to an updated list of feature vectors. The
linear transformation O is used to combine the attention heads’ outputs into a single feature vectoryi ∈ R384 .
s1ij , s2ij , s3ij = split(Vj DV ij )


N
X
(7)
yi = O 
Aij · s3ij jThe attention mechanism’s output, therefore, corresponds to the updated scalar feature vectors yi andscalar filters s1ij and s2ij , which are used to weight the directional information inside the update layer.
C.2

Update Layer

The update layer (Figure 1b) is used to compute interactions between atoms (attention block) andexchange information between scalar and vector features. The updated scalar features yi from theattention block are split up into three feature vectors qi1 , qi2 , qi3 ∈ RF . The first feature vector, qi1 ,takes the role of a residual around the scaled vector features. The resulting scalar feature update ∆xi
of this update layer is then defined as∆xi = qi1 + qi2

hU1~vi , U2~vi i(8)where hU1~vi , U2~vi i denotes the scalar product of vector features ~vi , transformed by linear projections
U1 and U2 .
On the side of the vector features, scalar information is introduced through a multiplication between
qi3 and a linear projection of the vector features U3~vi . The representation is updated with equivariant
features using the directional vector between two atoms. The edge-wise directional information
is multiplied with scalar filter s2ij and added to the rescaled vector features s1ij · ~vj . The result isaggregated inside each atom, forming w~ i . The final vector feature update ∆~vi for the current update
layer is then produced by adding the weighted scalar features to the equivariant features w
~ i.
w
~i =

N
X

s1ij

~vj + s2ij

j∆~vi = w~ i + qi3

U3~vi

9

~ri − ~rjk~ri − ~rj k(9)