Equivariant Transformers for Neural Network based Molecular Potentials

1 INTRODUCTION . Quantum mechanics are essential for the computational analysis and design of molecules and materials . However , the complete solution of the Schrö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 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 more accurate 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 quantum Monte Carlo ( Pfau et al. , 2020 ; Hermann et al. , 2020 ) . Other methods include force-field and semiempirical 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 & Parrinello , 2007 ) and the seminal work from Rupp et al . ( 2012 ) . 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 field of materials science and quantum chemistry , uses more modern graph convolutions ( Schütt et al. , 2018 ; Unke & Meuwly , 2019 ; Qiao et al. , 2020 ; Schütt et al. , 2021 ) . SchNet ( Schütt et al. , 2017b ; 2018 ) , for example , uses continuous filter 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 ; Schütt 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 ( Schütt et al. , 2017b ; Qiao et al. , 2020 ; Klicpera et al. , 2020b ; Unke & Meuwly , 2019 ; Liu et al. , 2020 ; Schütt et al. , 2021 ) . For example , Qiao et al . ( 2020 ) first perform a low-cost mean-field 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 TorchMD-NET , 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 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 ( MD17 ( Chmiela et al. , 2017 ) and normal mode sampled data ( ANI-1 ( Smith et al. , 2017b ) ) . 2 METHODS . 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 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 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 , a layer normalization ( Ba et al. , 2016 ) followed by an output network computes scalar atomwise predictions using gated equivariant blocks ( Weiler et al. , 2018 ; Schü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 the architecture is given in Figure 1 . 2.1 NOTATION . To differentiate between the concepts of scalar and vector features , this work follows a certain notation . Scalar features are written as x ∈ RF , while we refer to vector features as ~v ∈ R3×F . The vector norm ‖·‖ and scalar product 〈· , ·〉 of vector features are applied to the spatial dimension , while all other operations act on the feature dimension . Upper case letters denote matrices A ∈ RN×M . 2.2 EMBEDDING LAYER . 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 distance filter . This operation resembles a continuous-filter convolution ( Schütt et al. , 2017b ) but , as it is used in the first 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 filter is generated from expanded interatomic distances using a linear transformation WF . First , the distance dij between two atoms i and j is expanded via a set of exponential normal radial basis functions eRBF , defined as eRBFk ( dij ) = φ ( dij ) exp ( −βk ( exp ( −dij ) − µk ) 2 ) ( 1 ) where β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 initialized as ( 2K−1 ( 1 − exp ( −dcut ) ) ) −2 for all k as proposed by Unke & Meuwly ( 2019 ) . The cutoff distance dcut was set to 5Å . 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 φ ( dij ) = { 1 2 ( cos ( πdij dcut ) + 1 ) , if dij ≤ dcut 0 , if dij > dcut . ( 2 ) The neighborhood embedding ni for atom i is then defined as ni = N∑ j=1 embednbh ( zj ) WF eRBF ( dij ) ( 3 ) with embednbh 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 [ embedint ( zi ) , ni ] , resulting in xi = W C [ embedint ( zi ) , ni ] + bC ( 4 ) with embedint being the intrinsic embedding function . The vector features ~vi are initially set to 0 . 2.3 MODIFIED ATTENTION MECHANISM . We use a modified 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 rij , are projected into two multidimensional filters DK and DV , according to DK = σ ( WD K eRBF ( rij ) + b DK ) DV = σ ( WD V eRBF ( rij ) + b DV ) ( 5 ) 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 DK : Q = WQxi and K = WKxi ( 6 ) dot ( Q , K , DK ) = F∑ k Qk Kk DKk ( 7 ) 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 . A = SiLU ( dot ( Q , K , DK ) ) · φ ( dij ) ( 8 ) 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 √ dk −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 filter graph convolution ( Schü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 representation is split into three equally sized vectors s1ij , s 2 ij , s 3 ij ∈ 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 vector yi ∈ R384 . s1ij , s 2 ij , s 3 ij = split ( Vj DV ij ) yi = O N∑ j Aij · s3ij ( 9 ) The attention mechanism ’ s output , therefore , corresponds to the updated scalar feature vectors yi and scalar filters s1ij and s 2 ij , which are used to weight the directional information inside the update layer .

Summary: 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.