PaperlayerView exact revision on arXiv

Source-linked AI summary

Accurate Learning of Graph Representations with Graph Multiset Pooling

Jinheon Baek, Minki Kang, Sung Ju Hwang

arXiv:2102.11533v4cs.LG

TL;DR

Existing pooling methods may ignore task relevance and graph structure or fail to preserve distinctions captured by the WL test. The paper formulates pooling as graph multiset encoding and proposes GMT, an attention-based pooling mechanism. GMT is theoretically as powerful as the WL test and performs strongly across classification, reconstruction, and generation tasks.

  • Problem

    Existing sum or average pooling ignores task relevance and structural dependencies, while hierarchical methods may fail to preserve distinctions between graphs separated by the WL test.

  • Method

    GMT formulates graph pooling as graph multiset encoding and uses graph-structured multi-head attention to select representative nodes and model their interactions.

  • Results

    GMT is as powerful as the WL test and outperforms most pooling baselines across graph classification, reconstruction, and generation tasks.

  • Takeaways & Limitations

    GMT provides a compact graph representation that incorporates task relevance and global structure while retaining information useful for reconstruction and generation.

Abstract

from arXiv · show

Graph neural networks have been widely used on modeling graph data, achieving impressive results on node classification and link prediction tasks. Yet, obtaining an accurate representation for a graph further requires a pooling function that maps a set of node representations into a compact form. A simple sum or average over all node representations considers all node features equally without consideration of their task relevance, and any structural dependencies among them. Recently proposed hierarchical graph pooling methods, on the other hand, may yield the same representation for two different graphs that are distinguished by the Weisfeiler-Lehman test, as they suboptimally preserve information from the node features. To tackle these limitations of existing graph pooling methods, we first formulate the graph pooling problem as a multiset encoding problem with auxiliary information about the graph structure, and propose a Graph Multiset Transformer (GMT) which is a multi-head attention based global pooling layer that captures the interaction between nodes according to their structural dependencies. We show that GMT satisfies both injectiveness and permutation invariance, such that it is at most as powerful as the Weisfeiler-Lehman graph isomorphism test. Moreover, our methods can be easily extended to the previous node clustering approaches for hierarchical graph pooling. Our experimental results show that GMT significantly outperforms state-of-the-art graph pooling methods on graph classification benchmarks with high memory and time efficiency, and obtains even larger performance gain on graph reconstruction and generation tasks.

1 INTRODUCTION

Existing graph pooling methods can lose task-relevant node information, graph structure, or injective distinctions. GMT addresses these limitations through graph multiset encoding and attention-based pooling, with theoretical and empirical support.

  • Sum and average pooling treat all nodes equally, limiting task-specific graph representations and hierarchical compression.
  • Node drop methods lose information by discarding nodes, while node clustering methods incur high computational complexity from dense assignments.
  • GMT formulates pooling as graph multiset encoding with structural dependencies and uses attention to select representative nodes and encode their relationships.
  • GMT is evaluated on graph classification, reconstruction, and generation tasks, with larger gains reported for reconstruction than classification.
  • The proposed pooling operation can be as powerful as the WL test and can extend to node clustering approaches with learnable clusters.

2 RELATED WORK

Related work develops GNN node representations, graph pooling, and permutation-invariant set or multiset encoders. These lines of research motivate task-specific pooling that preserves graph information and structure.

  • Graph Neural Network: GNNs generally encode nodes by aggregating neighborhood features, while transformer-based GNNs additionally model relatedness between nodes.
  • Graph Pooling: Simple average pooling can lose information because it weights all node information equally without considering graph-specific key features.
  • (Multi-)Set Representation Learning: Graph nodes form a multiset because different nodes may have identical feature vectors, motivating multiset encoding for graph representation learning.

3 GRAPH MULTISET POOLING

GMT combines graph-structured attention with multiset pooling to compress node representations while preserving structural dependencies. Its design is permutation-invariant, injective under stated assumptions, WL-level in expressive power, and efficient in space.

  • Graph Multiset Transformer: GMT treats graph representation learning as multiset encoding with auxiliary structural dependencies and uses graph-structured attention to encode them.
  • Preliminaries: Graph pooling maps node representations to an entire graph representation after message passing.
  • Preliminaries: Average and sum readouts ignore graph structures, while node drop and clustering methods respectively discard nodes or face scalability limits.
  • Graph Multiset Transformer: Graph multi-head attention weights relevant node values through multiple projected query, key, and value subspaces.
  • Graph Multiset Transformer: GMH constructs attention keys and values with GNNs using node features and adjacency information, explicitly incorporating graph structure.
  • Graph Multiset Transformer: GMPool uses learnable seed vectors and attention to compress n nodes into k representative nodes, while SelfAtt models interactions among condensed nodes.
  • Connection with Weisfeiler-Lehman Graph Isomorphism Test: GMT can map different graphs to distinct embeddings with WL-level expressive power, while Graph Multiset Pooling uses O(nk) space that can be optimized to O(n).
  • Graph Multiset Transformer: The full GMT architecture combines GNN and pooling layers, and its pooling function can be extended to hierarchical clustering with learnable k cluster centroids.

4 EXPERIMENT

GMT is evaluated on graph classification, reconstruction, generation, ablation, and efficiency tasks. Across these experiments, it generally improves performance and retains graph information while using practical memory and time resources.

  • Experimental Setup: GMT is evaluated on classification, reconstruction, generation, and efficiency tasks using synthetic and real-world graph datasets.Classification uses TU and OGB datasets; reconstruction uses ring, grid, and ZINC graphs; generation includes molecule generation and retrosynthesis.
  • Graph Classification: GMT outperforms most graph-pooling baselines or matches the best baseline on graph classification benchmarks.The classification results cover 10 datasets, while some non-parametric GNN baselines still perform strongly on individual datasets.
  • Ablation Study: Graph attention with self-attention substantially improves performance over mean pooling, while GMT without message passing still captures graph multiset structure.The ablation removes graph attention, self-attention, and message-passing operations to identify sources of improvement.
  • Efficiency: GMT is memory-efficient through sparse-graph compatibility and takes less than or nearly about one second on large graphs.The efficiency comparisons use Erdos-Renyi graphs; HaarPool and EdgePool are described as slower models.
  • Graph Reconstruction: GMPool produces almost perfect reconstruction on synthetic ring and grid graphs and largely outperforms baselines on ZINC validity, exact match, and accuracy.The ZINC experiments vary the compression ratio, and the results indicate that the pooled representations retain information needed for reconstruction.
  • Graph Reconstruction: GMT yields more calibrated clustering than MinCutPool on ZINC examples, capturing detailed molecular substructures associated with successful reconstruction.The visualization marks assigned atom clusters and incorrectly predicted atoms in reconstructed molecules.
  • Graph Generation: GMT stabilizes valid molecule generation earlier than baselines and improves retrosynthesis performance when used with GLN.The generation experiments use MolGAN for molecule generation and GLN as the retrosynthesis backbone.

5 CONCLUSION

The paper proposes GMT to address task-insensitive aggregation and non-injective hierarchical pooling. It theoretically relates GMT to WL expressiveness and reports gains across classification, reconstruction, and generation tasks.

  • 5 CONCLUSION: GMT encodes node embeddings as a multiset while incorporating global graph structure and task relevance during compression.The method is intended to distinguish different graphs through distinct embeddings.
  • 5 CONCLUSION: GMT is theoretically justified as powerful as the Weisfeiler-Lehman test and can extend to node-clustering schemes.The conclusion also reports validation on 10 graph-classification datasets and additional reconstruction and generation tasks.
  • 5 CONCLUSION: GMT outperforms state-of-the-art graph-pooling models on most classification datasets and is superior on reconstruction and generation tasks.The latter tasks require more accurate graph representations than classification tasks according to the paper.

A.1 PROOFS REGARDING WEISFEILER-LEHMAN TEST

The proofs establish that GMPool and GMT can encode graph multisets injectively while preserving permutation invariance, yielding graph representations as powerful as the WL test.

  • A.1 PROOFS REGARDING WEISFEILER-LEHMAN TEST: GMPool maps different bounded-size multisets to unique outputs, including when each block produces multiple elements.The proof assumes a countable input feature space and constructs injective multiset mappings through attention-based pooling.
  • A.1 PROOFS REGARDING WEISFEILER-LEHMAN TEST: GMPool uses GNN-derived node representations and structural adjacency information to construct an injective function over the input multiset.The proof represents the pooled features as transformed sums over GNN-processed multiset elements.
  • A.1 PROOFS REGARDING WEISFEILER-LEHMAN TEST: Because GMPool ignores element order, it satisfies permutation invariance over the multiset.This property complements injectiveness in the graph-level pooling function.
  • A.1 PROOFS REGARDING WEISFEILER-LEHMAN TEST: GMT combines multiple GMPool and SelfAtt blocks to map two different graphs into distinct embedding spaces.This establishes the stated WL-level expressiveness for the resulting GNN with the proposed pooling functions.
  • A.1 PROOFS REGARDING WEISFEILER-LEHMAN TEST: SelfAtt can represent an injective element-wise transformation, allowing the sequence of GMPool and SelfAtt blocks to remain injective.The construction uses a zero attention-output matrix and row-wise feedforward processing, with MLPs supplying the required functions.

A.2 PROOFS REGRADING NODE CLUSTERING

The node-clustering proofs show that GMPool supports learnable hierarchical clustering while reducing pooling space requirements and avoiding the quadratic storage required by clustering approaches.

  • A.2 PROOFS REGRADING NODE CLUSTERING: GMPool condenses n nodes to k nodes in O(nk) space, further optimized to O(n) when k ≪ n.The attention query-key multiplication and weighted value aggregation both use O(nk) space.
  • A.2 PROOFS REGRADING NODE CLUSTERING: Sparse GNNs with GMPool require O(n + m) total space, whereas node clustering approaches require O(n2) space.Here n denotes nodes and m denotes edges in the graph.
  • A.2 PROOFS REGRADING NODE CLUSTERING: GMPool can perform hierarchical node clustering with k learnable cluster centroids generated from seed vectors.Its attention weights inherently form a soft assignment matrix C = w(QK^T).
  • A.2 PROOFS REGRADING NODE CLUSTERING: The proposed GMPool significantly outperforms previous node-clustering approaches because it explicitly learns data-dependent cluster centroids.The centroids are learned through the learnable seed vectors S.
  • A.2 PROOFS REGRADING NODE CLUSTERING: After pooling, GMPool can generate a new adjacency matrix for compressed nodes through node clustering, but this construction has quadratic scalability issues.The method therefore replaces adjacency with the identity matrix after the first block for multiple GMPools, approximating MH rather than GMH.

C EXPERIMENTAL SETUP

The experimental section introduces baselines, the proposed model, and evaluation procedures for graph classification, reconstruction, and generation.

  • C EXPERIMENTAL SETUP: The experiments begin by introducing the baselines and the proposed model.
  • C EXPERIMENTAL SETUP: The evaluation covers graph classification, graph reconstruction, and graph generation tasks.
  • C EXPERIMENTAL SETUP: Experimental details are described separately for each of the three task types.

C.1 BASELINES AND OUR MODEL

The section compares graph pooling baselines with GMT across global, hierarchical, and reconstruction architectures. GMT globally condenses nodes and models their interactions, while remaining compatible with hierarchical pooling.

  • Global graph classification: GCN and GIN use mean or sum pooling with a graph convolutional or graph isomorphism message-passing layer.
  • Global graph classification: Set2Set recurrently encodes all nodes using content-based attention.
  • Node drop baselines: SortPool, SAGPool, TopKPool, and ASAP are node-drop methods that remove low-scoring nodes or clusters.
  • Node clustering baselines: DiffPool, MinCutPool, HaarPool, and StructPool are node-clustering methods that coarsen nodes using learned, spectral, wavelet, or conditional-random-field-based assignments.
  • Our model: GMT first condenses nodes with GMPool and then models interactions within the resulting set; its global READOUT layer can replace the final layer of hierarchical methods.

C.2 GRAPH CLASSIFICATION

The classification experiments use biochemical, social, and molecular graph benchmarks with standard dataset-specific features, metrics, splits, and efficiency protocols. Memory and time comparisons vary graph size and pooling settings under controlled conditions.

  • Datasets and metrics: The study evaluates six TU datasets across biochemical and social domains using classification accuracy.
  • Datasets and metrics: Four OGB molecule datasets are evaluated with ROC-AUC and their additional atom and bond features.
  • Evaluation protocol: TU experiments use 10-fold cross-validation with conventional splits and reserve 10 percent of training data for validation.
  • Implementation settings: Most TU and OGB settings use hidden size 128, while D&D uses hidden size 32 and batch size 10 because large graphs constrain clustering with large batches.
  • Efficiency experiments: Memory experiments use Erdos-Renyi graphs with m = 2n and compress n nodes to k = 4, whereas time experiments use m = n^2/10 and a 25% pooling ratio for most models.

C.3 GRAPH RECONSTRUCTION

The reconstruction experiments assess whether pooling preserves node and molecular information across synthetic and real-world graphs. They compare common pooling baselines with GMPool under controlled compression, training, and evaluation settings.

  • Synthetic graphs: Synthetic ring and grid graphs use 2-D node coordinates and intact adjacency matrices to test restoration of all node locations after pooling.
  • Molecule graphs: Molecule reconstruction uses 12K ZINC graphs to restore the exact atom types of all nodes across compression ratios.
  • Architecture and baselines: The reconstruction architecture uses two message-passing layers before pooling and after unpooling, with one sequential pooling and unpooling operation.
  • Architecture and baselines: Comparisons include TopKPool, DiffPool, MinCutPool, and GMPool, with GMPool used without SelfAtt in the node-clustering-style setup.
  • Training settings: Synthetic experiments use a 25% pooling ratio and MSE, while molecule experiments use 5%, 10%, 15%, and 25% ratios with cross-entropy loss.
  • Evaluation metrics: Molecule reconstruction is measured by validity, exact match, and atom-type accuracy.
  • Generation extensions: Generation experiments replace graph embedding functions with GMT-based GMPool1 in MolGAN and GLN, using QM9 and USPTO-50k respectively.

D ADDITIONAL EXPERIMENTAL RESULTS

Additional results report validation and leaderboard comparisons, adjacency-reconstruction analyses, and molecule-clustering visualizations. They show strong classification and reconstruction behavior while identifying distinct information requirements and a remaining adjacency-reconstruction limitation.

  • Graph classification: GMT outperforms most baselines or matches the best baseline on validation graph-classification results.
  • Graph classification: Reproduced OGB results are nearly identical to leaderboard results despite different hyperparameters and random seeds.
  • Graph reconstruction: Synthetic reconstruction errors are large when the training objective and evaluation metric target different node or adjacency information.
  • Graph reconstruction: For GMT on ring graphs, adjacency errors are 0.0331 when minimizing node discrepancies and 0.0324 when minimizing adjacency discrepancies.
  • Limitations: Adjacency reconstruction remains limited because continuous predictions must be transformed into discrete 0-or-1 edges; further reconstruction and visualization are left for future work.
  • Molecule reconstruction: Molecule visualizations show similar atoms or connected structures receiving similar clusters, including distinct clusters for oxygen-, fluorine-, and ring-related patterns.
Loading 2102.11533v4…