Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Mohammed Haroon Dupty, Yanfei Dong, Wee Sun Lee

TL;DR

This work tackles the expressivity gap of $1$-WL-based GNNs by embedding exact isomorphism solver ideas into a differentiable learning framework. PF-GNN builds a probabilistic representation over IR colorings via Sequential Monte Carlo, approximating a universal graph representation $f(oldsymbol{G}) = hoig(ig extstyle abla oldsymbol{ ext{(path embeddings)}}ig)$ and learning end-to-end through a learned node-selection policy, differentiable resampling, and a mean-particle readout. The method achieves permutation-invariant, unique-identification capabilities under universal approximators and shows strong empirical performance on synthetic isomorphism benchmarks and real-world datasets, often outperforming 1-WL and competing with higher-order models. Its practical impact lies in providing a flexible, backbone-agnostic framework that enhances GNN expressivity with linear runtime growth and parallelizable path sampling, enabling robust graph understanding across diverse tasks.

Abstract

Message passing Graph Neural Networks (GNNs) are known to be limited in expressive power by the 1-WL color-refinement test for graph isomorphism. Other more expressive models either are computationally expensive or need preprocessing to extract structural features from the graph. In this work, we propose to make GNNs universal by guiding the learning process with exact isomorphism solver techniques which operate on the paradigm of Individualization and Refinement (IR), a method to artificially introduce asymmetry and further refine the coloring when 1-WL stops. Isomorphism solvers generate a search tree of colorings whose leaves uniquely identify the graph. However, the tree grows exponentially large and needs hand-crafted pruning techniques which are not desirable from a learning perspective. We take a probabilistic view and approximate the search tree of colorings (i.e. embeddings) by sampling multiple paths from root to leaves of the search tree. To learn more discriminative representations, we guide the sampling process with particle filter updates, a principled approach for sequential state estimation. Our algorithm is end-to-end differentiable, can be applied with any GNN as backbone and learns richer graph representations with only linear increase in runtime. Experimental evaluation shows that our approach consistently outperforms leading GNN models on both synthetic benchmarks for isomorphism detection as well as real-world datasets.

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

TL;DR

This work tackles the expressivity gap of -WL-based GNNs by embedding exact isomorphism solver ideas into a differentiable learning framework. PF-GNN builds a probabilistic representation over IR colorings via Sequential Monte Carlo, approximating a universal graph representation and learning end-to-end through a learned node-selection policy, differentiable resampling, and a mean-particle readout. The method achieves permutation-invariant, unique-identification capabilities under universal approximators and shows strong empirical performance on synthetic isomorphism benchmarks and real-world datasets, often outperforming 1-WL and competing with higher-order models. Its practical impact lies in providing a flexible, backbone-agnostic framework that enhances GNN expressivity with linear runtime growth and parallelizable path sampling, enabling robust graph understanding across diverse tasks.

Abstract

Message passing Graph Neural Networks (GNNs) are known to be limited in expressive power by the 1-WL color-refinement test for graph isomorphism. Other more expressive models either are computationally expensive or need preprocessing to extract structural features from the graph. In this work, we propose to make GNNs universal by guiding the learning process with exact isomorphism solver techniques which operate on the paradigm of Individualization and Refinement (IR), a method to artificially introduce asymmetry and further refine the coloring when 1-WL stops. Isomorphism solvers generate a search tree of colorings whose leaves uniquely identify the graph. However, the tree grows exponentially large and needs hand-crafted pruning techniques which are not desirable from a learning perspective. We take a probabilistic view and approximate the search tree of colorings (i.e. embeddings) by sampling multiple paths from root to leaves of the search tree. To learn more discriminative representations, we guide the sampling process with particle filter updates, a principled approach for sequential state estimation. Our algorithm is end-to-end differentiable, can be applied with any GNN as backbone and learns richer graph representations with only linear increase in runtime. Experimental evaluation shows that our approach consistently outperforms leading GNN models on both synthetic benchmarks for isomorphism detection as well as real-world datasets.
Paper Structure (27 sections, 2 theorems, 26 equations, 5 figures, 17 tables, 1 algorithm)

This paper contains 27 sections, 2 theorems, 26 equations, 5 figures, 17 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. Search Tree of colorings
  5. Proposed Method
  6. Probabilistic representation
  7. Particle Filter Graph Neural Network
  8. Experiments
  9. Graph isomorphism detection
  10. Graph properties:
  11. Real world benchmarks
  12. Ablation studies
  13. Conclusion
  14. Appendix
  15. Particle filter algorithm
  16. ...and 12 more sections

Key Result

Theorem 1

Consider any $n$-vertex graph ${\mathcal{G}}$ with no distinct node attributes. Assume we use universal multiset function approximators for target-cell selection and graph pooling function $\psi$, and a GNN with 1-WL equivalent expressive power for color refinement, then the representation $f({\math

Figures (5)

  • Figure 1: Two 1-WL equivalent graphs with different colorings after one step of individualization and refinement.
  • Figure 2: An example search tree of colorings generated by exact graph isomorphism solvers. Initial coloring is produced by 1-WL refinement. PF-GNN approximates the tree by sampling multiples paths from root to leaf.
  • Figure 3: Runtime ratio of PF-GNN on SR25 dataset.
  • Figure 4: Example equitable partitions. The middle coloring in the second row is both equitable and discrete coloring.
  • Figure 5: IR process: First a vertex is distinguished from rest of the vertices and this information is propagated to other nodes via message passing to generate a refined equitable coloring of the graph.

Theorems & Definitions (4)

  • Theorem 1
  • Theorem 2
  • proof
  • proof