ASAP: Adaptive Structure Aware Pooling for Learning Hierarchical Graph Representations

TL;DR

ASAP introduces a novel, adaptive, and structure-aware pooling method for GNNs that improves graph classification performance by effectively capturing substructure and scaling to large graphs.

Contribution

The paper proposes ASAP, a new sparse, differentiable pooling method with self-attention and modified GNNs, addressing limitations of previous pooling techniques.

Findings

Achieves state-of-the-art results on multiple graph classification benchmarks.

Improves performance by an average of 4% over existing sparse hierarchical methods.

Effectively captures graph substructure and scales to large graphs.

Abstract

Graph Neural Networks (GNN) have been shown to work effectively for modeling graph structured data to solve tasks such as node classification, link prediction and graph classification. There has been some recent progress in defining the notion of pooling in graphs whereby the model tries to generate a graph level representation by downsampling and summarizing the information present in the nodes. Existing pooling methods either fail to effectively capture the graph substructure or do not easily scale to large graphs. In this work, we propose ASAP (Adaptive Structure Aware Pooling), a sparse and differentiable pooling method that addresses the limitations of previous graph pooling architectures. ASAP utilizes a novel self-attention network along with a modified GNN formulation to capture the importance of each node in a given graph. It also learns a sparse soft cluster assignment for nodes at each layer to effectively pool the subgraphs to form the pooled graph. Through extensive experiments on multiple datasets and theoretical analysis, we motivate our choice of the components used in ASAP. Our experimental results show that combining existing GNN architectures with ASAP leads to state-of-the-art results on multiple graph classification benchmarks. ASAP has an average improvement of 4%, compared to current sparse hierarchical state-of-the-art method.