Geometry-Aware Edge Pooling for Graph Neural Networks

Summary

New experiments with preliminary results and additions:

• Computational efficiency: Evaluate runtime and memory usage across pooling methods. Report and contrast pre-training and training costs. Highlight the computational improvements compared to EdgePool.
• Regression experiments: Report and compare the performance of our pooling method at regression tasks during revisions.

Clarifications and revisions:

• SOTA results: Clarify the purpose of graph pooling. Contextualise the reported accuracy gains w.r.t. to baseline pooling methods and existing literature [1, 2].
• Novelty of diffusion distances: Explain the conceptual benefits of using diffusion distances on graphs. Highlight potential extensions of our methods based on established theory on diffusion maps [4, 5].
• Typos: Correct minor typos in the spelling of MagEdgePool.

Extended discussion:

• Computational efficiency: Add an extended discussion on the computational efficiency of our pooling method in comparison to alternative methods.
• Large-graphs tasks: Discuss potential extensions to large-scale graphs.

---

We would like to clarify the points raised by Reviewer GecT.

Improved Classification Accuracy: Achieves superior performance on most datasets, with significant improvements noted on specific datasets like DHFR.

We are thankful that this review acknowledges the superior classification accuracy and robust performance achieved by our pooling methods.

Inconsistent Performance: Although beneficial for most datasets, the proposed method underperforms on certain datasets (e.g., Mutagenicity), with only half achieving state-of-the-art (SOTA) results.

The purpose of our evaluation is not to achieve SOTA results across all datasets and all potential GNN architectures, but to give a consistent benchmark that enables a fair comparison of pooling methods across tasks. For that purpose we plug each pooling layer into the same GNN architecture as commonly done to compare pooling methods. However, it is unrealistic to believe that this setup could or should achieve SOTA results on every datasets without dataset-specific modifications.

Nevertheless, our results show consistent and robust performance across tasks, and reach very high accuracies above what is reported in most pooling benchmarks [1, 2]. In fact, on Mutagenicity our pooling method outperforms all alternative pooling layers included in our comparison. For this reason, we argue that our pooling methods reach SOTA results amongst baseline pooling methods and do not underperform.

High Computational Complexity: Despite being training-free, PCA computations introduce considerable complexity, potentially limiting scalability to large-scale graphs.

We would like to clarify that our pooling method does not compute PCA. We compute the magnitude or spread of a graph as defined in Equations (1) and (2) for edge pooling and discuss the computational complexity of these calculations in our manuscript. However, this is not equivalent to PCA. This seems to be a misunderstanding about the computation of our algorithm.

Further, we provide additional results on the memory cost and computational complexity in our response to AV7X. We thus find that the computational performance of our proposed pooling method is competitive for most standard graph classification datasets. Overall, our pooling approach notably improves on the computational efficiency of EdgePool. Thus, we address scalability limitations of existing (edge) pooling methods, such as EdgePool, DiffPool, and MinCutPool, which scale notably worse to large-scale graphs.

Runtime and Memory Usage: Could the authors provide detailed runtime and memory usage data? Comparisons with other training-free methods (e.g., TopK and MinCut) would be valuable.

We would like to clarify that TopK and MinCut are trainable rather than training-free methods. Considering the distinction between, trainable (TopK, MinCut, SAGPool, DiffPool) and training-free methods (ours, Graclus, NDP, NMF), we are happy to provide runtime and memory comparisons in the response to reviewer AV7X.

Specifically, we find that e.g. on IMDB-M our method SpreadEdgePool (99.7s per run) allows for notably faster training than MinCutPool (132.8s) and moderately faster training than TopKPool (110.9s). Thus, we find that the improved computational efficiency of GNN training is an advantage of our pooling method.

Lack of Novelty in Diffusion Distance: The diffusion distance concept applied by the authors is not novel; rather, the contribution lies in its specific application to graph pooling.

Our pooling method indeed leverages the diffusion distance, a well-studied metric on general measure spaces that leverages both geometrical and topological aspects of the input data. The novelty of our method lies in its formulation as a technique based on metric space magnitude, a recently introduced invariant for metric spaces. As reviewer 7pUv pointed out, we are the first paper to study magnitude-informed approaches in the context of graph learning. The use of the diffusion distance, which users can replace if they need to, is rather a strength of our method.

For further research, we believe that the usage of diffusion distances offers the chance to leverage a rich theory on approximation methods via landmarks [4], or localised diffusion computations [5], which can lead to further computational improvements and extension of our methods. We thus believe the established theory on diffusion maps is a benefit of our methodology.

Spelling Errors: The manuscript contains several spelling errors (e.g., "MagEgdePool" instead of "MagEdgePool" on line 195; repeated "SpreadEgdePool" errors on lines 220, 229, and 341), suggesting the need for careful proofreading.

We thank the reviewer for pointing out specific typos and look forward to correcting them in the final manuscript. We are happy to rectify these minor mistakes in our revision.

Scalability to Large-Scale Graphs: Current experiments focus on small graphs. Testing on large-scale datasets, such as PDB for Enzyme Commission (EC) number prediction or Gene Ontology (GO) term prediction [1], would better demonstrate practical applicability.

Experiments on large-scale graphs would indeed be of interest, but truly massive graphs are rarely encountered for graph classification tasks [3] and offer a particular set of challenges regarding memory and computational costs. Existing pooling methods that process dense adjacency matrices, such as MinCut or DiffPool, are unlikely to scale to such tasks without major modifications [2, 3].

Given these necessary architecture adjustments, large-scale experiments remain outside the scope of our work for now. However, for future work, we note that our edge pooling method could indeed be scaled to larger graphs by e.g. computing the edge scores in a localised or mini-batched manner on carefully chosen subgraphs rather than on the whole graph. The sparse pre-trained edge assignment could then be used efficiently for training on large-scale graphs.

We also note that already, our pooling method notably outperforms EdgePool or dense pooling methods such as MinCut or DiffPool in terms of memory costs and computational efficiency, improving the scalability of edge pooling in most practical applications to standard graph classification tasks.

Extension to Regression Tasks: The study currently focuses exclusively on classification. Could the proposed method also be extended to regression tasks, such as Protein Localization Prediction on UniProt [2]?

We would like to note that protein localisation prediction on UniProt is a node classification task rather than a regression task. Working with the suggested datasets, PDB or UniProt, would require domain-specific preprocessing and careful evaluation beyond what is feasible in the scope of rebuttals.

Our pooling architectures can be extended to either node classification or graph regression problems. However, we consider graph-level tasks to be the most relevant applications of pooling methods [1] and thus choose to focus on graph classification tasks for our analysis.

Overall, we hope that our answer could clarify any misunderstandings.

References

[1] Wang, P., et al., 2024. A comprehensive graph pooling benchmark: Effectiveness, robustness and generalizability.

[2] Grattarola, D., et al., 2022. Understanding pooling in graph neural networks.

[3] Bianchi, F.M. and Lachi, V., 2023. The expressive power of pooling in graph neural networks.

[4] Long, A.W. and Ferguson, A.L., 2019. Landmark diffusion maps (L-dMaps): Accelerated manifold learning out-of-sample extension.

[5] David, G. and Averbuch, A., 2012. Hierarchical data organization, clustering and denoising via localized diffusion folders.

Summary

New experiments with preliminary results and additions:

• Regression experiments: Report and compare the performance of our pooling method at regression tasks during revisions.

Clarifications and revisions:

• Definition of a graph: Correct the definition of a graph to be consistent across the manuscript.
• Changes in the abstract: Clarify the sentence in the abstract.
• Figure 1: Double-check the number of nodes. Specify that graphs are pooled to approximately half their size.

---

We thank the reviewer for their positive feedback, in particular for highlighting that our pooling approach demonstrates robust and generalisable performance.

In some parts, the graph is defined as G = (X, E), while in some other parts it is G = (V, E). Besides, the definitions of X, V, and E are not given.

In the abstract, I am confused about this sentence: "However, existing pooling operations often optimise for the learning task at the expense of fundamental graph structures and interpretability."

We thank the reviewer for pointing out specific revisions in the text and we will clarify the definition of a graph and reformulate the sentence in the abstract. Specifically, the sentence refers to the observation that existing trainable pooling operators tend to optimise for a task-specific loss. However, such methods might not preserve the structure of the graph during pooling and make counterintuitive choices on which graph properties to preserve. This is illustrated by the comparison between pooling methods in Figure 1, which shows that pooling methods such as TopKPool, SAGPool, DiffPool, or MinCutPool do not necessarily lead to interpretable graph coarsening.

In the visualization results from Figure 1, the pooling ratio of different pooling methods are not consistent.

The pooling ratios in Figure 1 are consistently set to 50%. However, some pooling methods (e.g. Graculus or NDP) do not give exact control over the number of vertices, but pool graphs to approximately half of their original size. This is why the exact number of nodes differ slightly. We will clarify this in the manuscript and aim to provide a revised version of Figure 1 with more comparable node numbers whenever possible.

Is the proposed method in this paper effective for the graph regression task?

Our method can be extended to regression tasks and flexibly plugged into any GNN. For revisions, we plan to incorporate regression tasks into our analysis and adjust our model architecture and evaluation to fit these tasks. We expect regression performance to agree with the classification results. The arguments why our proposed edge pooling methods perform well at graph classification, namely that features are aggregated in an expressive manner and that graph structure is preserved during pooling, similarly hold for graph regression tasks.

Summary

Extended experiments with preliminary results and additions:

• Computational efficiency: Evaluate runtime and memory usage across pooling methods. Report and contrast pre-training and training costs.
• Computational speed-ups: Propose and benchmark approximation techniques for estimating diffusion distances to speed up graph pooling.

Extended discussion:

• Computational efficiency: Discuss further strategies for ensuring computational scalability.
• Reproducibility and usability: Include details on the reproducibility and open-source implementation of our pooling method.

---

We thank the reviewer for their positive feedback and concrete questions, which we are happy to address. Specifically, we are thankful for the review highlighting the rigorous mathematical and theoretical grounding and strong performance of our proposed pooling method.

Can you provide more experimental results regarding the efficiency of the method?

We will include extended runtime and memory comparisons in our revisions and want to present experimental results on the computational efficiency of our pooling method. Specifically, we train the GNN architecture specified in our manuscript using GIN layers across 200 epochs using 10-fold stratified cross-validation. We compare our edge pooling methods (MagEdgePool , SpreadEdgePool ) to trainable pooling methods from torch_geometric (EdgePool, TopKPool, SAGPool) or from torch-geometric-pool (DiffPool, MinCutPool). We record the mean and standard deviation of the runtimes in seconds and GPU memory usage in MB per cross-validation run.

We note that in particular our proposed edge pooling methods MagEdgePool and SpreadEdgePool allow for significantly more efficient GNN training than EdgePool. Similarly, our methods generally improve on the runtimes for dense pooling methods such as DiffPool and MinCutPool, and improve on the computational efficiency of TopKPool and SAGPool.

Table of training times in seconds compared across pooling methods. The fastest methods are marked in bold.

Table of memory usage in MB compared across pooling methods. The most efficient methods are marked in bold.

Having observed that pre-computed edge-pooling speeds up GNN training times, we further compare with the runtime and memory costs of other non-trainable pooling operators prior to training. We report the computational costs of computing the pooling assignment for all graphs in the datasets.

As reported in our manuscript, we find that SpreadEdgePool is generally more efficient than MagEdgePool. Beyond the exact computations of our pooling methods, we also approximate the change in diffusion distance during edge contraction by updating the contracted node’s distances to be the minimum distance to the original nodes. These approximate versions denoted MagEdgePool* and SpreadEdgePool* notably reduce computational costs.

Table of pre-training times in seconds compared across non-trainable pooling methods. The fastest method is marked in bold. The fastest approximation of our pooling method is marked in cursive.

Table of pre-training memory usage in MB compared across non-trainable pooling methods. The most efficient method is marked in bold. The fastest approximation of our pooling method is marked in cursive.

The presented results show one of the limitations discussed in our manuscript, namely that our edge pooling approach scales in the number of edges as well as the size of the graphs. Nevertheless, our proposed pooling methods still notably outperform EdgePool in terms of computational efficiency when considering both the preprocessing and training costs as reported above. Further, the pretraining costs of our method range in a number of seconds, only need to be computed once per dataset, and the memory requirements remain below what is required by GNN training. Hence, we find that our proposed pooling methods scale sufficiently well to standard graph datasets.

We look forward to presenting extended experiments on computational efficiency in revisions. For future work, we believe there is a strong potential to further speed up our pooling methods by investigating alternative graph metrics, testing sampling heuristics to restrict the edge score computations to a subset of candidate edges, or computing edge score on local subgraphs rather than on the whole graph.

Is the code open-source and reproducible?

We will publish a light-weight implementation of our pooling method as an open-source package to allow users to try our pooling method on their own datasets.

This is in addition to our supplementary code submission, which can be used to reproduce our experiments.

Thank you for the response. However, it appears that the author has overlooked the second weakness regarding scalability with the increasing number of graph nodes. Moreover, it is valuable to include a comparison and discussion of global pooling methods such as SOPool [1] and DKEPool [2], etc., as this would make the analysis more comprehensive.

[1] Wang, Z., & Ji, S. (2023). Second-order pooling for graph neural networks. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(6), 6870–6880.
[2] Chen, K., Song, J., Liu, S., et al. (2023). Distribution knowledge embedding for graph pooling. IEEE Transactions on Knowledge and Data Engineering, 35(8), 7898–7908.

Summary

Extended experiments:

• Comparison with EdgePool: Extend experiments with EdgePool on IMDB-M, IMDB-B, and further datasets. Highlight runtime and memory efficiency improvements.
• Randomised baselines: Add randomised edge pooling experiments (already partly reported for NCI1, NCI109). Expand to further datasets.
• Spectral preservation: Reproduce Figure 3 analyses for further datasets to confirm consistent magnitude/spectral distance patterns.

Clarifications and revisions:

• Positioning vs. EdgePool: Emphasise that the main advantages are computational scalability and flexible pooling ratios, not raw accuracy gains.
• Accuracy interpretation: Clarify that pooling aims at efficiency + structure preservation; contextualise improvements with stratified 10-fold CV and prior baselines.
• Link between magnitude/spread and accuracy: Elaborate theoretical reasoning behind structure preservation; connect to expressivity results in existing work [3].

Extended discussion:

• Add a dedicated subsection on expressivity vs. structure preservation.
• Incorporate recent findings [2,3,4,6] to contextualise when structural preservation matters.
• Discuss conditions under which randomised baselines can remain competitive, clarifying limits of structural information use.

---

We thank the reviewer for the insightful and detailed questions and find they offer a great opportunity to discuss the reasons behind the robust performance of our pooling methods.

The authors say that EdgePool [1] is the most successful edge-based. Although they argue its major drawback of pooling graphs to half their size, there is no experimental comparison to support the superior performance of MagEdgePool/SpreadEdgePool with another edge-pooling method.

To extend our evaluation, we compare with the pytorch_geometric implementation of EdgePool. We find that there is no significant difference in performance between our methods and EdgePool when computed on e.g. IMDB-M, where EdgePool reaches an accuracy of 47.6 ± 3.3; and 70.9 ± 4.5 on IMDB-B. We expect this same pattern to hold across further datasets and will report extended results during revisions.

As a major advantage, our methods have notably lower computational costs during training than EdgePool (cf. the response to reviewer AV7X). Our methods thus make edge-contraction pooling scalable to larger datasets and enable significantly faster GNN training. On IMDB-B for example, EdgePool takes 362.2 seconds and 666.6 MB for 200 epochs during training, but our method SpreadEdgePool only requires 62.7 seconds and 283.4 MB.

We thus find that learning feature-based edge scores as done by EdgePool, is not necessary to ensure classification performance, but rather adds a notable computational burden. The main advantage of our pooling methods over EdgePool is thus improved scalability as well as the flexible choice of pooling ratio.

The author cites [2] but they do not debate their argument that pooling methods (e.g GRACLUS and DIFFPOOL) are not significantly superior to baselines built upon randomized cluster assignments.

We thank the reviewer for their insightful suggestion and we will include further discussion and randomised ablations to the paper.

We find that this discussion and the observations by [2] are closely connected to the following challenges: Standard graph learning tasks do not necessarily require the graph structure to ensure high performance [4, 6]. Randomised modifications of expressive pooling operators can retain the expressivity of the underlying pooling operation [3].

Hence, if all task-relevant information is already captured by the learnt features before pooling, randomised baselines can reach competitive performance. Nevertheless, destroying the structure of graphs during pooling can hinder the effectiveness of the preceding message passing layers [3]. Our pooling approach aims to mitigate exactly this risk by presenting geometry-aware and expressive edge-pooling methods.

As preliminary results, we add randomised edge pooling to our main experiment reported in Table 1. We find that for NCI109 randomised edge pooling reaches an accuracy of 68.0 ± 3.0 as compared to 71.8 ± 1.8 for SpreadEdgePool. On NCI1 random edge pooling reaches an accuracy of 71.1 ± 2.2 compared to 73.4 ± 2.5 for SpreadEdgePool. Thus, the random baseline performs worse than any alternative pooling method on NCI109 and decreases the accuracy on NCI1 slightly. We hypothesise that this decrease in accuracy is observed because graph structure is relevant for these learning tasks [4] and our pooling method better preserves structural properties.

The accuracy reported on Table 1 presents only slightly better results, i.e, the experimental results are not convincing.

The accuracies in Table 1 are computed for GNNs whose only difference is the choice of pooling layer to allow for a comprehensive and fair comparison across pooling methods. In this context, it is in fact a challenging task to achieve very high performance gains solely based on the choice of pooling strategy. We therefore argue that the goal of pooling is not necessary to beat any non-pooling baseline, but to enable more efficient training while preserving task-relevant structure or feature information.

Given we use stratified 10-fold cross validation to ensure a fair comparison, the reported accuracies are comparable to the baseline performance reported by [5] and similar or higher than the degree of improvement reported by established pooling methods, such as EdgePool [1]. Thus, our results show robust performance and realistic accuracy gains.

We are happy to extend our experimental setup with additional comparison partners or datasets in case the reviewer has specific suggestions.

The DHFR dataset was the single case where an expressive improvement was obtained. Do the authors have an explanation for this?

DHFR most benefits from the regularising effect of structure-aware edge pooling due to two potential factors. First of all, DHFR shows structural and geometric differences between graphs representing active and inactive compounds. Our pooling method helps preserve key graph structures, such as rings and graph connectivity, during pooling as visualised in Figure 2 in our manuscript. This can aid the preceding layers’ capability to discriminate between the coarsened graphs. Second, our pooling method gives an expressive pooling operator as discussed above. Amongst other expressive operators, we hypothesise that the aggregation via edge collapse pooling might be most suitable for this task specifically.

Figure 3 indicates that preserving magnitude corresponds to lower spectral distances and better retention of spectral properties. Would you say this also corresponds to the lower difference between the MagEdgePool/SpreadEdgePool and No Pooling results reported in Table 1? Since Figure 3 only shows the results for the NCI1, could report the same analyses for another dataset such as DHFR?

Yes, we can reproduce the results from Figure 3 for other graph datasets, such as DHFR, and will report further datasets as ablations in the revisions. The same patterns we report for NCI1 hold for other datasets in the sense that we can consistently show that MagEdgePool and SpreadEdgePool induce both low spectral distances and low differences in magnitude during pooling indicating superior preservation of the graph structure.

We argue that the more faithful preservation of graph properties does contribute to the better performance reported in Table 1. This is particularly relevant for classification tasks such as NCI1 that rely on the geometry of the input graph for task performance [3, 4].

In general, do you have an intuition/experiment/theorem that explicitly relates the magnitude/spread (or magnitude’s difference between input and reduced graph) with the accuracy?

There is a strong link between preserving magnitude/spread during pooling and preserving the graph structure of the input graph. Magnitude can be expressed as a polynomial with variables given by the eigenvalues of the similarity matrix Z = [exp(-d(xi,xj))], which we compute from the diffusion distances d. We believe that under mild conditions small changes in the Laplacian matrix lead to small changes in the similarity matrix which then induce small spectral distances between the input and the pooled graph. We aim to analyse this relationship further in future work.

Further, we find theoretical links between the expressivity of pooling operators and the accuracies reported in Table 1. Edge contraction pooling when for example using sum pooling for aggregating the features fulfills sufficient conditions laid out by [3] for preserving the expressive power of the preceding message passing layers. Thus our pooling methods fulfil theoretical guarantees for both expressive feature retention as well as graph structure preservation.

We hypothesise that structure preservation will be particularly beneficial in case labels (that is, tasks) are directly related to the structure. As recent research shows [6], this prerequisite is not necessarily given, warranting careful study. We will add an extended discussion on expressivity and structure-preservation during revisions.

References

[1] Diehl, F., et al., 2019. Towards graph pooling by edge contraction.

[2] Mesquita, D., et al., 2020. Rethinking pooling in graph neural networks.

[3] Bianchi, F.M. and Lachi, V., 2023. The expressive power of pooling in graph neural networks.

[4] Coupette, C., et al., 2025. No Metric to Rule Them All: Toward Principled Evaluations of Graph-Learning Datasets.

[5] Grattarola, D., et al., 2022. Understanding pooling in graph neural networks.

[6] Bechler-Speicher, M., et al., 2024. Graph neural networks use graphs when they shouldn't.