PF-GNN: Differentiable particle filtering based approximation of universal graph representations

1 INTRODUCTION . In recent years , Graph Neural Networks ( GNNs ) have emerged as learning models of choice on graph structured data . GNNs operate on a message passing paradigm ( Kipf & Welling , 2016 ; Defferrard et al. , 2016 ; Veličković et al. , 2017 ; Gilmer et al. , 2017 ) , where nodes maintain latent embeddings which are updated iteratively based on their neighborhood . This way of representation learning on graphs provides the necessary inductive bias to encode the structural information of graphs into the node embeddings . The process of message passing in GNNs is equivalent to vertex color-refinement procedure or the 1-dimensional Weisfeiler-Lehman ( WL ) test used to distinguish non-isomorphic graphs ( Xu et al. , 2018 ; Morris et al. , 2019 ) . Consequently , GNNs suffer from the limitations of 1-WL color-refinement in their expressive power . In each step of 1-WL color-refinement , two vertices get different colors if the colors of neighboring vertices are different . The procedure stabilizes after a few steps when the colors can not be further refined . Due to the symmetry in graph structures , certain non-isomorphic graphs induce same colors upon 1-WL refinement . Higher-order WL refinement and their neural k-GNN versions break some of the symmetry by operating on k-tuples of nodes . They are more expressive but require exponentially increasing computation time and hence , are not practical for large k. Motivated by the fact that a fully expressive graph representation learning model should be able to produce embeddings that can distinguish any two non-isomorphic graphs ( Chen et al. , 2019b ) , we turn to exact graph isomorphism solvers for better inductive biases in our learning algorithm . Most of the practical graph isomorphism solvers use 1-WL in combination with the traditional technique of individualization and refinement ( IR ) ( McKay & Piperno , 2014 ; Junttila & Kaski , 2011 ) for coloring the graph . Individualization is the process of artificially introducing asymmetry by recoloring a vertex and thereby , distinguishing it from the rest of the vertices . Refinement refers to 1-WL refinement which can propagate this information to recolor the rest of the graph . The two graphs shown in Fig . 1 are not distinguishable after 1-WL refinement but induce different colorings after one IR step . The IR process is repeated for each refinement until every vertex gets a unique color . However , in order to maintain permutation-invariance , whenever a vertex is individualized , other vertices that have the same color need to individualized as well and thereafter refined . This process generates a search tree with colorings as nodes , and can grow exponentially in worst case . In this work , we propose to learn graph representations with the inductive bias of the search tree generated by the graph-isomorphism solvers . However , generating the complete search tree is computationally expensive . Isomorphism-solvers prune the search tree by detecting automorphisms on the fly as they generate the tree . Nevertheless , detecting automorphisms is non-trivial from the perspective of end-to-end discriminative learning and hence , we need to approximate it . To this end , we first define a universal neural graph representation based on the search tree of colorings . Then we take a probabilistic view and approximate it by sampling multiple paths from root to leaves of the search tree . We observe that the process of sampling a vertex , individualizing it and the subsequent refinement resembles the state transition in sequential state estimation problems.Hence , we model our sampling approach on the principled technique of Sequential Monte Carlo or Particle Filtering . We introduce Particle Filter Graph Neural Networks ( PF-GNN ) , an end-to-end learnable algorithm which maintains a weighted belief distribution over a set of K graph colorings/embeddings . With each step of IR , PF-GNN transitions to a new set of embeddings . It then updates the belief by re-weighting each particle with a discriminatively learned function that measures the quality of the refinement induced after the transition . With this inductive bias , the belief evolves over time to be more discriminative of the input graph . After a few steps , we can use the belief along with the embeddings to generate the final representation of the graph . Our approach is simple , efficient , parallelizable , easy to implement and can learn representations to distinguish non-isomorphic graphs beyond 1-WL GNNs . We evaluate PF-GNN over diverse set of datasets on tasks which are provably not learnable with 1-WL equivalent GNNs . Furthermore , our experiments on real-world benchmark datasets show the strong performance of PF-GNN over other more expressive GNNs . 2 RELATED WORK . It was established that GNNs are limited in expressive power , and can not go beyond 1-WL test for graph isomorphism by Xu et al . ( 2018 ) and Morris et al . ( 2019 ) . Later analysis of GNN has shown other limits of GNNs like counting substructures and detecting graph properties ( Arvind et al. , 2020 ; Chen et al. , 2020 ; Loukas , 2019 ; Dehmamy et al. , 2019 ; Srinivasan & Ribeiro , 2019 ) . Chen et al . ( 2019b ) further formalizes the intuition that there is equivalence between learning universal graph representations and solving the graph isomorphism problem . Thereafter , many models have been proposed that are more expressive than 1-WL . Prominent ones are k-GNNs and their equivalent models ( Morris et al. , 2019 ; Maron et al. , 2019 ; Vignac et al. , 2020 ; Morris et al. , 2020 ) , but they are difficult to scale beyond 3-WL . Other methods need to preprocess the graph to find substructures ( Bouritsas et al. , 2020 ; Li et al. , 2020 ) , which are not task-specific and may incur more computation time . Another way to improve expressivity of GNNs is to introduce randomness ( You et al. , 2019 ; Sato et al. , 2020 ; Abboud et al. , 2020 ; Zambon et al. , 2020 ) . However , adding uniform randomness interferes with the learning task , and hence these models have not shown good performance on real-world datasets . In PF-GNN , randomness is controlled as only a subset of nodes are sampled . Furthermore , since all components are learned discriminatively , the representation is tuned towards the end task . Our approach of learning with particle filter updates is inspired from recent works which make particle filters differentiable ( Karkus et al. , 2018 ; Ma et al. , 2019 ; 2020 ) . 3 PRELIMINARIES . Let Gn be the set of all graphs of n vertices with vertex set V = { 1 , 2 , . . . , n } and edge set E , Isomorphism between any two graphs G , H ∈ Gn is a bijection f : VG → VH such that ( u , v ) ∈ EG ⇐⇒ ( f ( v ) , f ( u ) ) ∈ EH . An automorphism of G is an isomorphism that maps G onto itself . One way of identifying non-isomorphic graphs is by generating unique colorings for graphs based on their structures in a permutation-invariant way and then , comparing them . A colouring of the graph G ∈ Gn is a surjective function π : V → N , which assigns each vertex to a natural number . The number of colors is denoted by |π| . The set of vertices sharing the same color form a color cell in the coloring . We denote the set of colored vertices with π = { p1 , p2 , . . . , pk } where pi is a color cell . A coloring in which every vertex gets a distinct color is called a discrete colouring i.e . |π| = n. For any two colorings π , π′ , we say that π′ is finer than or equal to π , written π′ π , if π ( v ) < π ( w ) ⇒ π′ ( v ) < π′ ( w ) for all v , w ∈ V . i.e . each cell of π′ is a subset of a cell of π , but the converse is not true . Coloring is also loosely called a partition , since it partitions V into cells . A coloring is an equitable partition when any two vertices of the same color are adjacent to the same number of vertices of each color . 1-dimensional Weisfeiler Lehman test is a simple and fast procedure to color graphs . It starts with the same initial color for all vertices . Then , it iteratively refines the coloring of the graph by mapping the the tuple of the color of a vertex and its neighbors to a distinct new color . i.e at step t , πt+1 ( v ) = HASH ( πt ( v ) , { { πt ( u ) , u ∈ N ( v ) } } ) , where { { } denotes a multiset and N ( v ) is the set of adjacent vertices of v. The algorithm terminates when π forms an equitable partition . 3.1 SEARCH TREE OF COLORINGS Equitable coloring can not be further refined due to the symmetry in the graph structure . To break the symmetry , exact isomorphism solvers employ the technique of individualization-refinement to generate further refined colorings . Individualization is the technique of breaking the symmetry in the graph by distinguishing a vertex with a new unique color . Once a vertex is individualized , 1-WL refinement is used to refine the coloring until a further refined equitable partition is reached . However , this comes at a cost . In order to maintain permutation-invariance , we have to individualize and refine all the vertices with the same color . As a result , we get as many refined colorings as the number of vertices with the chosen color . This procedure can be repeated iteratively for each of the refined coloring and the process takes shape of a rooted search tree of equitable colorings . The search tree has initial 1-WL coloring at the root . Then a non-singleton color cell called the target cell is chosen and all vertices of the target cell are individualized in parallel and thereafter , colorings are refined . The refined equitable colorings form the child nodes of the root . After sufficient repetitions , we get a set of discrete colorings at the leaves . This search tree is unique to the isomorphism class of graphs , i.e . all non-isomorphic graphs produce distinct search trees , and consequently , discrete colorings at the leaves . An example of search tree is shown in Fig . 2 . 4 PROPOSED METHOD . In this section , we first define a universal representation of any n-vertex graph , then propose a practical approximation of the representation . For any arbitrary n-vertex graph G , we aim to learn its neural representation f ( G ) that can uniquely identify G. To achieve this , we can construct a computation graph which mimics the search tree of colorings . We propose to design a neural equivalent of the search tree . To do this , we can model the target-cell selector using a learnable function that outputs scalar scores on the vertices based on their embeddings , and then individualize those which have the maximum score . Note that node embeddings in GNN are analogous to colors in 1-WL refinement . We can then adopt GNN to produce refined set of embeddings . If we have discrete embeddings after T iterations , then f ( G ) can be computed as , f ( G ) = ρ ( ∑ ∀I ψ ( G , πIT ) ) ( 1 ) where I is a sequence of vertices individualized iteratively ( I identifies a single path from root to a leaf in the search tree ) and the summation is over all such paths . πIT is the discrete coloring ( discrete embeddings ) of G after individualizing and refining with vertices in I. ψ is an multiset function approximator and ρ is an MLP . Theorem 1 presents the expressive power of f ( G ) . Theorem 1 . Consider any n-vertex graph G with no distinct node attributes . Assume we use universal multiset function approximators for target-cell selection and graph pooling function ψ , and a GNN with 1-WL equivalent expressive power for color refinement , then the representation f ( G ) in Eqn . 1 is permutation-invariant and can uniquely identify G. Proof of Theorem 1 is provided in Appendix A.3 .

Summary: This paper presents a neural version of individual-refinement (IR) architecture for improving the expressiveness of GNN in terms of isomorphism test. As IR is the dominant approach of practical graph isomorphism test, adapting IR to GNN is a novel and important idea. As IR suffers from the exponential number of branches, the paper adapts particle filtering algorithm to sample K paths to approximate the full version of IR algorithm.  Simulation and real-world datasets are used to demonstrate the improvement over base GNN.