PaperlayerView exact revision on arXiv

Source-linked AI summary

Structure-Aware Transformer for Graph Representation Learning

Dexiong Chen, Leslie O'Bray, Karsten Borgwardt

arXiv:2202.03036v3stat.MLcs.LG

TL;DR

Existing graph Transformers can avoid GNN limitations but positional or absolute encodings do not necessarily capture structural similarity between nodes. SAT adds subgraph-based structural information to self-attention, achieving state-of-the-art benchmark performance and improving base GNNs, while retaining quadratic attention complexity.

  • Problem

    Positional or absolute encodings in graph Transformers generally do not measure structural similarity between nodes and their neighborhoods.

  • Method

    SAT incorporates rooted subgraph representations into self-attention, combining attributed and structural similarity and allowing existing GNNs to extract subgraph information.

  • Results

    SAT achieves state-of-the-art performance on five graph and node prediction benchmarks and consistently improves the base GNN it builds upon.

  • Takeaways & Limitations

    SAT combines local structural information from GNNs with global interactions from Transformers and provides better interpretability than classic Transformers.

  • Takeaways & Limitations

    SAT retains the Transformer's quadratic self-attention complexity, while k-subgraph SAT also requires more memory than k-subtree SAT.

Abstract

from arXiv · show

The Transformer architecture has gained growing attention in graph representation learning recently, as it naturally overcomes several limitations of graph neural networks (GNNs) by avoiding their strict structural inductive biases and instead only encoding the graph structure via positional encoding. Here, we show that the node representations generated by the Transformer with positional encoding do not necessarily capture structural similarity between them. To address this issue, we propose the Structure-Aware Transformer, a class of simple and flexible graph Transformers built upon a new self-attention mechanism. This new self-attention incorporates structural information into the original self-attention by extracting a subgraph representation rooted at each node before computing the attention. We propose several methods for automatically generating the subgraph representation and show theoretically that the resulting representations are at least as expressive as the subgraph representations. Empirically, our method achieves state-of-the-art performance on five graph prediction benchmarks. Our structure-aware framework can leverage any existing GNN to extract the subgraph representation, and we show that it systematically improves performance relative to the base GNN model, successfully combining the advantages of GNNs and Transformers. Our code is available at https://github.com/BorgwardtLab/SAT.

1. Introduction

GNNs capture local graph structure but face expressiveness, over-smoothing, and over-squashing limitations. SAT addresses the challenge of encoding structural information into Transformers through structure-aware self-attention.

  • GNNs use multilayer message passing to aggregate local neighborhood information, supporting applications including drug discovery, protein design, and social network analysis.
  • Limited expressiveness, over-smoothing, and over-squashing constrain message-passing GNNs and motivate architectures beyond neighborhood aggregation.Over-smoothing makes node representations converge, whereas over-squashing compresses distant-node messages through graph bottlenecks.
  • Transformers can capture interactions between any node pair in one self-attention layer while avoiding the structural inductive bias of GNNs.
  • Shortest-path positional encodings can assign identical encodings to nodes with different local structures, such as a node forming a triangle with its neighbors.
  • SAT introduces flexible structure-aware self-attention that explicitly considers graph structure and captures structural interactions between nodes.
  • The proposed methods generate subgraph representations, achieve at least the expressiveness of those representations theoretically, and outperform state-of-the-art GNNs and Transformers on five benchmarks.

2. Related Work

Related work studies message-passing GNNs, absolute encodings, and graph Transformers. SAT combines local and global information through absolute encoding and a novel structure-aware relative encoding.

  • Message-passing GNNs became a predominant framework for graph representation learning after reformulating earlier graph neural networks around message passing.
  • Absolute encoding adds or concatenates positional or structural graph representations to input node features, addressing limitations in GNN expressiveness.
  • Graph Transformer methods use absolute encodings such as Laplacian eigenvectors while differing in whether attention operates locally or across the full graph.
  • Relative encoding incorporates positional or distance relationships directly into self-attention, including methods based on graph kernels and diffusion kernels.
  • SAT combines absolute encoding with a novel structure-aware relative encoding to incorporate both local and global graph information.

3. Background

Transformers infer node relations from attributes and can overlook graph structure, motivating positional or structural encodings. SAT reframes attention as kernel smoothing and extends it to compare local subgraphs.

  • A graph is represented as G = (V, E, X), with node attributes x_u and an attribute matrix X for n nodes.
  • Unlike GNNs, Transformers infer node relations from attributes and treat the graph as a multiset of nodes rather than explicitly using its structure.
  • Transformer self-attention is permutation-equivariant, so nodes with identical attributes receive identical representations regardless of graph location or surrounding structure.
  • Absolute encoding adds positional or structural graph representations to input features, but generally does not measure structural similarity between nodes and their neighborhoods.
  • Self-attention can be rewritten as a kernel smoother with a learned exponential kernel on node features and a linear value function.
  • A diffusion-kernel relative encoding captures positional similarity, but remains position-aware rather than structure-aware.

4. Structure-Aware Transformer

SAT incorporates local subgraph representations into Transformer self-attention so node interactions reflect both attributed and structural similarity. Its extractors support theoretical expressiveness guarantees and flexible trade-offs between computation and structural detail.

  • 4.1. Structure-Aware Self-Attention: Structure-aware self-attention replaces feature-only similarity with a kernel comparing subgraph representations centered at nodes.The generalized kernel accounts for local substructures in addition to node attributes.
  • 4.1. Structure-Aware Self-Attention: SAT computes structure-aware queries and keys using an extractor that maps each node and graph to a vector representation of a centered subgraph.The framework can combine with any subgraph representation model, including GNNs and differentiable graph kernels.
  • 4.1. Structure-Aware Self-Attention: The k-subtree GNN extractor represents a rooted k-subtree efficiently, with small k performing well without over-smoothing or over-squashing.Its expressiveness remains bounded by the Weisfeiler–Lehman limitation of message-passing GNNs.
  • 4.1. Structure-Aware Self-Attention: The k-subgraph GNN extractor pools updated representations across each node’s entire k-hop subgraph, yielding more expressive representations than the k-subtree extractor.It requires enumerating all k-hop subgraphs and therefore scales less well to large datasets.
  • 4.1. Structure-Aware Self-Attention: SAT can also learn trainable hidden anchor subgraphs whose matching to node neighborhoods identifies predictive structural motifs and supports model interpretation.Visualizing the learned hidden graphs can provide insights into dataset-specific structural motifs.
  • 4.3. Combination with Absolute Encoding: Absolute positional encodings do not guarantee structurally similar representations, whereas SAT bounds representation distance using overall graph features and subgraph similarity.Theorem 2 further states that suitable parameters can make output representations distinguishable under specified attribute and subgraph conditions.

5. Experiments

Experiments across five graph and node prediction benchmarks show that SAT consistently outperforms strong GNN and graph Transformer baselines. Ablations indicate that structure-aware attention is the main performance driver, while SAT also improves its underlying GNNs and yields more focused attention patterns.

  • Comparison to state-of-the-art methods: SAT achieves state-of-the-art performance on graph and node classification tasks, outperforming state-of-the-art graph Transformers and sparse GNNs.The evaluation covers five medium-to-large benchmark datasets: ZINC, CLUSTER, PATTERN, OGBG-PPA, and OGBG-CODE2.
  • Comparison to state-of-the-art methods: On CODE2, SAT models outperform state-of-the-art methods by a large margin despite relatively few parameters and minimal hyperparameter tuning.The authors report that this result would place SAT first on the OGB leaderboard.
  • SAT models versus sparse GNNs: Both k-subtree and k-subgraph SAT consistently bring large performance gains over the sparse GNN used to extract subgraph representations.Across different GNN extractors, SAT acts as a systematic enhancer; PNA performs best with SAT, and k-subgraph SAT is at least as good as k-subtree SAT in almost all cases.
  • Hyperparameter studies: Structural information produces substantial improvement on ZINC, with optimal performance around k = 3 for both k-subtree and k-subgraph extractors.For k-subtree, performance deteriorates beyond k = 4, whereas k-subgraph is less affected; k-subtree remains more computationally efficient and scalable.
  • Hyperparameter studies: Absolute positional encoding improves SAT, but its gain is far smaller than that from structure-aware attention, while readout choice has very little impact.The ablation compares SAT with and without RWPE or Laplacian positional encoding and evaluates multiple readout methods on ZINC.
  • Model interpretation: SAT learns sparser and more informative attention than a Transformer with RWPE, assigning more attention to known mutagenic motifs in Mutagenicity molecules.The Transformer misses some important H atoms in NH2 groups, whereas SAT highlights only the H atoms in that group.

6. Discussion

SAT incorporates structural information into Transformer attention and reports strong empirical performance with minimal hyperparameter tuning. Its main limitations are higher memory requirements for k-subgraph SAT and the quadratic complexity inherited from self-attention.

  • Discussion: SAT incorporates structural information into the Transformer and addresses limitations of absolute encoding.The model also provides better interpretability than the Transformer.
  • Limitations: k-subgraph SAT requires more memory than k-subtree SAT, which can restrict use when high-memory GPUs are unavailable.SAT also inherits the Transformer’s quadratic self-attention complexity.
  • Discussion: SAT can combine with any GNN, while the framework can also accommodate other structure extractors such as differentiable graph kernels.Future work includes structure extractors more expressive than the 1-WL test and alternatives beyond GNNs.
  • Future work: Reducing self-attention’s high memory cost and time complexity is identified as an important direction for future work.The discussion points to linear Transformers as an approach with linear time and space complexity.
  • Discussion: SAT’s theoretical analysis bounds distances between structure-aware attention representations using assumptions on the mapping and structure extractor.The proof combines bounds on separate terms and then takes the infimum over node permutations.
  • Discussion: The analysis defines each node’s subgraph representation and treats the collection of these representations as multisets compared by a matching metric.The notation uses h_w = ϕ(w, G) and a permutation set between node sets.

B.2. Expressivity Analysis

The expressivity analysis establishes conditions under which structure-aware attention distinguishes node representations. It relies on assumptions about node attributes, subgraph representations, and the mapping used in attention.

  • Expressivity Analysis: The mapping f is assumed to be continuous and approximated in practice by an MLP through the universal approximation theorem.The proof selects parameters and mappings to exploit differences in node attributes and subgraph representations.
  • Expressivity Analysis: Theorem 2 states that structure-aware attention can distinguish two node representations when the graph contains suitably unique attributes and subgraph representations.The theorem assumes countable node attributes and existence of nodes unique in attributes and subgraph representation.
  • Expressivity Analysis: The proof analyzes injectivity of the original dot-product attention with respect to the query.It uses a contrapositive argument based on equality of attention outputs for all parameters and mappings.
  • Expressivity Analysis: Setting W_Q = W_K = 0 makes all softmax coefficients identical, allowing the proof to compare multisets of node attributes.The argument assumes equal graph sizes, adding virtual isolated nodes to the smaller graph if necessary.
  • Expressivity Analysis: The exponential dot-product kernel is treated through its feature mapping into a reproducing kernel Hilbert space.Its universality supports density arguments used in the expressivity proof.

C.2. Datasets Description

The experiments use five graph datasets spanning molecular regression, synthetic node classification, protein-network graph classification, and source-code subtoken classification. Standard splits and evaluation metrics are used for the datasets.

  • Datasets: The experimental suite includes ZINC, CLUSTER, PATTERN, OGBG-PPA, and OGBG-CODE2, each evaluated with its standard protocol.The datasets use standard train, validation, and test splits and corresponding evaluation metrics.
  • ZINC: ZINC is a molecular graph-regression dataset for predicting constrained solubility, using a 12K-molecule subset.The experiments follow the dataset’s established splits.
  • Dataset-related setup: Table 4 summarizes SAT hyperparameters across the datasets and defines RWPE-p as a p-dimensional positional representation from p random-walk steps.The table is a hyperparameter reference rather than a dataset description.
  • PATTERN and CLUSTER: PATTERN and CLUSTER are synthetic Stochastic Block Model datasets for node classification.PATTERN detects membership in a predefined pattern, whereas CLUSTER identifies graph communities.
  • OGBG-PPA: OGBG-PPA contains protein-protein association networks classified into 37 species-category classes.Nodes represent proteins, edges represent associations, and edge attributes include association information such as co-expression.
  • OGBG-CODE2: OGBG-CODE2 contains Python source-code abstract syntax trees, with the task of classifying subtokens comprising method names.The experiments use the standard splits provided for the dataset.

C.3. Hyperparameter Choices and Reproducibility

SAT experiments use limited hyperparameter search, standard Transformer-style optimization, and comparisons across structure extractors and readout methods. The reported additional results find PNA consistently strongest among tested GNN extractors, while readout choice has little effect.

  • Hyperparameter choice: Only dropout rate and subgraph size k are tuned, with k selected from {1, 2, 3, 4}; other hyperparameters are fixed for simplicity.RWPE is fixed for ZINC, PATTERN, and CLUSTER, and validation sets select dropout and k.
  • Optimization: Models use AdamW, Transformer-style warm-up, task-dependent L1 or cross-entropy loss, and dataset-specific learning-rate schedulers.The Transformer scheduler is used on ZINC, PATTERN, and CLUSTER, while cosine scheduling is used on the larger OGB datasets.
  • Computation: Table 5 reports parameter counts and per-epoch training times for k-subtree SAT models using the selected hyperparameters and different base GNNs.SAT uses fewer parameters on OGB datasets than most state-of-the-art methods, according to the accompanying passage.
  • Additional ZINC results: Table 6 compares SAT structure extractors and readout methods on ZINC using test MAE.The caption identifies test MAE as the reported metric.
  • Additional experiments: Additional experimental results are provided for ZINC, OGBG-PPA, and OGBG-CODE2, and OGBG-PPA comparisons use 10 random-seed runs.Table 7 summarizes k-subtree SAT with different GNNs against state-of-the-art methods on OGBG-PPA.
  • Additional ZINC results: PNA-based SAT models consistently outperform other tested GNN extractors, while readout methods have little impact on prediction performance.The comparison covers different structure extractors and readout methods on ZINC.

C.4.3. ADDITIONAL RESULTS ON OGBG-CODE2

This section reports additional OGBG-CODE2 results and identifies the tables summarizing comparisons between k-subtree SAT variants, different GNNs, and state-of-the-art methods.

  • C.4.3. ADDITIONAL RESULTS ON OGBG-CODE2: Table 8 summarizes k-subtree SAT results with different GNNs against state-of-the-art methods on OGBG-CODE2.The results are computed from 10 runs using different random seeds.
  • C.4.3. ADDITIONAL RESULTS ON OGBG-CODE2: Table 7 presents a comparison of SAT and state-of-the-art methods on OGBG-PPA.All results are computed from 10 different runs.
  • C.4.3. ADDITIONAL RESULTS ON OGBG-CODE2: Table 8 presents a comparison of SAT and state-of-the-art methods on OGBG-CODE2.All results are computed from 10 different runs.
  • C.4.3. ADDITIONAL RESULTS ON OGBG-CODE2: The section also includes implementation details about model visualization.

D.1. Dataset and Training Details

The study uses stratified splits of the Mutagenicity molecular-graph dataset and trains a two-layer vanilla Transformer with RWPE for visualization.

  • D.1. Dataset and Training Details: Mutagenicity contains 4337 molecular graphs labeled according to their mutagenic effect.
  • D.1. Dataset and Training Details: The dataset is randomly split into train, validation, and test sets in an 80/10/10 stratified proportion.
  • D.1. Dataset and Training Details: A two-layer vanilla Transformer using RWPE is trained first for the visualization setup.
  • D.1. Dataset and Training Details: The visualization model uses hidden dimension 64, 8 attention heads, and CLS pooling as the readout method.

D.2. Additional Results

The additional results include attention visualizations comparing SAT with a vanilla Transformer and discussion of SAT’s relationship to subgraph networks and graph pooling.

  • D.2. Additional Results: SAT generally learns sparser and more informative [CLS] attention weights than a vanilla Transformer on Mutagenicity examples.The text notes this pattern even for a very large graph shown in the left panel of the middle row.
  • D.2. Additional Results: The visualization examples use atom symbols including C, O, Cl, H, N, F, Br, S, P, I, Na, K, Li, and Ca.
  • D.2. Additional Results: Figure 5 compares [CLS] attention weights learned by SAT with those learned by a classic Transformer using RWPE.The SAT weights appear in the middle column and the Transformer weights in the right column.
  • D.2. Additional Results: The section clarifies SAT’s relationship and differences with Subgraph Neural Networks and the broader topic of graph pooling.
  • D.2. Additional Results: Subgraph Neural Networks generate representations at the subgraph level for subgraph prediction, whereas SAT models structural interactions between nodes in a Transformer.The passage contrasts SNN’s explicit use of position, neighborhood, and structural information with SAT’s node representations.
  • D.2. Additional Results: GNNs traditionally incorporate structural information through neighborhood aggregation, while local pooling provides a supplemental approach typically based on graph clustering.Local pooling coarsens the adjacency matrix by clustering nodes and replacing the adjacency representation with cluster assignments.
Loading 2202.03036v3…