[Paper Review] Multi-Level Attention Pooling for Graph Neural Networks: Unifying Graph Representations with Multiple Localities

Takeshi Itoh, Takatomi Kubo+1 more|arXiv (Cornell University)|Mar 2, 2021

Advanced Graph Neural Networks85 references29 citations

TL;DR

This paper proposes Multi-Level Attention Pooling (MLAP), a GNN architecture that enhances graph-level classification by learning and aggregating attention-based graph representations from multiple message-passing layers. By preserving layer-wise structural information and combining it via weighted or summed aggregation, MLAP improves representation discriminability and achieves state-of-the-art performance on molecular and synthetic graph datasets, particularly when combined with normalization techniques like GraphNorm.

ABSTRACT

Graph neural networks (GNNs) have been widely used to learn vector representation of graph-structured data and achieved better task performance than conventional methods. The foundation of GNNs is the message passing procedure, which propagates the information in a node to its neighbors. Since this procedure proceeds one step per layer, the range of the information propagation among nodes is small in the lower layers, and it expands toward the higher layers. Therefore, a GNN model has to be deep enough to capture global structural information in a graph. On the other hand, it is known that deep GNN models suffer from performance degradation because they lose nodes' local information, which would be essential for good model performance, through many message passing steps. In this study, we propose multi-level attention pooling (MLAP) for graph-level classification tasks, which can adapt to both local and global structural information in a graph. It has an attention pooling layer for each message passing step and computes the final graph representation by unifying the layer-wise graph representations. The MLAP architecture allows models to utilize the structural information of graphs with multiple levels of localities because it preserves layer-wise information before losing them due to oversmoothing. Results of our experiments show that the MLAP architecture improves the graph classification performance compared to the baseline architectures. In addition, analyses on the layer-wise graph representations suggest that aggregating information from multiple levels of localities indeed has the potential to improve the discriminability of learned graph representations.

Motivation & Objective

• To address the performance degradation in deep GNNs caused by oversmoothing, which erases essential local structural information.
• To improve graph representation discriminability by unifying information from multiple levels of locality in a graph.
• To develop a generalizable architecture that enhances graph-level classification without modifying the underlying GNN backbone.
• To explore whether attention-based pooling at each layer can better preserve and utilize local and global structural patterns than single-layer pooling.

Proposed method

• Introduces an attention pooling layer after each message-passing layer in a GNN to compute layer-wise graph representations.
• Uses a learnable attention mechanism to assign importance scores to nodes in each layer, enabling selective aggregation of informative substructures.
• Aggregates the layer-wise graph representations using either sum or learnable weighted pooling to produce the final graph-level representation.
• Applies the aggregation mechanism in a manner compatible with various GNN backbones (e.g., GCN, GIN, GAT) and normalization techniques like GraphNorm.
• Employs a hierarchical aggregation strategy inspired by the Jumping Knowledge network, allowing models to dynamically select relevant levels of locality.
• Designs the architecture to be modular and extensible, supporting multi-stage training and alternative aggregator functions beyond sum and weighted pooling.

Experimental results

Research questions

• RQ1 Can aggregating graph representations from multiple layers of a GNN improve classification performance compared to single-layer pooling?
• RQ2 Does preserving layer-wise representations through attention pooling mitigate the loss of local structural information in deep GNNs?
• RQ3 How does MLAP compare to existing aggregation methods like Jumping Knowledge (JK) in terms of discriminability and performance?
• RQ4 Does the integration of MLAP with normalization techniques like GraphNorm further enhance model generalization?
• RQ5 What is the impact of different aggregation strategies (e.g., sum vs. weighted) on training stability and performance?

Key findings

• MLAP significantly improves graph classification accuracy on the MCF-7 and synthetic datasets, outperforming both baseline GNNs and the JK aggregation method.
• The performance gain is particularly notable when combined with GraphNorm, where MLAP + GraphNorm achieves superior results compared to naive GNN + GraphNorm and JK + GraphNorm.
• t-SNE visualization confirms that layer-wise representations from MLAP are more discriminative, with clearer separation between positive and negative samples.
• MLAP-Weighted performed worse than MLAP-Sum in some datasets, suggesting that learnable weights may introduce training instability.
• The architecture is compatible with various GNN backbones and normalization techniques, demonstrating broad applicability and potential for integration into existing models.
• Analyses show that aggregating representations from multiple levels of locality enhances the discriminability of learned graph embeddings, supporting the core hypothesis of the method.