T ORCH MD-ET: E QUIVARIANT T RANSFORMERS FOR N EURAL N ETWORK BASED M OLECULAR P OTENTIALS
A BSTRACT
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 efﬁciency comparable with classical force ﬁelds.
In this work we propose a novel equivariant Transformer architecture (TorchMDET), outperforming state-of-the-art on MD17, ANI- and many QM9 targets in both accuracy and computational efﬁciency.
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
I NTRODUCTION
Quantum mechanics are essential for the computational analysis and design of molecules and materials.
However, the complete solution of the Schr ¨ odinger 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 atomic systems according to density functional theory (DFT), which can provide energy estimates with sufﬁciently 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 efﬁciency to be applied on a 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-ﬁeld and semiempirical quantum mechanical theories, which provide very efﬁcient estimates but lack accuracy.
The ﬁeld of machine learning molecular potentials is relatively novel.
The ﬁrst important contributions are rooted in the Behler-Parrinello (BP) representation (Behler & Parrinello, 2007) and the seminal work from Rupp et al.
One of the best transferable machine learning potentials for biomolecules, called ANI (Smith et al., 2017a), is based on BP.
A second class of methods, mainly developed in the ﬁeld of materials science and quantum chemistry, uses more modern graph convolutions (Sch¨utt et al., 2018; Unke & Meuwly, 2019; Qiao et al., 2020; Sch¨utt et al., 2021).
SchNet (Sch¨utt et al., 2017b; 2018), for example, uses continuous ﬁlter convolutions in a 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 extended to coupled-cluster solvers (Hermann et al., 2020) and protein folding using coarse-grained systems (Wang et al., 2019;
Husic et al., 2020; Doerr et al., 2021).
Recently, other work has shown that a shift towards rotationally equivariant networks (Anderson et al., 2019; Fuchs et al., 2020;
Sch¨utt et al., 2021), particularly useful when the predicted quantities are vectors and tensors, can also improve the accuracy on 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 a scaling advantage to large amounts of data, while kernel methods usually work best in a 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¨utt et al., 2017b;
Qiao et al., 2020; Klicpera et al., 2020b; Unke & Meuwly, 2019; Liu et al., 2020; Sch¨utt et al., 2021).
For example, Qiao et al.
ﬁrst perform a low-cost mean-ﬁeld electronic 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 attention (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 an equivariant Transformer 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 the Transformer’s attention weights and comparing latent representations between different types of data such as energy-minimized (QM9 (Ramakrishnan et al., 2014)), molecular dynamics (MD (Chmiela et al., 2017) and normal mode sampled data (ANI-1 (Smith et al., 2017b)).
M ETHODS
The traditional Transformer architecture as proposed by Vaswani et al.
(2017) operates on a sequence of tokens.
In the context of chemistry, however, the natural data structure for the representation of molecules is a graph.
To work on graphs, one can interpret self-attention as constructing a fully connected graph over input tokens and computing interactions between nodes.
We leverage this concept and extend it to include information stored in the graph’s edges, corresponding to interatomic distances in the context of molecular data.
This requires the use of a modiﬁed attention mechanism, which we introduce in the following sections, along with the overall architecture of our equivariant
Transformer (TorchMD-ET).
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 x i .
Then, a series of update layers compute interactions between pairs of atoms through a modiﬁed 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¨utt 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 the architecture is given in Figure 1.
2. N OTATION
To differentiate between the concepts of scalar and vector features, this work follows a certain notation.
Scalar features are written as x ∈ R F , while we refer to vector features as (cid:126)v ∈ R 3 × F .
The vector norm (cid:107)·(cid:107) and scalar product (cid:104)· , ·(cid:105) of vector features are applied to the spatial dimension, while all other operations act on the feature dimension.
Upper case letters denote matrices A ∈ R N × M .
2. E MBEDDING LAYER
The embedding layer assigns two learned vectors to each atom type z i .
One is used to encode information speciﬁc 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 distance ﬁlter.
This operation resembles a continuous-ﬁlter convolution (Sch ¨ utt et al., 2017b) but, as it is used in the ﬁrst layer, allows the model to 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 ﬁlter is generated from expanded interatomic distances using a linear transformation
W F .
First, the distance d ij between two atoms i and j is expanded via a set of exponential normal radial basis functions e RBF , deﬁned as
[Equation]
where β k and µ k are ﬁxed parameters specifying the center and width of radial basis function k .
The µ vector is initialized with values equally spaced between exp( − d cut ) and 1, β is initialized as (2 K − 1 (1 − exp( − d cut )))
− 2 for all k as proposed by Unke & Meuwly (2019).
The cutoff distance d cut was set to 5 ˚A.
The cosine cutoff φ ( d ij ) is used to ensure a smooth transition to 0 as d ij approaches d cut in order to avoid jumps in the regression landscape.
It is given by
[Equation]
The neighborhood embedding n i for atom i is then deﬁned as
[Equation]
with embed nbh being the neighborhood embedding function and N the number of atoms in the graph.
The ﬁnal atomic embedding x i is calculated as a linear projection of the concatenated intrinsic embedding and neighborhood embedding (cid:2) embed int ( z i ) , n i (cid:3) , resulting in
[Equation]
with embed int being the intrinsic embedding function.
The vector features (cid:126)v i are initially set to 0.
2. M ODIFIED A TTENTION M ECHANISM
We use a modiﬁed multi-head attention mechanism (Figure 1c), extending dot-product attention, in order to include edge data into the calculation of attention weights.
First, the feature vectors are passed through a layer normalization.
Then, edge data, i.e. interatomic distances r ij , are projected into two multidimensional ﬁlters D K and D V , according to
[Equation]
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 K :
[Equation]
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 d cut do not interact.
[Equation]
Traditionally, the resulting attention matrix A is passed through a softmax activation, however, we replace this step with a SiLU function to preserve the distance cutoff.
The softmax scaling factor of √ d k − 1 , 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-like functions might even improve accuracy, supporting the idea of switching to SiLU in this case.
We place a continuous ﬁlter graph convolution (Sch¨utt 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 representation is split into three equally sized vectors s 1 ij , s 2 ij , s 3 ij ∈ R F .
The vector s 3 ij 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 vector
y i ∈ R 384 .
[Equation]
The attention mechanism’s output, therefore, corresponds to the updated scalar feature vectors y i and scalar ﬁlters s 1 ij and s 2 ij , which are used to weight the directional information inside the update layer.
2. U PDATE L AYER
The update layer (Figure 1b) is used to compute interactions between atoms (attention block) and exchange information between scalar and vector features.
The updated scalar features y i from the attention block are split up into three feature vectors q 1 i , q 2 i , q 3 i ∈ R F .
The ﬁrst feature vector, q 1 i , takes the role of a residual around the scaled vector features.
The resulting scalar feature update ∆ x i of this update layer is then deﬁned as
[Equation]
) where (cid:104) U 1 (cid:126)v i , U 2 (cid:126)v i (cid:105) denotes the scalar product of vector features (cid:126)v i , transformed by linear projections U 1 and U 2 .
On the side of the vector features, scalar information is introduced through a multiplication between q 3 i and a linear projection of the vector features U 3 (cid:126)v i .
The representation is updated with equivariant features using the directional vector between two atoms.
The edge-wise directional information is multiplied with scalar ﬁlter s 2 ij and added to the rescaled vector features s 1 ij · (cid:126)v j .
The result is aggregated inside each atom, forming (cid:126)w i .
The ﬁnal vector feature update ∆ (cid:126)v i for the current update layer is then produced by adding the weighted scalar features to the equivariant features (cid:126)w i .
[Equation]
2. T RAINING
Models are trained using mean squared error loss and the Adam optimizer (Kingma & Ba, 2017) with parameters β 1 = 0 .
9 , β 2 = 0 .
999 and (cid:15) =
10 − 8 .
Linear learning rate warm-up is applied as suggested by Vaswani et al.
(2017) by scaling the learning rate with ξ = step n steps .
After the warm-up period, we systematically decrease the learning rate by scaling with a decay factor upon reaching 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 on energies and forces, we apply exponential smoothing to the energy’s train and validation loss.
New losses are discounted with a factor of α = 0 .
See Appendix A for a more comprehensive summary of hyperparameters.
E XPERIMENTS AND R ESULTS
We evaluate the equivariant Transformer on the QM9 (Ramakrishnan et al., 2014), MD17 (Chmiela et al., 2017) and ANI-1 (Smith et al., 2017b) benchmark datasets.
QM9 comprises 133,885 small organic molecules with up to nine heavy atoms of type C, 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 1 compares the equivariant Transformer’s results on QM9 with the invariant architectures
SchNet (Sch ¨ utt et al., 2018), PhysNet (Unke & Meuwly, 2019) and DimeNet++ (Klicpera et al., 2020a), the covariant Cormorant (Anderson et al., 2019) architecture and equivariant EGNN (Satorras et al., 2021), LieTransformer (we compare to their best variant, LieTransformer-T3+SO3 Aug) (Hutchinson et al., 2020) and PaiNN (Sch ¨ utt 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 µ , both scalar and vector features are used in the ﬁnal calculation.
The output MLP consists of two gated equivariant blocks (Weiler et al., 2018;
Sch¨utt et al., 2021) with the same layer sizes as in the otherwise used output network.
The updated scalar features x i and vector features (cid:126)v i are then used to compute µ as
[Equation]
[Equation]
where (cid:126)r is the center of mass of the molecule.
For the prediction of the electronic spatial extent (cid:10)
R 2 (cid:11) , scalar features are transformed using gated equivariant blocks as described above, yielding scalar atomic predictions x i , and multiplied by the squared norm of atomic positions
[Equation]
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 to atomic coordinates (cid:126)F i = − ∂
ˆ E/∂(cid:126)r i .
To evaluate the architecture’s performance in a limited data setting, the model is trained on onlysamples from which 50 are 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 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 the invariant models
[Table]
[Table]
Toluene
[Table] Uracil [Table] SchNet (Sch¨utt et al., 2017b), PhysNet (Unke & Meuwly, 2019) and DimeNet (Klicpera et al., 2020b), as well as the equivariant architectures PaiNN (Sch¨utt et al., 2021) and NequIP (Batzner et al., 2021) can be found in Table 2.
To evaluate the architecture’s capabilities on a large collection of off-equilibrium conformations, we train and evaluate the equivariant Transformer on the ANI-1 dataset.
It contains 22,057, conﬁgurations 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 ﬁtted 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.
Figurecompares the equivariant Transformer’s performance to previous methods DTNN (Sch¨utt et al., 2017a), SchNet (Sch¨utt et al., 2017b), MGCN (Lu et al., 2019) and ANI (Smith et al., 2017a).
[Figure]
3. A TTENTION W EIGHT A NALYSIS
Neural network predictions are notoriously difﬁcult 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 on the ANI-1, QM9, and MD17 test sets for all molecules and extract each sample’s attention matrix from all attention heads in all 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 4 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 over each unique combination of interacting atom types (hydrogen, carbon, oxygen, nitrogen, ﬂuorine).
This generates two attention scores for each pair of atom types, one from the perspective of atom type z 1 attending z 2 and vice versa.
The attention scores are compared to the probability of this bond occurring in the respective dataset, making sure the network’s attention is not simply proportional to the relative frequency of the interaction.
Figure 3 presents a 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 conformation, we expect it to pay attention to atoms that are displaced from the equilibrium conformation, i.e. the energy-minimized structure, of the molecule.
We test this hypothesis by comparing the attention weights of displaced atoms to those of the equilibrium conformation.
Using the QM9 test set as the source of equilibrium conformations, we displace single atoms by 0.4 ˚A in a random direction and compare the absolute magnitude of attention weights involving the displaced atom to the remaining attention weights.
We ﬁnd 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 hydrogen in its equilibrium position.
Attention for displaced hydrogen atoms even 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 a detailed overview of the results, see Appendix C.
It is interesting to see that the training dataset inﬂuences attention.
For static structures, like in QM9, attention analysis shows that very little importance is attributed to hydrogens, while core 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. A BLATION S TUDIES
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 a regular atom type embedding, which we evaluate by comparing the testing MAE on U 0 from QM9.
The ablation causes a drop in accuracy of roughly 6% (from 6.24 to 6.60 meV).
We also try replacing the neighbor embedding by an additional update layer as the neighbor embedding resembles a graph convolution operation, which is the equivalent operation to an update layer in graph convolutional networks such as SchNet.
Here, the drop in accuracy is even more pronounced with a decrease of around 10% (from 6.24 to 6.85 meV).
However, this may also be the result of overﬁtting as each update layer has about 4.6 times more parameters than a neighbor embedding layer. B ond p r obab ili t i e s B ond p r obab ili t i e s B ond p r obab ili t i e s
H
C
N O F H
C
N
O
F z i
ANIH
C
N
O
F
MD
H
C
N
O
F
H
C
N
O
F
QM
A tt en t i on sc o r e s
A tt en t i on sc o r e s
A tt en t i on sc o r e s
H
C
N
O
F z j
H
C
N
O
F z i
H
C
N
O
F z j
H
C
N
O
F z j
H
C
N
O
F
The hyperparameter set used for MD17 and ANI-1 results in a model size of 1.34M trainable parameters, which is signiﬁcantly more than recent similar architectures such as PaiNN (600k) or NequIP (290k).
To rule out the possibility that TorchMD-ET’s results is simply caused by an increased number of parameters, we train smaller versions of TorchMD-ET, which are comparable to the size of PaiNN and NequIP.
We ﬁnd that smaller versions of the TorchMD-ET are still competitive and outperform previous state-of-the-art results on MD17.
See Appendix D for details on the results, hyperparameters and computational efﬁciency.
3. C OMPUTATIONAL E FFICIENCY
To assess the computational efﬁciency 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 3).
We report times for different sizes of TorchMD-ET, differing in the number of update layers, the feature dimension, and the size of the RBF distance expansion.
TorchMD-ET-large uses the QM9 hyperparameter set, while TorchMD-ET-small is constructed using MD17/ANI-1 hyperparameters (see Table 4).
The measurements were conducted using just-in-time (JIT) compiled models.
The JIT-compiled versions of TorchMD-ET-small and TorchMD-ET-large are 1.5 and 1.8 times respectively more efﬁcient 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.0ms per batch using TorchMD-ET-small and 48.5ms using TorchMD-ET-large.
[Figure]
[GDB 105807] [GDB 11383]
[GDB 120698]
[GDB 5605]
[GDB 96582]
[Figure]
[Figure]
ned
[Figure]
[Table]
D ISCUSSION
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 show a high degree of accuracy on the QM9 benchmark dataset, however, the architecture’s effectiveness is particularly clear when looking at the prediction of energies and atomic forces in the context of molecular dynamics.
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.
As described in previous work (Sch¨utt et al., 2021; Fuchs et 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 is capable of predicting tensorial properties.
By extracting and analyzing the model’s attention weights, we gain insights into the molecular representation, which is characterized by the nature of the corresponding training data.
We show that the model does not pay much attention to the location of hydrogen when trained only on energy-minimized molecules, while a model trained on data including off-equilibrium conformations focuses to a large degree on hydrogen.
Furthermore, we validate the learned representation by analyzing attention weights involving atoms displaced from their equilibrium location.
We demonstrate that off-equilibrium conformations in the training data play an important role in the accurate prediction of a molecule’s energy.
This highlights the importance of conﬁgurational diversity in the evaluation of neural network potentials.
S OFTWARE AND D ATA
The equivariant Transformer is implemented in PyTorch (Paszke et al., 2019), using PyTorch Geometric (Fey & Lenssen, 2019) as the underlying framework for geometric deep learning.
Training is done using pytorch-lightning (Falcon & The PyTorch Lightning team, 2019), a high-level interface for training PyTorch models.
The datasets QM9 1 , MD17 2 and ANI-1 3 are publicly available and all source code for training, running and analyzing the models presented in this work is available at github.com/torchmd/torchmd-net. A CKNOWLEDGMENTS
GDF thanks the project PID2020-116564GB-I00 funded by MCIN/ AEI /10.13039/501100011033, Unidad de Excelencia Mar´ıa de Maeztu” funded by the MCIN and the AEI (DOI: 10.13039/501100011033) Ref: CEX2018-000792-M and the European Union’s Horizon 2020 research 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.
A A RCHITECTURAL DETAILS AND HYPERPARAMETERS
All models in this work were trained using distributed training across two NVIDIA RTX
Ti
GPUs, using the DDP training protocol.
This leads to training times of 16h for QM9, 10h for
MD17 and 83h for ANI-1.
[Table]
While TorchMD-ET follows a similar idea as the SE(3)-Transformer introduced by Fuchs et al.
(2020), there are signiﬁcant architectural differences.
The SE(3)-Transformer relies heavily on expensive features, such as Clebsch-Gordan coefﬁcients and spherical harmonics while TorchMD-ET only requires interatomic distances.
Additionally, we split scalar and equivariant features into two pathways, which exchange information inside the update layer while the SE(3)-Transformer computes message passing updates for each type of feature vector (scalar or equivariant).
Finally, our modiﬁed attention mechanism and update step differ signiﬁcantly from the SE(3)-Transformer’s message passing layer, which, for example, does not handle self-interactions in the attention mechanism and applies only linear transformations to distance features.
B I MPORTANCE OF HYDROGEN
While TorchMD-ET trained on QM9 mostly attends to carbon atoms, models trained on molecular dynamics trajectories show a strong focus on carbon-hydrogen 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 U 0 ) and MD (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 5.
[Table]
[Table]
[Table]
C A TOM DISPLACEMENT
Figure 5 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 on equilibrium structures (QM9) focuses on carbon and neglects hydrogen.
The models trained on off-equilibrium conformations show a higher degree of attention for displaced hydrogen atoms.
We 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.
Q M 0.
H
C
O
M D 0.
A N I e q u
i l i b r i u m d i s p l a c e d 0.
e q u i l i b
r
i u m d i s p l a c e d e q u i l i
b r i u m d
i s p l a c e d m ode l t r a i ned on
D E QUIVARIANT T RANSFORMERS WITH REDUCED PARAMETER COUNT
We compare TorchMD-ET with a reduced number of parameters, matching those of PaiNN (600k) and NequIP (290k), to state-of-the-art models on the MD17 benchmark.
This ensures that our results are not solely a consequence of a larger model size but correspond to an improved architecture.
The smaller TorchMD-ET models are still competitive and outperform state-of-the-art on most MD targets.
Table 6 provides an overview of the adjusted architectural hyperparameters and Tablesummarizes the results on MD17.
[Table]
TorchMD
[Table]
ET TorchMD
[Table]
ET
[Table]
[Table]
E A BLATION OF EQUIVARIANT FEATURES
We perform an ablation of TorchMD-ET’s equivariance and compare the performance of the resulting rotationally invariant model to that of the equivariant Transformer.
Without equivariance, the MAE in total energy U 0 in QM9 increases from 6.24 to 6.64 meV (6%).
On aspirin inside the MD17 benchmark, removing the equivariance causes the energy MAE to rise from 5.37 to 13. meV (146%), while force errors go up from 11.05 to 30.27 meV/ ˚A (174%).
As, without equivariance, the error increases much more drastically on dynamical data, we hypothesize that equivariant features are particularly useful when dealing with non-zero forces.
F M OLECULAR REPRESENTATION BY DATASET
Figures 6a, 6b and 6c show a three dimensional visualization of three random molecules from the datasets QM9, MD17 (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 a visualized conformation during training.
The red and blue lines between 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 correspond to negative attention weights, which occur due to using SiLU activations inside the attention mechanism.
While the model trained on QM9 focuses largely on carbon-carbon interactions, it is clear that models trained on MD17 and ANI-1 have a strong focus on hydrogen-carbon and hydrogen-oxygen interactions.
This corresponds to the results found in the attention weights averaged over pairwise interactions between elements.
The attention weights are normalized inside each molecule.
A N Imoleculemoleculemolecule
M D ( u r a c il )
Q M
Molecules from QM m ode l t r a i ned on
[Figure]
[Figure]
A N Imoleculemoleculemolecule
M D ( u r a c il )
Q M
Molecules from MD17 (aspirin) m ode
l t r a i ned on
A N Imoleculemoleculemolecule
M D ( u r a c il )
Q M
Molecules from ANIm ode l t r a i ned on
[Figure]
G F REQUENCY OF E LEMENTS
Figure 7 shows the distribution of elements in the datasets QM9, MD17 and ANI-1 where MD 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 ﬂuorine as MD17 contains only a single molecule with nitrogen (uracil) and ﬂuorine is only found in QM9.
[Figure]
[Equation]
H
C
N
O
F
MD
H
C
N
O
F
QM
H MD17 A TTENTION W EIGHTS
Figures 8a 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 of attention for hydrogen atoms is common for most of the molecules.
However, models trained on aspirin, malondialdehyde and uracil additionally exhibit a strong focus on oxygen.
The sum over each row equals one, meaning that the probabilities represent the conditional probability
[Equation]
[Equation]
where N bonded ( z i , z j ) is the total number of bonds between atom types
z i and z j found in the collection of molecules.
B ond p r obab ili t i e s
B ond p r obab ili t i e s B ond p r obab ili t i e s
B ond p r obab ili t i e s
H
C
N
O
H
C
N
O z i
MD17Aspirin
H
C
N
O
MD17Benzene
H
C
N
O
MD17Ethanol
H
C
N
O
H
C
N
O
[Table]
A tt e n t i on sc o r e s A tt e n t i on sc o r e s
A tt e n t i on sc o r e s A tt e n t i on sc o r e s
H
C
N
O z j
H
C
N
O z i
H
C
N
O z j
H
C N
O z j
H
C
N
O z j
H
C
N
O (a) Attention scores and
B ond p r obab ili t i e s B ond p r obab ili t i e s
B ond p r obab ili t i e s B ond p r obab ili t i e s
H
C
N
O
H
C
N
O z i
MD17 Naphthalene
H
C
N
O
MD17 Salicylic Acid
H
C
N
O
MD17Toluene
H
C
N
O
H
C
N
O
MD17Uracil
A tt en t i on sc o r e s A tt en t i on sc o r e s
A tt en t i on sc o r e s A tt en t i on sc o r e s
H
C
N
O z j
H
C
N
O z i
H
C
N
O z j
H
C
N
O z j
H
C
N
O z j
H
C
N
O
[Figure]