An end-to-end attention-based approach for learning on graphs
arXiv:2402.10793 · doi:10.1038/s41467-025-60252-z
Abstract
There has been a recent surge in transformer-based architectures for learning on graphs, mainly motivated by attention as an effective learning mechanism and the desire to supersede handcrafted operators characteristic of message passing schemes. However, concerns over their empirical effectiveness, scalability, and complexity of the pre-processing steps have been raised, especially in relation to much simpler graph neural networks that typically perform on par with them across a wide range of benchmarks. To tackle these shortcomings, we consider graphs as sets of edges and propose a purely attention-based approach consisting of an encoder and an attention pooling mechanism. The encoder vertically interleaves masked and vanilla self-attention modules to learn an effective representations of edges, while allowing for tackling possible misspecifications in input graphs. Despite its simplicity, the approach outperforms fine-tuned message passing baselines and recently proposed transformer-based methods on more than 70 node and graph-level tasks, including challenging long-range benchmarks. Moreover, we demonstrate state-of-the-art performance across different tasks, ranging from molecular to vision graphs, and heterophilous node classification. The approach also outperforms graph neural networks and transformers in transfer learning settings, and scales much better than alternatives with a similar performance level or expressive power.
Now published in Nature Communications
References in corpus (10)
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- The Open Catalyst 2020 (OC20) Dataset and Community Challenges
- Recipe for a General, Powerful, Scalable Graph Transformer
- Set Transformer: A Framework for Attention-based Permutation-Invariant Neural Networks
- OGB-LSC: A Large-Scale Challenge for Machine Learning on Graphs
- GraphNorm: A Principled Approach to Accelerating Graph Neural Network Training
- Weisfeiler and Leman go Machine Learning: The Story so far
- Accurate GW frontier orbital energies of 134 kilo molecules
- FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision
- Graph Inductive Biases in Transformers without Message Passing