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: The paper introduces PF-GNN, a new method through which to individualize node representations with the aim of increasing GNN expressiveness. It is inspired by graph isomorphism algorithms' individualize and refine (IR) approaches, which individualize individual nodes then refine colorings, and then aggregate all colorings across the search tree of possible individualization paths. Specifically, PF-GNN approximately emulates IR by repeating a random sampling k times and returning an aggregate coloring from each sample. For every sample, T iterations are made, such that a sample node is selected for individualization, and then representations are updated using a GNN. However, PF-GNN introduces learnable components to this mechanism, such that affinities for node selection (yielding a distribution over all possible individualization candidate nodes) are learned, rather than simple uniform sampling. Furthermore, beliefs for sampling are iteratively updated in keeping with a particle filtering approach (PF), and node individualization itself is made using a learnable function, to enable more flexibility and data-driven individualization. Finally, the model is empirically evaluated on a series of datasets and compared with existing standard GNNs and individualization techniques, e.g., random node initialization (RNI). In these experiments, PF-GNN is shown to improve model expressiveness, is more resilient than RNI for larger numbers of nodes, and achieves strong performance on real-world datasets.