[Paper Review] Representing Long-Range Context for Graph Neural Networks with Global Attention

Zhanghao Wu, Paras Jain+4 more|arXiv (Cornell University)|Jan 21, 2022

Advanced Graph Neural NetworksComputer Science125 citations

TL;DR

GraphTrans augments a standard GNN with a permutation-invariant Transformer to learn long-range node relationships and introduces a special CLS readout for graph-level prediction, achieving state-of-the-art results on multiple graph classification benchmarks.

ABSTRACT

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.

Motivation & Objective

• Motivate the need to model long-range dependencies in graphs beyond stacking GNN layers or hierarchical pooling.
• Propose a simple, framework-agnostic architecture that combines a GNN backbone with a Transformer for global reasoning.
• Investigate a novel readout mechanism using a <CLS> token to produce a graph-level embedding.
• Demonstrate the effectiveness and generality of the approach across biology, chemistry, and computer programming graphs.

Proposed method

• Use a standard GNN backbone to learn local, short-range structure.
• Apply a permutation-invariant Transformer on top of GNN embeddings to learn all-pairs interactions.
• Project node embeddings into Transformer dimension without positional encodings and process with multi-head self-attention.
• Adopt a special <CLS> token as a learnable readout to produce a single graph embedding for classification.
• Train end-to-end with standard optimization (Adam) and demonstrate scalability and efficiency comparisons.

Experimental results

Research questions

• RQ1 Can a Transformer module applied on GNN embeddings capture long-range dependencies better than deeper GNNs or pooling approaches?
• RQ2 Does a special <CLS> token readout provide a more effective graph embedding than global pooling or virtual node methods?
• RQ3 Is GraphTrans effective across diverse graph domains (biological, chemical, programming) and scalable to large graphs?

Key findings

• GraphTrans achieves state-of-the-art results on several Open Graph Benchmark graph classification tasks.
• On NCI1 and NCI109 biology datasets, GraphTrans (small) improves accuracy over strong baselines by substantial margins (e.g., 81.3% vs 68–82% range in baselines).
• On molpcba, GraphTrans improves over GIN and GIN-Virtual baselines, showing the benefit of combining GNN local structure with global attention.
• On Code2, GraphTrans outperforms the DAGNN state-of-the-art and other baselines, demonstrating effective long-range reasoning in programming graphs.
• Ablations show the <CLS> readout is the most effective aggregation method among simple sequence aggregations, and that the Transformer learns complementary information to the GNN.