[Paper Review] Accurate Learning of Graph Representations with Graph Multiset Pooling

Jinheon Baek, Minki Kang+1 more|arXiv (Cornell University)|Feb 23, 2021

Advanced Graph Neural Networks48 references68 citations

TL;DR

The paper introduces Graph Multiset Transformer (GMT), a graph pooling method that treats graph representation as a multiset encoding and uses graph-structured attention to produce WL-test-like, permutation-invariant graph embeddings. GMT achieves state-of-the-art results on graph classification, reconstruction, and generation while being memory-efficient.

ABSTRACT

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.

Motivation & Objective

• Formulate graph pooling as a multiset encoding problem with graph structure awareness.
• Develop Graph Multiset Transformer (GMT) that uses graph-structured attention to produce permutation-invariant, injective graph embeddings.
• Prove GMT's expressiveness is at most as powerful as the Weisfeiler-Lehman (WL) test and extensible to node clustering.
• Empirically validate GMT on graph classification, reconstruction, and generation tasks across biochemical, social, and molecular datasets.

Proposed method

• Define GMPool as a seed-based, attention-driven pooling that condenses n node representations into k representative nodes using Graph Multi-head Attention (GMH).
• Use GMH to compute keys/values via graph-aware GNNs, not plain linear projections, to capture structural information.
• Incorporate a Self-Attention block to model interactions among pooled nodes and among all nodes, via multi-head attention.
• Prove injectiveness of GMPool and SelfAtt to align GMT with the WL-test power bound.
• Connect GMT to node clustering with a sparse, scalable approach that avoids dense adjacency operations.
• Evaluate GMT in graph classification, reconstruction, and molecule generation tasks; analyze memory/time efficiency.

Experimental results

Research questions

• RQ1 Can graph pooling be cast as a multiset encoding problem that remains permutation-invariant and injective?
• RQ2 Does Graph Multiset Transformer (GMT) achieve WL-test-level discriminative power while maintaining computational efficiency on large graphs?
• RQ3 How does GMT perform on graph classification, reconstruction, and generation compared to state-of-the-art pooling methods?
• RQ4 What are the memory and time trade-offs of GMT relative to clustering-based pooling methods?

Key findings

• GMT outperforms most graph pooling baselines on graph classification benchmarks with favorable memory and time efficiency.
• GMT yields larger performance gains on graph reconstruction and molecule generation tasks, indicating it preserves more graph information.
• Ablation shows graph attention and self-attention components contribute significantly to performance; GMT can operate effectively with or without additional message-passing layers.
• GMT is memory-efficient and scalable on sparse graphs, outperforming memory-heavy pooling baselines in practice.