Что такое graph pooling (TopK, DiffPool, SAGPool)?

Graph pooling — это методы агрегации и снижения размерности графовых данных в графовых нейронных сетях (GNN), позволяющие выделить важные подграфы или узлы и упростить структуру для последующей обработки.

• TopK pooling выбирает топ-K узлов с наибольшими значениями некоторой скалярной оценки важности (например, на основе признаков узлов), отбрасывая остальные. Это позволяет сохранить наиболее информативные узлы.

• DiffPool — дифференцируемый метод кластеризации узлов, который обучается создавать мягкое разбиение узлов на кластеры, формируя иерархическую структуру графа. Он генерирует матрицу присвоения узлов к кластерам, что позволяет уменьшить размер графа и сохранить его глобальную структуру.

• SAGPool (Self-Attention Graph Pooling) использует механизм самовнимания для вычисления важности узлов, выбирая наиболее релевантные для задачи узлы. Это улучшает качество отбора по сравнению с простым TopK, учитывая контекст узлов.

Пример использования TopK pooling в PyTorch Geometric:

import torch
from torch_geometric.nn import TopKPooling

pool = TopKPooling(in_channels=16, ratio=0.5)
x = torch.randn((num_nodes, 16))  # признаки узлов
edge_index = ...  # ребра графа

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

Таким образом, graph pooling помогает эффективно обрабатывать большие графы, выделяя ключевые структуры и снижая вычислительную сложность.