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: * The paper proposes to use a type of networks that can be shown to be universal approximators for functions that are permutation equivariant, rotationally and translationally invariant. * This property is required to model functions that work on molecular structures. In essence, this provides a very strong prior on the functions we want to learn. * The experiments show that by enforcing this structure learning on molecular graphs can be improved (higher quality/label efficiency) compared to previous work. While, the experimental section is quite compact and heavily compressed nothing stood out as obviously wrong but I might be missing a detail.  Overall I like the high level idea behind the paper and the results seem good. While, at a high level I am able to understand the paper, I am really struggling to understand the details and I do not feel like I would be able to reproduce the results or understand the nuances of the architecture. I am aware that the code is released, but this should not be a replacement for a clearly written paper. For this reason I suggest a score of marginally below the acceptance threshold. In the main review I have highlighted the issues and points were I got stuck.