Towards Sparse Hierarchical Graph Classifiers

Introduction and Related Work

Here we study the problem of graph classification; the task of learning to categorise graphs into classes. This is a direct generalisation of image classification , as images may be easily cast as a special case of a “grid graph” (with each pixel of an image connected to its eight immediate neighbours). Therefore, it is natural to investigate and generalise CNN elements to graphs .

Generalising the convolutional layer to graphs has been a very active area of research, with several graph convolutional layers proposed in recent times, significantly advancing the state-of-the-art on many challenging node classification benchmarks (analogues of image segmentation in the graph domain), as well as link prediction. Conversely, generalising pooling layers has received substantially smaller levels of attention by the community.

The proposed strategies broadly fall into two categories: 1) aggregating node representations in a global pooling step after each or after the final message passing step , and 2) aggregating node representations into clusters which coarsen the graph in a hierarchical manner . Apart from , all earlier works in this area assume a fixed, pre-defined cluster assignment, that is obtained by running a clustering algorithm on the graph nodes, e.g. using the GraClus algorithm to obtain structure-dependent cluster assignments or finding clusters via k-means on node features . The main insight by recent works is that intermediate node representations (e.g. after applying a graph convolution layer) can be leveraged to obtain both feature- and structure-based cluster assignments that are adaptive to the underlying data and that can be learned in a differentiable manner.

The first end-to-end trainable graph CNN with a learnable pooling operator was recently pioneered, leveraging the DiffPool layer . DiffPool computes soft clustering assignments of nodes from the original graph to nodes in the pooled graph. Through a combination of restricting the clustering scores to respect the input graph’s adjacency information, and a sparsity-inducing entropy regulariser, the clustering learnt by DiffPool eventually converges to an almost-hard clustering with interpretable structure, and leads to state-of-the-art results on several graph classification benchmarks.

The main limitation of DiffPool is computation of the soft clustering assignments—while the assignments eventually converge, during the early phases of training, an entire assignment matrix must be stored; relating nodes from the original graph to nodes from the pooled graph in an all-pairs fashion. This incurs a quadratic $\mathcal{O}(kV^{2})$ storage complexity for any pooling scheme with a fixed pooling ratio $k$, and is therefore prohibitive for large graphs.

In this work, we leverage recent advances in graph neural network design to demonstrate that sparsity need not be sacrificed to obtain good performance on end-to-end graph convolutional architectures with pooling. We demonstrate performance that is comparable to variants of DiffPool on four standard graph classification benchmarks, all while using a graph CNN that only requires $\mathcal{O}(V+E)$ storage (comparable to the storage complexity of the input graph).

Model

To specify a CNN-inspired neural network for graph classification, we first require a convolutional and a pooling layer. In addition, we require a readout layer (analogous to a flattening layer in an image CNN), that converts the learnt representations into a fixed-size vector representation, to be used for final prediction (e.g. a simple MLP). These layers are specified in the following paragraphs.

Given that our model will be required to classify unseen graph structures at test time, the main requirement of the convolutional layer in our architecture is that it is inductive, i.e. that it does not depend on a fixed and known graph structure. The simplest such layer is the mean-pooling propagation rule, as similarly used in GCN or Const-GAT :

Pooling layer

To make sure that a graph downsampling layer behaves idiomatically with respect to a wide class of graph sizes and structures, we adopt the approach of reducing the graph with a pooling ratio, $k\in(0,1]$. This implies that a graph with $N$ nodes will have $\lceil kN\rceil$ nodes after application of such a pooling layer.

Unlike DiffPool, which attempts to do this via computing a clustering of the $N$ nodes into $\lceil kN\rceil$ clusters (and therefore incurs a quadratic penalty in storing cluster assignment scores), we leverage the recently proposed Graph U-Net architecture , which simply drops $N-\lceil kN\rceil$ nodes from the original graph.

The choice of which nodes to drop is done based on a projection score against a learnable vector, $\vec{p}$. In order to enable gradients to flow into $\vec{p}$, the projection scores are also used as gating values, such that retained nodes receiving lower scores will experience less significant feature retention. Fully written out, the operation of this pooling layer (computing a pooled graph, $\left({\bf X^{\prime}},{\bf A^{\prime}}\right)$, from an input graph, $\left({\bf X},{\bf A}\right)$), may be expressed as follows:

Here, $\|\cdot\|$ is the $L_{2}$ norm, top-$k$ selects the top-$k$ indices from a given input vector, $\odot$ is (broadcasted) elementwise multiplication, and $\cdot_{\vec{i}}$ is an indexing operation which takes slices at indices specified by $\vec{i}$. This operation requires only a pointwise projection operation and slicing into the original feature and adjacency matrices, and therefore trivially retains sparsity.

Readout layer

Lastly, we seek a “flattening” operation that will preserve information about the input graph in a fixed-size representation. A natural way to do this in CNNs is global average pooling, i.e. the average of all learnt node embeddings in the final layer. We further augment this by performing global max pooling as well, which we found strengthened our representations. Lastly, inspired by the JK-net architecture , we perform this summarisation after each conv-pool block of the network, and aggregate all of the summaries together by taking their sum.

Concretely, to summarise the output graph of the $l$-th conv-pool block, $({\bf X}^{(l)},{\bf A}^{(l)})$:

where $N^{(l)}$ is the number of nodes of the graph, $\vec{x}_{i}^{(l)}$ are the $i$-th node’s feature vector, and $\|$ denotes concatenation. Then, the final summary vector (for a graph CNN with $L$ layers) is obtained as the sum of all those summaries (i.e. $\vec{s}=\sum_{l=1}^{L}\vec{s}^{(l)}$) and submitted to an MLP for obtaining final predictions.

We find that the aggregation across layers is important, not only to preserve information at different scales of processing, but also to handle efficiently retaining information on smaller input graphs that may quickly be pooled down to a too small number of nodes.

The entire pipeline of our model may be visualised in Figure 1.

Experiments

To assess how well our sparse model can hierarchically compress the representation of a graph while still producing features relevant for classification, we evaluate the graph neural network architecture on several well-known benchmark tasks: biological (Enzymes, Proteins, D&D) and scientific collaboration (Collab) datasets . We report the performance achieved from carrying out 10-fold cross-validation on each of these, in relation to the results presented by Ying et al. .

Model parameters

Our graph neural network architecture comprises three blocks, each of them consisting of a graph convolutional layer with 128 (Enzymes and Collab) or 64 features (D&D and Proteins), followed by a pooling step (refer to Section 2 for details). We ensure that there is enough information after each coarsening stage by preserving 80% of the existing nodes. A learning rate of 0.005 was used for Proteins and 0.0005 for all other datasets. The model was trained using the Adam optimizer for 100 epochs on Enzymes, 40 on Proteins, 20 on D&D and 30 on Collab.

Results

Table LABEL:table:results illustrates our comparison to the performances reported by Ying et al. . In all cases, our algorithm significantly outperforms the GraphSAGE sparse aggregation method , while successfully competing at most within 1 percentage point of accuracy with the three variants of DiffPool , the recent singular development in hierarchical graph representation learning. Unlike the latter, our method does not require quadratic memory, paving the way to deploying scalable hierarchical graph classification algorithms on larger real-world datasets.

We also verify this claim empirically—through experiments on random inputs—in Figure 2, where we demonstrate that our method compares favourably to DiffPool on larger-scale graphs, even if the pooling layer doesn’t drop any nodes (compared to a 0.25 retain rate for the DiffPool).

Acknowledgements

We would like to thank the developers of PyTorch . CC acknowledges funding by DREAM CDT. PV and PL have received funding from the European Union’s Horizon 2020 research and innovation programme PROPAG-AGEING under grant agreement No 634821. TK acknowledges funding by SAP SE. We specially thank Jian Tang and Max Welling for the extremely useful discussions.

References

Appendix A Qualitative analysis

We qualitatively investigate the distribution of graph summaries, using a pre-trained model on a fold of the Collab dataset to produce 499 outputs across all 3 classes. Figure 3 shows that an evident clustering can be achieved, once the graph has been processed by the sequence of convolution and pooling layers leveraged by our architecture.