ENAHPool: The Edge-Node Attention-based Hierarchical Pooling for Graph Neural Networks

Q1: There is an ablation study but I don't feel it covers all the claims and components of the proposed methodology.

A1: Due to space limitations, please refer to our response Q2 to Reviewer foip.

Q2: The contribution to the basic ML research.

A2: This paper proposes a new graph pooling that can overcomes the drawbacks of many existing pooling methods. Specifically, we would like to clarify that GNNs are always a crucial research direction for machine learning, and their development relies on continuously optimizing existing technologies.

Specifically, graph pooling is an essential subfield of GNNs, and most existing methods still suffer from some theoretically drawbacks. For instance, global pooling methods fail to capture hierarchical structural features of graphs. Hierarchical pooling methods based on Top-K selection discard a significant amount of node information, influencing the connectivity construction of the coarsened graph. Cluster-based hierarchical pooling methods can solve the above problems, but they typically aggregate node features and edge connectivity strengths using simple summation, ignoring the differences between nodes and edges. Thus, we propose applying attention mechanisms for weighted aggregation, resulting in more effective coarsened graph representations.

The attention mechanism has been widely applied across various fields. Its core idea is to dynamically assign weights to capture key information, enhancing the model's representation capability. Recently, attention mechanisms have gained traction in graph tasks. For instance, [1] and [2] employ multi-head attention to assess neighbor importance, enabling differentiated aggregation for effective node representations. In graph pooling, [3] uses soft attention to weight node importance, while [4] applies multi-head attention to identify task-relevant parts of the input data and learn each node’s global significance after pooling. [5] integrates attention into second-order pooling, further boosting model expressiveness. Inspired by these works, we introduce a node attention mechanism akin to the self-attention mechanism in [6], whose success highlights its ability to focus on graph-relevant nodes, providing a solid theoretical foundation for our approach.

Among existing methods, edge attention mechanisms remain underexplored. A notable example is [7], which leverages attention to compute edge contraction scores. Another approach, [8], suggests assigning higher importance to edges connecting more different nodes. Inspired by these works, we introduce an edge attention mechanism akin to the edge information scoring function in [9]. The success of [9] demonstrates its ability to identify important edges in a graph, providing a strong theoretical foundation for our edge attention aggregation strategy.

However, it is important to note that while GAT can be seen as an exploration of edge attention mechanisms, it essentially computes node attention to selectively aggregate neighboring node information, without directly processing edges. This differs from ours, which is designed to aggregate edge connectivity strengths between clusters. The connectivity strengths in the coarsened graph reflects the influence of the information propagation between clusters, determined by both edge importance and quantity. For instance, in a social network, the number of edges between the leadership teams of Companies A and B may be the same as that between the employees of Companies A and C, yet B and C’s influence on A could differ. Our edge attention mechanism aims to capture such differences.

In conclusion, we believe this study provides both practical improvements and a new perspective for theoretical research in machine learning.

[1] Velickovic, P. et al. Graph attention networks. In ICLR, 2018. URL https://openreview.net/forum?id=rJXMpikCZ.

[2] Zhang, J. et al. Gaan: Gated attention networks for learning on large and spatiotemporal graphs. In UAI, pp. 339–349, 2018.

[3] Li, Y. et al. Gated graph sequence neural networks. In ICLR, 2016. URL http://arxiv.org/abs/1511.05493.

[4] Xu Y. et al. Multistructure graph classification method with attention-based pooling. In IEEE TCSS, 10(2): 602-613, 2022.

[5] Wang Z. et al. Second-Order Pooling for Graph Neural Networks. In TPAMI, vol. 45, no. 6, pp. 6870-6880, 2023.

[6] Lee, J. et al. Self-attention graph pooling. In ICML, volume 97 of Proceedings of Machine Learning Research, pp. 3734–3743, 2019.

[7] Diehl, F. et al. Towards graph pooling by edge contraction. In ICML workshop, 2019.

[8] Gao, Z. et al. Lookhops: light multi-order convolution and pooling for graph classification. CoRR, 2020. URL https://arxiv.org/abs/2012.15741.

[9] Yu H. et al. Not all edges are peers: Accurate structure-aware graph pooling networks. In NN, 156: 58-66, 2022.

Q3: There are several references missing.

A3: Due to space limitations, please refer to our response Q1 to Reviewer PjzN.

Q1: References for the new pooling operations.

A1: Thank you for your valuable suggestion. We have further investigated more recently proposed pooling operations and will refine the related work section in the final version.

Recent research on graph pooling has primarily focused on cluster-based hierarchical methods. MathNet [1] leverages Haar-like wavelet multiresolution analysis to construct hierarchical graph representations. Tsitsulin et al. [2] introduced Deep Modularity Networks (DMoN), an unsupervised approach that optimizes the clustering process using the modularity measure. Bacciu et al. [3] proposed a pooling method based on Maximal k-Independent Sets (k-MIS), ensuring that selected nodes maintain a minimum pairwise distance of k. Zhou et al. [4] proposed Cross-View Graph Pooling, which integrates pooled graph information from both node and edge perspectives. WGDPool [5] utilizes a differentiable k-means clustering mechanism with Softmin assignments based on node-centroid distances.

[1] Zheng, X. et al. Mathnet: Haar-like wavelet multiresolution analysis for graph representation learning. Knowl. Based Syst., 273: 110609, 2023.

[2] Tsitsulin, A.et al. Graph clustering with graph neural networks. J. Mach. Learn. Res., 24:127:1–127:21, 2023.

[3] Bacciu, D. et al. Generalizing downsampling from regular data to graphs. In AAAI, pp. 6718–6727, 2023.

[4] Zhou, X. et al. Edge but not least: Cross-view graph pooling. In ECML PKDD, volume 13714 of Lecture Notes in Computer Science, pp. 344–359, 2022.

[5] Xiao, Z. et al. Wgdpool: A broad scope extraction for weighted graph data. Expert Syst. Appl., 249:123678, 2024.

Q2: Detailed time complexity analysis.

A2: In our proposed method, each pooling layer primarily involves the following steps. First, the MD-MPNN model performs convolution operations to obtain node embeddings. The computational complexity of MD-MPNN is $O(N^3)$ due to the need for matrix multiplications on the adjacency matrix to filter node information at different distances, helping to mitigate over-squashing issue. Second, the computational complexity of both node and edge attention mechanism are $O(N)$, as they are computed based on node features. Finally, the pooling operation has a time complexity of $O(KN^2)$, where $K$ represents the number of nodes in the next pooling layer, generally set as $rN$, with $r$ denoting the pooling ratio. Overall, the proposed method maintains an overall computational complexity of $O(N^3)$, which is comparable to other classic cluster-based hierarchical pooling methods such as StructPool [1] and MinCutPool [2].

[1] Yuan, H. et al. Structpool: Structured graph pooling via conditional random fields. In ICLR, 2020. URL https://openreview.net/forum?id=BJxg_hVtwH.

[2] Bianchi, F. M. et al. Spectral clustering with graph neural networks for graph pooling. In ICML, volume 119 of Proceedings of Machine Learning Research, pp. 874–883, 2020.

Q3: Improve the performance of other GNN associated with ENAHPool?

A3: Of course! Our proposed ENAHPool operation is compatible with any GNN architecture, as it focuses on aggregating nodes and edges during pooling, without constraining the choice of GNN. However, it is preferable to use a GNN that aggregates neighborhood information based on the adjacency matrix. Otherwise, it may fail to leverage the edge connectivity strengths between clusters learned by the edge attention mechanism (such as GAT), potentially causing our method to degrade into a purely node attention-based hierarchical pooling operation.

Q4: Address the over-squashing problem when perform the module for multiple times?

A4: To address the over-squashing issue, we analyzed the impact of the number of MD-MPNN layers on model performance (as shown in Figure 7 of the paper). The results demonstrate that, unlike traditional MPNNs, increasing the number of MD-MPNN layers gradually improves model performance, mainly due to its novel message-passing mechanism. However, once the number of layers exceeds a certain threshold, performance begins to decline slightly. This is likely because information from excessively distant nodes introduces noise, which diminishes the effectiveness of feature representation.

Q1: The abstract is a little long.

A1: We will update the abstract in the final version to make it more concise and easier to understand.

Q2: More dataset is preferred for Ablation Study.

A2: Thank you for your constructive suggestion. We have conducted the ablation experiments on all datasets. However, due to time constraints, we only report the average results from 5 runs of 10-fold cross-validation here.

Q3: Detailed discussion of the computational efficiency.

A3: In our proposed method, each pooling layer primarily involves the following steps. First, the MD-MPNN model performs convolution operations to obtain node embeddings. The computational complexity of MD-MPNN is $O(N^3)$ due to the need for matrix multiplications on the adjacency matrix to filter node information at different distances, helping to mitigate over-squashing issue. Second, the computational complexity of both node and edge attention mechanism are $O(N)$, as they are computed based on node features. Finally, the pooling operation has a time complexity of $O(KN^2)$, where $K$ represents the number of nodes in the next pooling layer, generally set as $rN$, with $r$ denoting the pooling ratio. Overall, the proposed method maintains an overall computational complexity of $O(N^3)$, which is comparable to other classic cluster-based hierarchical pooling methods such as StructPool [1] and MinCutPool [2].

It is worth noting that our paper repeatedly emphasizes the inefficiency of using attention mechanisms for node aggregation in existing cluster-based hierarchical pooling methods. For example, ABDPool [3] computes scaled dot-product self-attention for nodes within each cluster, while C2N-ABDP [4] extracts cluster representations via singular value decomposition and then computes attention between each cluster and its nodes for weighted aggregation. These methods require iterating over every cluster, leading to high computational complexity. In contrast, the attention mechanisms we propose are significantly more computationally efficient.

[1] Yuan, H. et al. Structpool: Structured graph pooling via conditional random fields. In ICLR, 2020. URL https://openreview.net/forum?id=BJxg_hVtwH.

[2] Bianchi, F. M. et al. Spectral clustering with graph neural networks for graph pooling. In ICML, volume 119 of Proceedings of Machine Learning Research, pp. 874–883, 2020.

[3] Liu, Y. et al. Abdpool: Attention-based differentiable pooling. In ICPR, pp. 3021–3026,2022.

[4] Ye, R. et al. C2N-ABDP: cluster-to-node attention-based differentiable pooling. In GbRPR, volume 14121 of Lecture Notes in Computer Science, pp. 70–80, 2023.

Q1: Convincing evidence is needed to claims that soft-assignment worsens the performance.

A1: Thanks for the suggestion. We conducted comparative experiments on all datasets to verify the positive impact of the hard assignment operation on model performance. However, due to time constraints, we only report the average results from 5 runs of 10-fold cross-validation here.

Q2: Convincing experiments and discussions are needed to verify the significance of different attention mechanisms.

A2: Theoretically, the aim of using edge-node attentions is to obtain a more meaningful coarsened graph, capturing more accurate hierarchical representations. These two attention mechanisms are indispensable. First, without the node attention, hierarchical pooling would flatten the aggregation of nodes within a cluster, disregarding their varying importance. Second, the edge connectivity strengths between clusters naturally serve as attention during graph convolution. To better leverage this property, we propose an edge attention mechanism, which enhances the aggregation weights of important edges and prevents the neglect of rare but highly valuable connections between each cluster.

The misunderstanding may have arisen because we only presented the ablation results for two datasets in the paper. To clarify this, we conducted comparative experiments on all datasets. However, due to time constraints, we only report the average results from 5 runs of 10-fold cross-validation here. The results show that both the node and edge attention mechanism have a positive impact on model performance, though their degree of influence varies across datasets. Moreover, the best results are consistently achieved when both node and edge attention mechanisms are used together, which demonstrates the effectiveness of our approach.

Q3: Evaluation criteria is reasonable, but the proposed methodology is not justified enough.

Due to space limitations, please refer to our response Q2 to Reviewer FuAf.

Q4: The experimental design is not clear.

To minimize the impact of data partitioning on model evaluation, we followed the testing methodology from [1]. Specifically, we performed 10 random splits of each dataset and conducted 10-fold cross-validation on each split. The model was evaluated on every validation set, resulting in a total of 100 evaluation results for each method on each dataset. The experimental setup of MinCutPool differs from ours. Although it also performs 10 random splits of the dataset, it does not conduct cross-validation. This operation may make the results more sensitive to the partition, which could explain the differences in reported accuracy.

In addition, for the baselines and benchmarks that were tested in [1], we directly cite the reported results, as we follow the same experimental setup. For those not covered in [1], we adhere to the hyperparameter settings specified in the paper for evaluation.

Additionally, regarding the auxiliary classifiers, we sincerely apologize for this oversight. In the initial version of our experiments, we attempted to include this module. However, after further discussion, we decided to remove it to ensure fair comparisons with the baselines. Unfortunately, due to insufficient proofreading of the final manuscript, this change was not properly reflected in the text. We will correct this in the final version and ensure that all experimental descriptions align with our code. If necessary, we can also submit the code with anonymous link, then you can check.

[1] Wu, J. et al. Structural entropy guided graph hierarchical pooling. In ICML, volume 162 of Proceedings of Machine Learning Research, pp. 24017–24030, 2022.