Hi-PART: Going Beyond Graph Pooling with Hierarchical Partition Tree for Graph-Level Representation Learning

Abstract

Graph pooling refers to the operation that maps a set of node representations into a compact form for graph-level representation learning. However, existing graph pooling methods are limited by the power of the Weisfeiler–Lehman (WL) test in the performance of graph discrimination. In addition, these methods often suffer from hard adaptability to hyper-parameters and training instability. To address these issues, we propose Hi-PART, a simple yet effective graph neural network (GNN) framework with Hi erarchical Par tition T ree (HPT). In HPT, each layer is a partition of the graph with different levels of granularities that are going toward a finer grain from top to bottom. Such an exquisite structure allows us to quantify the graph structure information contained in HPT with the aid of structural information theory. Algorithmically, by employing GNNs to summarize node features into the graph feature based on HPT’s hierarchical structure, Hi-PART is able to adequately leverage the graph structure information and provably goes beyond the power of the WL test. Due to the separation of HPT optimization from graph representation learning, Hi-PART involves the height of HPT as the only extra hyper-parameter and enjoys higher training stability. Empirical results on graph classification benchmarks validate the superior expressive power and generalization ability of Hi-PART compared with state-of-the-art graph pooling approaches.