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: The paper introduces the framework of locally permutation equivariant graph neural networks.  This framework applies permutation equivariant layers [Maron et. al. 2018] to local node neighborhoods by treating them as separate subgraphs and using a weight sharing scheme for subgraphs of the same size.  The authors build their framework by discussing the different choices made -- local neighborhoods, weight sharing, and representation space.  The authors also provide a category theory point of view of their framework.