What is graph pooling (TopK, DiffPool, SAGPool)?

Graph pooling is a set of methods for aggregating and reducing the dimensionality of graph data in graph neural networks (GNNs), allowing the extraction of important subgraphs or nodes and simplifying the structure for subsequent processing.

• TopK pooling selects the top-K nodes with the highest scalar importance scores (e.g., based on node features), discarding the rest. This helps retain the most informative nodes.

• DiffPool is a differentiable clustering method that learns to create soft partitions of nodes into clusters, forming a hierarchical graph structure. It generates an assignment matrix that maps nodes to clusters, reducing the graph size while preserving its global structure.

• SAGPool (Self-Attention Graph Pooling) uses a self-attention mechanism to compute node importance, selecting the most relevant nodes for the task. This improves the quality of selection compared to simple TopK by considering node context.

Example of using TopK pooling in PyTorch Geometric:

import torch
from torch_geometric.nn import TopKPooling

pool = TopKPooling(in_channels=16, ratio=0.5)
x = torch.randn((num_nodes, 16))  # node features
edge_index = ...  # graph edges

x, edge_index, _, batch, _, _ = pool(x, edge_index)

Thus, graph pooling helps efficiently process large graphs by highlighting key structures and reducing computational complexity.