PaperlayerView exact revision on arXiv

Source-linked AI summary

Representing Long-Range Context for Graph Neural Networks with Global Attention

Zhanghao Wu, Paras Jain, Matthew A. Wright, Azalia Mirhoseini, Joseph E. Gonzalez, Ion Stoica

arXiv:2201.08821v1cs.LG

TL;DR

GNNs have difficulty capturing long-range dependencies, while existing depth and pooling strategies remain limited. GraphTrans appends a permutation-invariant Transformer to a standard GNN, using global self-attention and a learned <CLS> readout. It achieves state-of-the-art results across several graph-classification tasks, including reported gains on NCI benchmarks.

  • Problem

    GNNs struggle to represent long-range dependencies because increasing depth can cause performance drops and oversmoothing, while pooling methods lack universal applicability.

  • Method

    GraphTrans adds a permutation-invariant Transformer atop a standard GNN to model all pairwise node interactions and aggregate them through a special <CLS> readout token.

  • Results

    GraphTrans achieves state-of-the-art results on several graph-classification tasks across OpenGraphBenchmark and NCI datasets.

  • Takeaways & Limitations

    The results suggest that purely learned, graph-structure-free attention may be suitable for high-level long-range relationships on graphs.

  • Takeaways & Limitations

    The Transformer relies on preceding GNN layers to encode structural information and therefore omits additive positional embeddings.

Abstract

from arXiv · show

Graph neural networks are powerful architectures for structured datasets. However, current methods struggle to represent long-range dependencies. Scaling the depth or width of GNNs is insufficient to broaden receptive fields as larger GNNs encounter optimization instabilities such as vanishing gradients and representation oversmoothing, while pooling-based approaches have yet to become as universally useful as in computer vision. In this work, we propose the use of Transformer-based self-attention to learn long-range pairwise relationships, with a novel "readout" mechanism to obtain a global graph embedding. Inspired by recent computer vision results that find position-invariant attention performant in learning long-range relationships, our method, which we call GraphTrans, applies a permutation-invariant Transformer module after a standard GNN module. This simple architecture leads to state-of-the-art results on several graph classification tasks, outperforming methods that explicitly encode graph structure. Our results suggest that purely-learning-based approaches without graph structure may be suitable for learning high-level, long-range relationships on graphs. Code for GraphTrans is available at https://github.com/ucbrise/graphtrans.

1 Introduction

GraphTrans addresses long-range graph dependencies by appending a permutation-invariant Transformer to a standard GNN. Its learned readout and global attention improve graph-classification performance across several datasets.

  • 1 Introduction: GNNs struggle with long-range dependencies because increasing depth can cause dramatic performance drops and oversmoothing.A K-hop dependency may require K layers, while the expanding receptive field can dilute distant-node signals.
  • 1 Introduction: GraphTrans replaces some explicit long-range structural operations with a learned attention-based module inspired by computer vision.The approach retains graph-based processing for local structure while learning global relationships.
  • 1 Introduction: The architecture adds a permutation-invariant Transformer after a standard GNN stack to compute global, long-range node interactions.Without positional encoding, the Transformer remains permutation-invariant and can be added atop existing GNN backbones.
  • 1 Introduction: GraphTrans achieves state-of-the-art results on several OpenGraphBenchmark and NCI biomolecular graph-classification datasets.The evaluation reports significant OpenGraphBenchmark accuracy improvements and substantial gains on NCI1.
  • 1 Introduction: The authors report that Transformer-based long-range reasoning improves GNN accuracy, especially for large graph-classification tasks.They suggest modeling all pairwise node interactions is particularly important in that setting.
  • 1 Introduction: The Transformer readout uses a special <CLS> token whose output embedding aggregates pairwise interactions into a single graph-classification vector.This learned readout outperforms global pooling, graph-specific pooling, and virtual-node approaches.

2 Related Work

Related work includes hierarchical and global graph pooling, DAGNN-style global aggregation, and Transformer applications that usually retain graph-specific attention or positional structure. GraphTrans differs by using permutation-invariant self-attention as a graph-level readout.

  • Graph Classification: Graph classification requires aggregating node or edge representations into a single graph embedding, using either global or local pooling.Local pooling progressively coarsens graphs, while global pooling directly produces a whole-graph encoding.
  • Graph Classification: Learned and non-learned hierarchical pooling methods have been proposed, but the effectiveness or necessity of graph coarsening remains unclear.Common global alternatives include mean or max pooling and virtual nodes.
  • Graph Classification: DAGNN aggregates over an entire directed acyclic graph within each layer and acts as learned global pooling into sink-node embeddings.GraphTrans is complementary because DAGNN uses global max-pooling over sink nodes for its final graph-level operation.
  • Transformers on Graphs: Prior graph Transformers commonly restrict attention to neighborhoods, whereas GraphTrans uses self-attention to collapse node encodings into one graph encoding.Related methods also use graph-specific positional encodings or alternative weighted aggregation schemes.
  • Efficient Transformers: Efficient Transformer variants target lower computation, memory usage, or resource requirements through architectural modifications and neural architecture search.The related-work discussion names LiteTransformer, Reformer, Performer, and NAS-based approaches.

3 Motivation: Modeling Long-Range Pairwise Interactions

The paper motivates global attention as an alternative to deep GNN stacks and hierarchical pooling for learning long-range graph relationships. GraphTrans lets every node attend to every other node while leaving local structural processing to the GNN.

  • Motivation: Modeling Long-Range Pairwise Interactions: Stacking GNN layers or using hierarchical pooling has not yet consistently improved long-range learning, while scaling single-layer receptive fields remains uncertain on very large graphs.The authors specifically note uncertainty for graphs with thousands of nodes.
  • Motivation: Modeling Long-Range Pairwise Interactions: Computer vision results motivate replacing explicit relational operations with learned attention for long-range correlations.The cited motivation distinguishes strong structural biases for local patterns from less structured modules for long-range patterns.
  • Motivation: Modeling Long-Range Pairwise Interactions: GraphTrans uses a traditional GNN backbone for local representations and a position-free Transformer for global reasoning.The Transformer is permutation-invariant because it omits positional encoding.
  • Motivation: Modeling Long-Range Pairwise Interactions: Every node can attend to every other node, allowing the Transformer to prioritize important relationships rather than nearby nodes.The preceding GNN handles immediate-neighborhood structure before this global attention stage.
  • Motivation: Modeling Long-Range Pairwise Interactions: In Code2 attention visualizations, nodes assign significant weight to distant nodes, including a pair separated by five hops.The <CLS> token also receives broad attention and serves as the graph-level readout.

4 Learning Global Information with GraphTrans

GraphTrans combines a standard GNN stack for local structural representations with a permutation-invariant Transformer for global pairwise reasoning and graph-level readout.

  • Architecture: GraphTrans consists of a GNN subnetwork followed by a Transformer subnetwork for learning global graph information.The two modules are the primary components of the architecture.
  • GNN module: The GNN module computes node representations using neighborhood aggregation, with few assumptions about the specific GNN layers.A generic layer stack updates each node from its neighborhood through a parameterized neural function.
  • Transformer module: The Transformer projects final per-node GNN encodings, applies layer normalization, and processes them without additive positional embeddings.The architecture expects the GNN embeddings to already encode structural information.
  • Readout: GraphTrans appends a learnable <CLS> embedding and uses its Transformer output as the representation of the whole graph.A linear projection followed by softmax produces the prediction from this graph embedding.
  • Readout: The special-token readout generalizes virtual-node readout by learning node-to-node relationships before distilling information into one graph representation.This can avoid the virtual node’s potential information bottleneck.

5 Experiments

GraphTrans was evaluated across biological and chemical graph-classification benchmarks using different GNN backbones and training setups. It improved over competitive baselines, including by 7.1% on NCI1 and 5.1% on NCI109 in the small model.

  • Datasets: GraphTrans was evaluated on graph-classification benchmarks spanning biology and chemistry, including NCI1, NCI109, and Molpcba.NCI1 and NCI109 contain about 4,000 biochemical-compound graphs each, while Molpcba contains 437,929 molecular graphs.
  • Biological benchmarks: 7.1% higher average accuracy on NCI1 and 5.1% higher average accuracy on NCI109 were achieved by the small GraphTrans model over its simple baseline.The experiments averaged test accuracy over 20 random seeds.
  • Biological benchmarks: 1.1% higher accuracy on NCI1 and 8.2% higher accuracy on NCI109 were obtained by GraphTrans (large) over the strong baseline.This model used GIN and four GNN layers instead of eight.
  • Chemical benchmarks: GraphTrans outperformed competitive baselines on Molpcba with two backbone GNN architectures.The table reports validation and test ROC-AUC, while the experiments used GIN as the baseline and encoder.
  • Chemical benchmarks: Transformer alone performed very badly on Molpcba, whereas GraphTrans improved the GIN and GIN-Virtual baselines.The reported result combines local graph structure from the GNN with long-range modeling from the Transformer.

5.3 Computer programming benchmark

On the Code2 programming benchmark, GraphTrans improves graph classification performance over strong baselines, including DAGNN, while remaining compatible with different GNN encoders.

  • The Code2 dataset contains 45,741 abstract syntax trees averaging 125 nodes, with the task of predicting method-name subtokens.
  • The benchmark uses a 4-layer GCN baseline and trains models for 30 epochs with batch size 16 because of the dataset scale.
  • With a frozen pretrained GNN, training only GraphTrans’s Transformer still improves Code2 F1-score, showing additional information can be learned from GNN embeddings.
  • GraphTrans significantly outperforms the state-of-the-art DAGNN on the Code2 leaderboard.
  • GraphTrans with a PNA encoder also significantly improves results, indicating compatibility with the choice of GNN module.

5.4 Transformers can capture long-range relationships

GraphTrans combines local structural representations from a GNN with global, position-agnostic Transformer attention. Ablations show that the Transformer can add information beyond a pretrained GNN.

  • With a frozen pretrained GNN, the Transformer improves validation F1-score by 0.0022 and test F1-score by 0.0042 on Code2.
  • With an unfrozen pretrained GNN, GraphTrans achieves an even higher F1-score while the modules can specialize in local and long-range information.

5.5 Effectiveness of <CLS> embedding

The <CLS> readout learns to aggregate information across graph nodes into a classification embedding. It outperforms alternative sequence-aggregation choices in the reported ablation.

  • The <CLS> token attends to important nodes across the graph to construct the whole-graph representation.
  • The <CLS> embedding is most effective among tested aggregation methods, improving test performance by 0.0275.
  • Concatenating input graph embeddings with Transformer inputs further raises validation and test F1-scores to 0.1670 and 0.1733.

5.6 Scalability

GraphTrans has runtime and memory behavior comparable to GCN in large-graph microbenchmarks, although its training overhead varies by dataset.

  • GraphTrans scales at least as well as GCN as graph size and edge density increase, with similar memory consumption.
  • Both GCN and GraphTrans encounter out-of-memory errors on large dense graphs.
  • GraphTrans trains faster than comparable GCN on NCI1 but is 7-11% slower on OGB-molpcba and OGB-Code2.

5.8 Number of parameters

GraphTrans improves accuracy with only a minor parameter increase overall, while achieving greater parameter efficiency than the GNN baseline on Code2.

  • GraphTrans only increases total parameters marginally for Molpcba and NCI.
  • The Transformer module adds minimal overhead on NCI1, OGBG-Molpcba, and OGBG-Code2.Speedup is measured as training iteration speed compared with a GNN-based model; larger values indicate faster running speed.
  • Overall, GraphTrans achieves improved accuracy with a minor increase in parameters.
  • On Code2, GraphTrans improves test F1 score from 0.1629 to 0.1810 while being substantially more parameter-efficient than the GNN.The Transformer reduces feature dimension before the expensive final prediction layer.

6 Conclusion

GraphTrans augments standard GNN stacks with a permutation-invariant Transformer readout to learn long-range relationships and summarize node interactions into a graph embedding. The framework improves state-of-the-art performance across several graph classification tasks and sometimes outperforms methods encoding domain-specific structure.

  • GraphTrans adds a permutation-invariant Transformer module after standard GNN layers to learn long-range relationships and produce a graph encoding.The Transformer learns pairwise node interactions and summarizes them into a special token’s embedding.
  • GraphTrans improves state-of-the-art performance across graph classification tasks involving program analysis, molecules, and protein association networks.
  • In some cases, GraphTrans outperforms methods that attempt to encode domain-specific structural information.
  • The paper identifies node and edge classification applications and improved Transformer scalability as next directions.
Loading 2201.08821v1…