GemNet: Universal Directional Graph Neural Networks for Molecules

1 Introduction . Graph neural networks ( GNNs ) have shown great promise for predicting the energy and other quantum mechanical properties of molecules . They can predict these properties orders of magnitudes faster than methods from quantum chemistry – at comparable accuracy . GNNs can thus enable the accurate simulation of systems that are orders of magnitude larger . However , they still exhibit severe theoretical and practical limitations . Regular GNNs are only as powerful as the 1-Weisfeiler Lehman test of isomorphism and thus can not distinguish between certain molecules [ 45 , 60 ] . Moreover , they require a large number of training samples to achieve good accuracy . In this work we first resolve the questionable expressiveness of GNNs by proving sufficient conditions for universality in the case of invariance to translations and rotations and equivariance to permutations ; and then extending this result to rotationally equivariant predictions . Simply using the full geometric information ( e.g . all pairwise atomic distances ) in a layer does not ensure universal approximation . For example , if our model uses a rotationally invariant layer we lose the relative information between components . Such a model thus can not distinguish between features that are rotated differently . This issue is commonly known as the “ Picasso problem ” : An image model with rotationally invariant layers can not detect whether a person ’ s eyes are rotated correctly . Instead , we need a model that preserves relative rotational information and is only invariant to global rotations . To prove universality in the rotationally invariant case we extend a recent universality result based on point cloud models that use representations of the rotation group SO ( 3 ) [ 18 ] . We prove that spherical representations are actually sufficient ; full SO ( 3 ) representations are not necessary . We then generalize this to rotationally equivariant predictions by leveraging a recent result on extending invariant to equivariant 1https : //www.daml.in.tum.de/gemnet 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . predictions [ 57 ] . We then discretize spherical representations by selecting points on the sphere based on the directions to neighboring atoms . We can connect this model to GNNs by interpreting these directions as directed edge embeddings . For example , the embedding direction of atom a would be defined by atom c , resulting in the edge embedding eca . Updating the spherical representation of atom a based on atom b then corresponds to two-hop message passing between the edges eca and edb via eba , with atoms c and d defining the embedding directions . This message passing formalism naturally allows us to obtain the molecule ’ s full geometrical information ( distances , angles , and dihedral angles ) , and the direct correspondence proves the model ’ s universality . We call this edge-based two-hop message passing scheme geometric message passing , and propose multiple structural enhancements to improve the practical performance of this formalism . Based on these changes we develop the highly accurate and sample-efficient geometric message passing neural network ( GemNet ) . We furthermore show that stabilizing the variance of GemNet ’ s activations with predetermined scaling factors yields significant improvements over regular normalization layers . We investigate the proposed improvements in a range of ablation studies , and show that each of them significantly reduces the model error . These changes introduce little to no computational overhead over two-hop message passing . Altogether , GemNet outperforms previous models for force predictions on COLL by 34 % , on MD17 by 41 % , and on OC20 by 20 % on average . We observe the largest improvements for the most challenging molecules , which exhibit dynamic , non-planar geometries . In summary , our contributions are : • Showing the universality of spherical representations and two-hop message passing with directed edge embeddings for rotationally equivariant predictions . • Geometric message passing : Symmetric message passing enhanced by geometric information . • Incorporating all proposed improvements in the Geometric Message Passing Neural Network ( GemNet ) , which significantly outperforms previous methods for molecular dynamics prediction . 2 Related work . Machine learning potentials . Research on predicting a molecule ’ s energy and forces ( so-called machine learning potentials ) started with hand-fitted analytical functions and then gradually moved towards fully learned models . Arguably , classical force fields are their very first instances . They use analytical functions with coefficients that were hand-tuned based on experimental data . A popular example for these is the Merck Molecular Force Field ( MMFF94 ) [ 29 ] . The next wave of methods used kernel ridge regression based on fixed , hand-crafted molecular representations [ 3 , 9 , 23 ] . Finally , modern research mostly focusses on fully end-to-end learnable models based on GNNs [ 28 , 53 ] . These models can also be combined with molecular features from quantum mechanical calculations to improve performance [ 48 ] . We consider this combination as orthogonal research . Directional GNNs . We can also achieve equivariance and invariance to rotations without relying on group representations . Directional GNNs achieve this by representing directional information explicitly [ 54 ] or in the form of angles [ 35 ] and dihedral angles [ 25 , 39 ] . Our work is focused on this class of models , proving their universality and proposing an improved variant , GemNet . Expressiveness of GNNs . A large part of GNN research has been focused on their ( limited ) expressiveness . Morris et al . [ 45 ] , Xu et al . [ 60 ] first proved that they are only as expressive as the Weisfeiler-Lehman test of isomorphism and Garg et al . [ 27 ] showed the limitations of basic directional message passing . Kondor et al . [ 37 ] , Maron et al . [ 43 ] , Morris et al . [ 45 , 46 ] then investigated higherorder representations to circumvent this issue . Finally , Azizian & Lelarge [ 2 ] , Maron et al . [ 42 ] showed that so-called folklore GNNs are the most expressive GNNs for a given tensor order . Equivariant neural networks . Equivariance and invariance have recently emerged as one of the foundational principles of modern neural networks [ 13 , 15 ] . This is especially relevant for models in physics , for which we often know the symmetries a priori . Equivariant models for the SO ( 3 ) group were first investigated in the context of spherical convolutions by Cohen et al . [ 14 ] , Esteves et al . [ 21 ] , Kondor et al . [ 36 ] . These methods leverage group representations to achieve full equivariance . They were then transferred to the context of 3D point clouds and molecules by Anderson et al . [ 1 ] , Thomas et al . [ 56 ] , Weiler et al . [ 58 ] , and further developed by Batzner et al . [ 4 ] , Finzi et al . [ 24 ] , Fuchs et al . [ 26 ] . Importantly , Yarotsky [ 61 ] proved the universality of 2D convolutional networks , and Bogatskiy et al . [ 5 ] extended this result to general groups . Maron et al . [ 44 ] proved universality for models invariant to Sn and equivariant to an additional symmetry . Dym & Maron [ 18 ] combined these results to prove universality for the joined group of translations , rotations , and permutations . Apart from reflections this is the exact group relevant for general molecules . 3 Universality of spherical representations . GNNs for molecules typically incorporate directional information in one of two ways : Via SO ( 3 ) representations [ 1 , 56 ] or by using directions in real space [ 35 , 54 ] . Directions in real space are associated with the three-dimensional S2 sphere , while the SO ( 3 ) group is double covered by the fourdimensional S3 sphere . Directional representations thus use one degree of freedom less than SO ( 3 ) representations , making them significantly cheaper . And , as we will prove in this section , directional representations actually provide the same expressivity as SO ( 3 ) representations for predictions in R3 . We achieve this by showing that the SO ( 3 ) -based tensor field network ( TFN ) [ 56 ] variant used by Dym & Maron [ 18 ] is equivalent to a similar model based on spherical representations , in the case of rotationally invariant predictions . We then generalize a recent result by Villar et al . [ 57 ] , which lets us extend our theorem to the rotationally equivariant case . Afterwards , we relate this universality to directional GNNs by interpreting them as a discretization of spherical representations . Preliminaries . We consider a point cloud with n points ( atoms ) , each associated with a position and a set of rotationally invariant features ( e.g . atom types ) , defined as X ∈ R3×n and Hin ∈ Rh×n . In this section we define model classes by sets of functions F . As a first step , we are interested in proving that the set F defining our model is equal to the full set of functions G′ that are invariant to the group of translations T3 and rotations SO ( 3 ) , and equivariant to the group of permutations Sn . We denote the codomain of functions in G′ as WnT , where WT is some representation of SO ( 3 ) . We denote a vector ’ s norm by x = ∥x∥2 , its direction by x̂ = x/x , and the relative position by xba = xb − xa . Proofs are deferred to the appendix . Note that this section is not intended as an introduction to the SO ( 3 ) group . For a concise introduction in the context of machine learning see e.g . Weiler et al . [ 58 , Section 3 ] or Kondor et al . [ 36 ] . Tensor field network . In order to show the equivalence of the TFN to spherical representations , we first need to define this model . Following Dym & Maron [ 18 ] , we split the model into two parts : Embedding functions Ffeat that lifts the input into an equivariant representation , and pooling functions Fpool that aggregate the results of multiple embedding functions on each point and computes the model output . The overall model is then defined as the set of functions FTFNK ( D ) , D = { f | f ( X , Hin ) = K∑ k=1 f ( k ) ∗ pool ( f ( k ) feat ( X , Hin ) ) , f ( k ) pool ∈ F TFN pool ( D ) , f ( k ) feat ∈ F TFN feat ( D ) } , ( 1 ) where D ∈ N denotes the function ’ s maximum polynomial degree , K ( D ) ∈ N is chosen such that Theorem 1 is fulfilled ( Dym & Maron [ 18 ] only prove the existence of this function ) , and f∗ denotes elementwise application of f on all points . We then define the set FTFNpool as all rotationally equivariant linear functions on the SO ( 3 ) group , i.e . all SO ( 3 ) convolutions [ 38 ] . Note that these are more expressive than the self-interaction layers used originally [ 56 ] . The embedding functions FTFNfeat ( D ) = { π2 ◦f ( 2D ) ◦ · · · ◦f ( 1 ) | f ( i ) ∈ FTFNprod } consist of an auxiliary function π2 ( X , H ) = H and a series of tensor product functions ( called convolution by Dym & Maron [ 18 ] ) FTFNprod = { f | f ( X , H ) = ( X , H̃TFN ( X , H ) ) } . The intermediate representations are H ∈ Wnfeat , where Wfeat is a representation of SO ( 3 ) indexed by the degree l and the order m. For Hin we have l=m=0 . The main update is defined by H̃TFN ( lo ) amo ( X , H ) = θH ( lo ) amo + ∑ lf , mf ∑ li , mi C ( lo , mo ) ( lf , mf ) , ( li , mi ) ∑ b∈Na F ( lf ) TFN , mf ( xb − xa ) H ( li ) bmi , ( 2 ) where θ is a ( learned ) scalar and Na are the neighbors of point a . The Clebsch-Gordan coefficients C ( lo , mo ) ( lf , mf ) , ( li , mi ) arise from decomposing the tensor product of two input SO ( 3 ) representations ( the filter and input representations ) into a sum of output representations . Their exact values are not relevant for this discussion . We index the output with degree lo and order mo , the learned filter with lf and mf , and the input with li and mi . F ( l ) TFN , m ( x ) = R ( l ) ( x ) Ylm ( x̂ ) is a rotationally equivariant filter , with a ( learned ) radial part R , which is any polynomial of degree ≤ D , and the real spherical harmonics Ylm with degree l and order m. The spherical harmonics are the basis for the Fourier transformation of functions on the sphere , analogously to sine waves for functions on R. We can prove universality for TFNs by using the universality of polynomial regression and showing that TFNs can fit any polynomial ( see Dym & Maron [ 18 ] for details ) , resulting in : Theorem 1 ( Dym & Maron [ 18 ] ) . Consider the set of functions G mapping R3×n+h×n → WnT that are equivariant to rotations and permutations and invariant to translations . For all n ∈ N , 1 . For D ∈ N0 , every polynomial p ∈ G of degree D is in FTFNK ( D ) , D. 2 . Every continuous function f ∈ G can be approximated uniformly on compact sets by functions in⋃ D∈N0 F TFN K ( D ) , D. Spherical networks . Instead of intermediate SO ( 3 ) representations we now switch to spherical representations , which are functions on the sphere H : S2 → R. We define the set of functions F sphereK ( D ) , D analogously to F TFN K ( D ) , D. However , for F sphere feat ( D ) we use H̃ spherea ( X , H ) ( r̂ ) = θHa ( r̂ ) + ∑ b∈Na Fsphere ( xb − xa , r̂ ) Hb ( r̂ ) , ( 3 ) with the filter function Fsphere ( x , r̂ ) = ∑ l , m R ( l ) ( x ) ℜ [ Y ( l ) ∗m ( x̂ ) Y ( l ) m ( r̂ ) ] , using the real part ℜ of the complex spherical harmonics Y ( l ) m . The set of pooling functions for invariant predictions is F spherepool = { f | f ( H ) = θpool ∫ S2 H ( r̂ ) dr̂ } , ( 4 ) with the learnable parameter θpool . We obtain the universality theorem by showing the equivalence between this model and TFN for rotationally invariant functions . The proof is based on the connection between spherical harmonics and the Clebsch-Gordan coefficients [ 51 , 3.7.72 ] ( see App . A ) . Theorem 2 . Consider the set of functions G′ mapping R3×n+h×n → WnT that are equivariant to permutations and invariant to translations and rotations . For all n ∈ N , 1 . For D ∈ N0 , every polynomial p ∈ G′ of degree D is in F sphereK ( D ) , D. 2 . Every continuous function f ∈ G′ can be approximated uniformly on compact sets by functions in⋃ D∈N0 F sphere K ( D ) , D. Next , we extend Theorem 2 to rotationally equivariant functions . We do this by generalizing a recent result by Villar et al . [ 57 ] to obtain ( see App . B ) : Theorem 3 . Let h : Rd×n+h×n → Rd×n be any function that is equivariant to permutations and rotations and invariant to translations . For all a ∈ [ 1 , n ] , let the set of relative vectors { xca | c ∈ [ 1 , n ] } not span a ( d−1 ) -dimensional space . Then there are n−1 functions f ( c ) : Rd×n+h×n → Rn such that ha ( X , H ) = n∑ c=1 c̸=a f ( c ) a ( X , H ) xca , ( 5 ) where f ( c ) is equivariant to permutations , but invariant to rotations and translations . This theorem lets us extend a rotationally invariant model to an equivariant one , while providing universality guarantees . Together , Theorem 2 and Theorem 3 ( with d = 3 ) thus show that we can approximate any rotationally equivariant function using only representations on the S2 sphere . We thus do not need SO ( 3 ) representations , spin-weighted spherical harmonics [ 22 ] , triplet embeddings , or complex-valued functions . This result puts theory back in line with practice , where the best results are currently achieved without relying on these more expensive representations [ 54 ] .

Summary: Many GNN models have recently been proposed for the problem of predicting molecular dynamics. These GNNs rely on rotationally invariant predictions since the underlying physics of the problem does not change with rotations. While such models have obtained good performance, there is limited theoretical understanding of these methods.  This paper makes two contributions to this task: first, it addresses the theoretical limitations by proving that GNNs with directed edge embeddings and two-hop message passing are universal approximators; second, it presents a new model called GemNet (an extension of Dimenet++ that uses two hop message passing), that obtains excellent results on COLL and MD17 datasets.