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