Graph Multiset Transformer (GMT)

• The paper introduces GMT, a novel graph pooling approach that reformulates node aggregation as a multiset encoding task augmented with explicit structural modeling.
• GMT employs a graph multi-head attention mechanism and Transformer-style self-attention to transform variable-sized node sets into fixed-dimensional embeddings with the power of the 1-WL test.
• Experimental results show GMT's superior performance and efficiency in graph tasks, delivering competitive accuracy on classification, reconstruction, and generation while scaling to large graphs.

The Graph Multiset Transformer (GMT) is a global graph pooling mechanism that formulates the graph pooling problem as a multiset encoding challenge augmented with explicit modeling of structural dependencies. GMT offers a permutation-invariant, injective, and computationally efficient approach for condensing arbitrary node sets into fixed-dimensional graph representations, using a graph-structured multi-head attention mechanism. The design directly addresses the limitations of both simple sum/average pooling and established hierarchical pooling protocols by integrating auxiliary graph structure within the pooling operation and achieving expressive power up to the 1-dimensional Weisfeiler–Lehman graph isomorphism test (Baek et al., 2021).

1. Motivation and Background

In graph neural networks (GNNs), effective node and edge representations are achieved via iterative neighborhood aggregation. However, numerous downstream tasks—such as graph classification, reconstruction, and generation—require a global, fixed-dimensional graph embedding. Common strategies for this readout step involve functionals across the set of node embeddings, including sum, mean, or max pooling. These functionals uniformly weigh nodes and cannot adapt to task relevance or exploit higher-order structural dependencies. Hierarchical pooling via node dropping (e.g., TopKPool, SAGPool, SortPool) risks critical information loss, while clustering-based schemes (e.g., DiffPool, MinCutPool) require dense assignments and adjacency coarsening, incurring $O(n^2)$ computational demands and potentially sacrificing injectiveness—so different, non-isomorphic graphs can collapse to indistinguishable representations.

Reconceptualizing the problem, GMT considers the node embeddings as a multiset and incorporates the adjacency matrix as auxiliary structure. The graph pooling objective is then to design a function:

$$\text{READOUT}: \left(\{h_v\}_{v=1}^{n}, A\right) \to \mathbb{R}^d$$

that is permutation-invariant, injective over pairs $(\text{multiset}, \text{adjacency})$, and computationally efficient.

2. Model Architecture and Pooling Mechanism

GMT proceeds by compressing $n$ node features into $k$ “seed” embeddings using a multi-head attention architecture tailored for graphs, followed (optionally) by Transformer-style self-attention among the condensed vectors.

Key Components:

• Seed Matrix $S\in\mathbb{R}^{k\times d}$: A fixed, learnable set of $k$ query vectors.
• Graph Multi-Head Attention (GMH): For $h$ attention heads, each computes projections:
  • $Q_i = S W_i^Q$, $K_i = \mathrm{GNN}_i^K(X, A) W_i^K$, $V_i = \mathrm{GNN}_i^V(X, A) W_i^V$
  • Attention weights: $A_i = \mathrm{softmax}(Q_i K_i^T / \sqrt{d_k})$
  • Output: $O_i = A_i V_i$
  • Aggregate: $$\mathrm{GMH}(S, X, A) = \mathrm{Concat}(O_1, ..., O_h) W^O$$
• Graph Multiset Pooling (GMPool$_k$): Applies a residual connection, row-wise feed-forward layer, and layer normalization:

$$\mathrm{GMPool}_k(X, A) = \mathrm{LN}\bigl(H + \mathrm{rFF}(H)\bigr)$$

where $H = \mathrm{GMH}(S, X, A)$.

• Self-Attention Over Condensed Vectors: A Transformer block allowing the $k$ pooled vectors to interact.
• Readout Pipeline: Typically,
  1. $H_0 = \mathrm{GMPool}_k(X, A)$
  2. $H_1 = \mathrm{SelfAtt}(H_0)$
  3. Final embedding $h_G = \mathrm{GMPool}_1(H_1, I_k)$ (with $I_k$ denoting $k\times k$ identity)

Notably, graph structure is injected into attention keys and values through the local GNN blocks, crucially distinguishing GMT from structure-agnostic attention pooling.

3. Theoretical Properties

GMT guarantees permutation invariance: Any permutation of node order induces a corresponding permutation of key and value rows, but the set of outputs is invariant as a function of the input multiset. The construction yields injectiveness, contingent on the injectiveness of the underlying GNN projections in each attention head. Specifically, for each fixed seed, the output vector is an injective function over the input node multiset. Consequently, the architecture replicates the expressive capacity of the 1-dimensional Weisfeiler–Lehman (1-WL) graph isomorphism test: two non-isomorphic graphs distinguishable by 1-WL will map to different outputs under GMT.

The following table summarizes these properties:

Property	Guarantee	Condition
Permutation invariance	Always	Applies to any input ordering
Injectiveness	Yes, up to 1-WL power	Each GNN$_i$ must be injective (e.g., GIN or stronger multilayer)
Computational efficiency	$O(nk)$ time and memory	$k\ll n$

4. Extension to Hierarchical Pooling

GMT's soft attention coefficients $$S_{\mathrm{soft}} = \mathrm{softmax}(S K^T / \sqrt{d_k})$$ can be recast as a soft assignment matrix for cluster-based pooling, where each seed acts as a cluster centroid. This enables hierarchical coarsening analogous to DiffPool, but with $O(nk)$ complexity rather than $O(n^2)$. Coarsened adjacency $$\tilde{A} = S_{\mathrm{soft}}^T A S_{\mathrm{soft}}$$ and features $\tilde{X} = S_{\mathrm{soft}}^T X$ are efficiently obtained. By stacking multiple levels with fresh seeds, a sparse, hierarchical pooling framework is achieved, inheriting the injectiveness and permutation invariance of the underlying GMT mechanisms.

5. Experimental Results

GMT demonstrates superior or competitive performance across a variety of graph learning tasks:

• Graph Classification: Across TU-datasets (DD, PROTEINS, MUTAG, etc.) and OGB datasets (HIV, Tox21, ToxCast, BBBP), GMT achieves or matches state-of-the-art accuracy, outperforming both node-drop and clustering-based pooling methods. For instance, on DD, PROTEINS, and MUTAG, GMT reports mean accuracies/ROC-AUCs of $78.72\pm0.59$, $75.09\pm0.59$, $83.44\pm1.33$ respectively, and $77.56\pm1.25$ on HIV (AUC), exceeding or matching all baselines (Baek et al., 2021).

• Ablation Studies: Removal of graph-structured attention, self-attention, or GNN-based key/value projections degrades performance, establishing all components as contributing significant gains.
• Efficiency: GMT scales to graphs with up to 10,000 nodes with $<8$ GB GPU RAM and $<1$ s per forward pass; dense cluster methods require $>16$ GB and more time.
• Graph Reconstruction and Generation: GMT yields near-perfect reconstruction on synthetic graphs and higher atom-type exact-match ratios ($\approx0.68$ at $10\%$ compression vs. $0.55$ for MinCutPool) on ZINC molecules. In generative settings (MolGAN, retrosynthetic prediction with Graph-Logic Network), replacing traditional pooling with GMT improves validity and top-$k$ accuracies, respectively.

6. Practical Usage and Extensions

Practical recommendations include:

• For graph classification, replacing standard pooling with GMPool$_1$ is effective in any GNN pipeline.
• A single GMPool$_k$ improves graph auto-encoder reconstruction when used as pooling and unpooling.
• Stacking GMPool$_k$ layers facilitates hierarchical abstraction, with graph coarsening via $S_{\mathrm{soft}}^T A S_{\mathrm{soft}}$ when memory allows.
• Standard hyperparameters: $k\approx\min(n,10)$ seeds, $h=4$–$8$ heads, $d_k = d_v = d/h$, and one or two self-attention layers typically suffice.

Future directions proposed include designing specialized GNN$_i^{K,V}$ for different trade-offs between receptive field size and locality, experimenting with seed initialization, and extending the framework to edge- or subgraph-level pooling using multiple query seeds. GMT's injectiveness, efficiency, and expressive power make it particularly suitable as a default choice for global graph representation learning across diverse domains (Baek et al., 2021).