Local Permutation Equivariance For Graph Neural Networks

1 INTRODUCTION . Many forms of data are naturally structured as graphs such as molecules , bioinformatics , social , or financial and it is therefore of interest to have algorithms which operate over graphs . Machine learning on graphs has received much interest in recent years , with the general framework of a message passing network providing both a useful inductive bias and scalability across a range of domains ( Gilmer et al. , 2017 ) . However , Xu et al . ( 2019 ) show that a model based on a message passing framework with permutation invariant aggregation functions is limited in expressive power . Therefore , there exists many non isomorphic graphs that a model of this form can not distinguish between . Figure 5 demonstrates two non-isomorphic graphs for which a message passing framework with max pooling would not be able to distinguish between the two graphs . More expressive graph networks exists and a common measure of expressivity is the WeisfeilerLehman ( WL ) test . One such method of building more expressive networks is based directly on the WL test . Here a neural network architecture is built based on variants of the WL test . Bouritsas et al . ( 2020 ) use a permutation invariant local update function , but incorporate permutation equivariant structural information into the feature space . Morris et al . ( 2019a ) build models based on different WL variants that consider local and global connections . Bodnar et al . ( 2021b ) introduce a WL test on simplicial complexes and incorporate this into a message passing scheme . Bodnar et al . ( 2021a ) extend work on simplicial complexes to cell complexes , which subsume simplicial complexes . On the other hand , rather than trying to directly incorporate techniques from WL tests directly into networks other model make use of permutation symmetries to build permutation equivariant graph neural networks ( Maron et al. , 2018 ) . This model can be built for k-order feature spaces and it was shown by Maron et al . ( 2019 ) that such models can distinguish between non-isomorphic graphs as well as the k-WL test . Natural graph networks are a different class of graph neural network , where the constraint placed upon the linear layer is that of naturality ( de Haan et al. , 2020 ) . The naturality constraint says that for each isomorphism class a map must be chosen that is equivariant to automorphisms . In general the task of learning on graphs consists of utilising many graphs of different sizes . Current methods for utilising permutation equivariant graph neural networks require that the graph be represented as an adjacency tensor , which limits there scalability . Furthermore , global natural graph networks also perform computations on entire graph features , which leads to a large computational complexity for large graphs . Local gauge symmetries have been considered to build models with local equivariance ( Cohen et al. , 2019 ) . This approach improves scalability of models by utilising local update functions , however for graphs we do not have a single local symmetry . Currently this is overcome in the majority of graph neural networks presented by utilising some form of message passing , but , in general , all works use a permutation invariant aggregation function leading to good scalability but poor expressivity . Local natural graph networks attempt to overcome the limited expressivity through placing a local naturality constraint on the message passing and having different message passing kernels on non-isomorphic edges . Through considering graph neural networks from an elementary category theory perspective and making use of aspects of group theory we present a framework for building local permutation equivariant models . This allows us to build a graph neural network model with local update functions that are permutation equivariant by considering restricted representations of the representation space of the whole graph . Further , we prove that this does not cause a loss of expressivity of the model and that this maintains the option to have a k-order feature space that ensures expressivity equal to k-WL test . Also , by constraining the kernel space under restricted representations , a natural weight sharing scheme becomes apparent , namely sharing weights across local graph neighbourhoods of the same degree . The approach of building models with a framework based on group theory makes clear the generality of the approach , where choices of representation space can be made for each convolutional layer without requiring prior information such as structural information to be encoded into the feature space . This framework can also be shown to include other leading methods as specific cases . 2 BACKGROUND . 2.1 GRAPH NETWORKS . Different graph neural networks express graphs in alternative forms . Generally , for a message passing model , a matrix of node features and a matrix of edge features is combined with a sparse edge index array specifying the connectivity of the graph . In other works , the graph is provided in a dense format , where the graph is given as a adjacency tensor with node and edge features held in one tensor . In this work we present the graph as follows : Definition 1 A Concrete Graph G is a finite set of nodes V ( G ) ⊂ N and a set of edges E ( G ) ⊂ V ( G ) × V ( G ) . The set of node ids may be non-contiguous and we make use of this here as we extract overlapping sub-graphs when performing the local updates . The same underlying graph can be given in may forms by a permutation of the ordering of the natural numbers of the nodes . Definition 2 A sub-Concrete Graph H is created by taking a node i ∈ V ( G ) , and extracting the nodes j ∈ V ( G ) and edges ( i , j ) ⊂ V ( G ) × V ( G ) , such that there is a connection between nodes i and j . Once a sub-concrete graph has been extracted , this same underlying sub-graph could be expressed through different permutations of the underlying numbering of the nodes . For brevity we refer to sub-concrete graphs as subgraphs throughout the paper . Definition 3 A Graph isomorphism , φ : G→ G′ is a bijection between the vertex sets of two graphs G andG′ , such that two vertices u and v are adjacent inG if and only if φ ( u ) and φ ( v ) are adjacent in G′ . This mapping is edge preserving , i.e . satisfies for all ( i , j ) ∈ V ( G ) × V ( G ) : ( i , j ) ∈ E ( G ) ⇐⇒ ( φ ( i ) , φ ( j ) ) ∈ E ( G′ ) An isomorphism from the graph to itself is known as an automorphism .. Relabelling of the graph by a permutation of the nodes is called a graph isomorphism , where an example of two isomorphic graphs is given in Figure 5 . We desire that the linear layers of the graph neural network respect the composition of graph isomorphisms . This requires us to define the feature space of the graphs and how feature spaces of isomorphic graphs are related . 2.2 PERMUTATION REPRESENTATIONS . The feature space of the graphs is a vector space V , where a representation of the group G is a homomorphism ρ : G → GL ( V ) of G to the group of automorphisms of V ( Fulton & Harris , 2013 ) . A map KG between two representations of G is a vector space map . The elements of the group g ∈ G can act on a vector v ∈ V by the representation matrix v → ρ ( g ) v. The symmetric subspace of the representation is the space of solutions to the constraint ∀g ∈ G : ρ ( g ) v = v. Here we are considering the symmetries of the symmetric group Sn . This constraint can be solved for different order representations ( Maron et al. , 2018 ; Finzi et al. , 2021 ) . We present the space of linear layers mapping from k-order representations to k′-order representations in Figure 2 . In addition , for the linear map KG , we require that if a graph is passed through KG and then transformed by permutation to an isomorphic graph this result is the same as if a graph is transformed by the same permutation to an isomorphic graph and then passed through KG . In short , this requires that permutation equivariance is satisfied . 2.3 CATEGORY THEORY . This section does not provide a complete overview of category theory , nor even a full introduction , but aims to provide a sufficient level of understanding to aid the reader with further sections of the paper , where we believe presenting the comparison between models from a category theory perspective makes more clear the distinctions between them . A category , C , consists of a set of objects , Ob ( C ) , and a set of morphisms ( structure-preserving mappings ) or arrows , f : A → B , A , B ∈ Ob ( C ) . There is a binary operation on morphisms called composition . Each object has an identity morphism . Categories can be constructed from given ones by constructing a subcategory , in which each object , morphism , and identity is from the original category , or by building upon a category , where objects , morphisms , and identities are inherited from the original category . A functor is a mapping from one category to another that preserves the categorical structure . For two categories C and D a functor F : C → D maps each object A ∈ Ob ( C ) to an object F ( A ) ∈ Ob ( D ) and maps each morphism f : A→ B in C to a morphism F ( f ) : F ( A ) → F ( B ) in D. Definition 4 A groupoid is a category in which each morphism is invertible . A groupoid where there is only one object is usually a group . 3 GLOBAL EQUIVARIANT GRAPH NETWORKS . 3.1 GLOBAL PERMUTATION EQUIVARIANCE . Global permutation equivariant models have been considered by Hartford et al . ( 2018 ) ; Maron et al . ( 2018 ; 2019 ) ; Albooyeh et al . ( 2019 ) , with Maron et al . ( 2018 ) demonstrating that for order-2 layers there are 15 operations that span the full basis for an permutation equivariant linear layer . These 15 basis elements are shown in Figure 2 with each basis element given by a different color in the map from representation ρ2 → ρ2 . Despite these methods , when solved for the entire basis space , having expressivity as good as the k-WL test , they operate on the entire graph . Operating on the entire graph features limits the scalability of the methods . In addition to poor scalability , global permutation appears to be a strong constraint to place upon the model . In the instance where the graph is flattened and an MLP is used to update node and edge features the model would have n4 trainable parameters , where n is the number of nodes . On the other hand , a permutation equivariant update has only 15 trainable parameters and in general 15 n4 . Viewing a global permutation equivariant graph network from a category theory perspective there is one object with a collection of arrows representing the elements of the group . Here the arrows or morphisms go both from and to this same single object . The feature space is a functor which maps from a group representation to a vector space . For a global permutation equivariant model the same map is used for every graph . G e g1 g2 Symmetric Group

Summary: Graph Neural Networks have recently become the state of the art for tasks on graphs due to their flexibility, scalability. In this work, the authors propose a framework to build GNN's that operate on local node neighborhoods in a permutation equivariant way - and argue that since LPEGN operates on lower dimensional spaces in comparison to regular GNN's the proposed technique offers significant improvements  in terms of GPU memory usage. The authors make use of category theory basics - and employ restricted representations of finite symmetric groups (i.e. fix some nodes while permuting other elements) based on the number of nodes in the neighborhood of the node - and ensure there is weight sharing between nodes with the same degree to achieve their objective.