Symmetry-driven graph neural networks

1 INTRODUCTION . Symmetries exist throughout nature . All the fundamental laws of physics are built upon the framework of symmetries , from the gauge groups describing the Standard Model of particle physics , to Einstein ’ s theories of general and special relativity . Once one understands the symmetry of a certain system , powerful predictions can be made . A notable example is that of Gell-Mann ’ s eightfoldway ( Gell-Mann , 1961 ) , built upon the symmetries observed in hadrons , that led to his prediction of the Ω− baryon , which was subsequently observed 3 years later ( Barnes et al. , 1964 ) . The study of symmetries and invariance in deep learning has recently become a field of interest to the community ( see , e.g. , ( Bronstein et al. , 2021 ) for a comprehensive overview ) , and rapid progress has been made in constructing architectures with group theoretic structures embedded within . Two fundamental architectures in machine learning , the convolutional and graph neural networks , are invariant to the translation and permutation groups respectively . Graph networks in particular are designed to learn from graph-structured data and are by construction invariant to permutations of the input nodes . They were originally proposed in ( Gori et al. , 2005 ; Scarselli et al. , 2008 ) and have received a lot of attention in the last years ( see , e.g. , ( Battaglia et al. , 2018 ; Hamilton , 2020 ; Wu et al. , 2020 ) for a comprehensive overview ) . Due to their properties , they find application in a broad range of problems like learning the dynamics of complex physical systems ( Sanchez-Gonzalez et al. , 2020 ; Pfaff et al. , 2021 ) , particle identification in particle physics ( Dreyer & Qu , 2021 ) , learning causal and relational graphs ( Kipf et al. , 2018 ; Li et al. , 2020 ) , discovering symbolic models ( Cranmer et al. , 2020 ) , as well as quantum chemistry ( Gilmer et al. , 2017 ) and drug discovery ( Stokes et al. , 2020 ) . In this work , we decouple node coordinates from other node attributes to obtain invariance ( and possibly equivariance ) to many transformations affecting the node coordinates , possibly belonging to important groups , in addition to permutation invariance . First , we define the distance preserving graph network ( DGN ) whose updated node , edge and global attributes are invariant to any transformation in the coordinate embeddings that preserves the distance between neighbouring nodes , while node coordinates can be updated in an equivariant way . Examples of such transformations are rotations and translations , but also transformations on the Hoberman sphere , whose 3D shape changes while preserving distances between connected nodes . Equivariance to dilations in the coordinates can be obtained via the conformal orthogonal group with the inclusion of an additional input layer . Then , we define the angle preserving graph network ( AGN ) , whose updated node , edge and global attributes ( resp . node coordinates ) are invariant ( resp . equivariant ) to any transformation preserving the angles between triples of neighbouring nodes in the graph ( a notable example being molecular conformations ) and , thus , to the n-dimensional conformal group on Rn . By constructing such architectures , we enable a wide range of possible transformations to be performed on graph-structured data , with the only requirement being that distances or angles between coordinates of neighbouring nodes are preserved . This means that the updated attributes/coordinates of both networks are invariant/equivariant to translations , rotations , reflections ( i.e. , the Euclidean group , E ( n ) ) , with the AGN allowing also invariance/equivariance to dilations , inversions and nonorthogonal rotations . In practice , this means the networks are able to filter out many copies of the same input whose coordinate embeddings have been transformed and therefore learn more efficiently than architectures which consider the inputs as distinct . In other words , a single sample contains the same information as many copies of it obtained by appropriately transforming it . Finally , while the two architectures we present are partially overlapping in terms of symmetries , it is important to consider specific use cases where the use of one architecture over another would be preferred . We test the architectures on a synthetic dataset consisting of n-dimensional geometric shapes and other benchmark datasets . Moreover , we show examples of cases in which using our architectures is counterproductive due to the lack of symmetries in the data . 2 BACKGROUND . 2.1 GROUP THEORY AND EQUIVARIANCE . Group theory provides the mathematical formulation for symmetries of systems , with symmetry operations represented by individual group elements . A group can be defined as a set G equipped with a binary operation , which enables one to combine two group elements to form a third , whilst preserving the group axioms ( associativity , identity , closure and inverse ) . A function unaffected by a group action is said to be invariant , which is particular instance of equivariance . Formally , let X ⊆ Rn and ϕg : X → X be a transformation on X for a group element g ∈ G. Then , the linear map Φ : X → Y , Y ⊆ Rn , is equivariant to G if ∀g ∈ G , ∃ϕ′g : Y → Y such that Φ ( ϕg ( x ) ) = ϕ ′ g ( Φ ( x ) ) , ∀ϕg : X → X . When ϕ′g is the identity , we say that Φ is invariant to G. The functional form of Φ is determined by the specific group of interest . It is worth pointing out that , whilst not often expressed in group theoretic notation , equivariance exists in common deep learning architectures ; the convolution operation used in CNNs is equivariant under the translation group and approximate invariance is typically achieved via pooling operations . Graph neural networks are invariant to the permutation group . Euclidean group The translation and orthogonal groups ( T ( n ) andO ( n ) respectively ) act through translations , orthogonal rotations and reflections . The semidirect product of these two groups E ( n ) = T ( n ) o O ( n ) is known as the Euclidean group , which consists of transformations that preserves distances . A function that is invariant under E ( n ) is ‖xi−xj‖22 , for any xi , xj ∈ En ( see Appendix A for a proof ) . Conformal group The conformal group Conf ( Rn,0 ) is the group of transformations ϕ : Rn → Rn that preserve angles . Formally , for any triple of vectors xj , xi , xk ∈ Rn , let us denote by ∠ ( xj , xi , xk ) the angle centred on xi with rays given by xj − xi and xk − xi . Then , a conformal transformation ϕ satisfies ∠ ( xj , xi , xk ) → ∠ ( ϕ ( xj ) , ϕ ( xi ) , ϕ ( xk ) ) = ∠ ( xj , xi , xk ) . ( 1 ) By definition , the transformations of the conformal group include translations , rotations , reflections ( which collectively form the Euclidean group E ( n ) , as well as dilations , inversions and other features . An important subgroup is the conformal orthogonal group , which requires transformations to be orthogonal ( and thus , inversions are not allowed ) . 2.2 GRAPHS AND COORDINATE EMBEDDINGS . A graph is defined as G ( V , E ) where V = { 1 , . . . , N } is the set of nodes and E = { ( j , i ) } ⊆ V × V is the set of ( directed ) edges connecting nodes in V ( where j and i denote the source and target nodes respectively ) . We define Ni , { j | ( j , i ) ∈ E } as the set of ( in- ) neighbours of node i . Moreover , we associate node and edge features , vi ∈ Rnv , ∀i ∈ V , and eji ∈ Rne , ∀ ( j , i ) ∈ E respectively , to each node and edge in the graph as well as a global attribute u ∈ Rnu . We assume that node features vi consist of node coordinates xi ∈ Rnx and additional features hi ∈ Rnh ( unrelated to coordinates ) , so that vi = [ xi , hi ] , nv = nx + nh . With a slight abuse of notation , we denote GX ( V , E ) and GY ( V , E ) as different coordinate embeddings of the same graph , meaning that they differ only in their node coordinates xi , while having the same node , edge and global attributes hi , eji and u . Finally , let A ∈ V ×V ×V be the set of ( ordered ) triples of nodes in a graph G that form an angle , i.e. , A , { ( j , i , k ) | j , k ∈ N ui , j 6= k , ∀i ∈ V } with N ui , { j | ( j , i ) ∈ E ∨ ( i , j ) ∈ E } being the set of in- and out-neighbors of node i . We define relative distance and angle preserving maps as follows ( see Figure 1 ) . Definition . Let GX ( V , E ) and GY ( V , E ) be such that ‖xi − xj‖2 = ‖yi − yj‖2 , ∀ ( i , j ) ∈ E . Let x+i = ψ ( i , GX ) and y+i = ψ ( i , GY ) for some function ψ : RK → Rn , K ∈ Z+ . We say that ψ is a relative distance preserving map if ‖x+i − x+j ‖2 = ‖y+i − y+j ‖2 , ∀ ( i , j ) ∈ E . Definition . Let GX ( V , E ) and GY ( V , E ) be such that ∠ ( xj , xi , xk ) = ∠ ( yj , yi , yk ) ∀ ( j , i , k ) ∈ A . Let x+i = ψ ( i , GX ) and y+i = ψ ( i , GY ) for some function ψ : RK → Rn , K ∈ Z+ . We say that ψ is a relative angle preserving map if ∠ ( x+j , x + i , x + k ) = ∠ ( y + j , y + i , y + k ) , ∀ ( j , i , k ) ∈ A . The simplest map satisfying the above definitions is the identity function x+i = ψ ( i , GX ) = xi ( 2 ) which keeps the coordinates unchanged during the update step . Updating all the coordinates in the same way also trivially preserves both relative distances and angles . In the case where the coordinate embedding X is a Conformal orthogonal transformation of Y a possible map ψ is defined by x+i = xi + ∑ j∈Ni aji ( xj − xi ) . ( 3 ) with aji possibly being a parametric function of other graph attributes . A thorough discussion and proofs are reported in in Appendix C .

Summary: This paper proposes two kinds of equivariant GNNs, including the distance preserving graph network (DGN) that is equivariant to the Euclidean transformations and angle preserving graph network (AGN) that is equivariant to the Conformal group. The proposed frameworks generalize several typical previous architectures. The advantage is experimentally evaluated on a synthetic dataset composed of n-dimensional geometric objects.