PaperlayerView exact revision on arXiv

Source-linked AI summary

SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations

Qitian Wu, Wentao Zhao, Chenxiao Yang, Hengrui Zhang, Fan Nie, Haitian Jiang, Yatao Bian, Junchi Yan

arXiv:2306.10759v5cs.LGcs.AIcs.SI

TL;DR

Large-graph representation learning must model interdependent nodes while controlling the computational burden of global attention and deep Transformer architectures. The paper proposes SGFormer, a one-layer linear global-attention model combined with graph propagation, and reports competitive benchmark performance, large speedups, and scaling to a 0.1B-node graph.

  • Problem

    Large graphs require effective node representations under limited computation, while existing deep multi-head Transformer designs hinder scalability and can be vulnerable to overfitting with scarce labels.

  • Method

    SGFormer uses single-layer global attention with linear complexity, combines it with GNN propagation, and avoids positional encodings, feature or graph preprocessing, and augmented loss.

  • Results

    SGFormer achieves highly competitive performance across diverse node property prediction benchmarks, scales to ogbn-papers100M with 0.1B nodes, and reports 30x acceleration over state-of-the-art Transformers.

  • Takeaways & Limitations

    A lightweight one-layer attention architecture can provide competitive large-graph representations while retaining efficient scaling and global information propagation.

  • Takeaways & Limitations

    The study assumes training and testing data come from identical distributions and leaves out-of-distribution generalization for future work.

Abstract

from arXiv · show

Learning representations on large-sized graphs is a long-standing challenge due to the inter-dependence nature involved in massive data points. Transformers, as an emerging class of foundation encoders for graph-structured data, have shown promising performance on small graphs due to its global attention capable of capturing all-pair influence beyond neighboring nodes. Even so, existing approaches tend to inherit the spirit of Transformers in language and vision tasks, and embrace complicated models by stacking deep multi-head attentions. In this paper, we critically demonstrate that even using a one-layer attention can bring up surprisingly competitive performance across node property prediction benchmarks where node numbers range from thousand-level to billion-level. This encourages us to rethink the design philosophy for Transformers on large graphs, where the global attention is a computation overhead hindering the scalability. We frame the proposed scheme as Simplified Graph Transformers (SGFormer), which is empowered by a simple attention model that can efficiently propagate information among arbitrary nodes in one layer. SGFormer requires none of positional encodings, feature/graph pre-processing or augmented loss. Empirically, SGFormer successfully scales to the web-scale graph ogbn-papers100M and yields up to 141x inference acceleration over SOTA Transformers on medium-sized graphs. Beyond current results, we believe the proposed methodology alone enlightens a new technical path of independent interest for building Transformers on large graphs.

1 Introduction

Large-graph representation learning must capture dependencies among massive numbers of nodes under limited computation, but Transformer designs built around deep multi-head attention challenge scalability. SGFormer argues that a single-layer attention model can remain competitive while scaling efficiently across graph sizes.

  • Large-graph learning seeks effective node representations under limited time and space budgets.
  • Global attention captures dependencies beyond graph neighborhoods, supporting strong performance on small-graph applications.
  • Deep multi-head attention produces large, data-hungry models that are difficult to scale to large graphs.
  • Quadratic all-pair attention and exponentially growing computation graphs make deep Transformers resource-intensive for million-node graphs.
  • SGFormer uses single-layer, single-head global attention with linear node complexity across twelve benchmarks spanning thousands to billions of nodes.
  • 37x/141x speedup in training/inference time costs is reported over scalable Transformers on medium-sized graphs, while ogbn-papers100M reaches 0.1B nodes.
  • A one-layer attention model can match the denoising effect of multi-layer attention and perform a steepest descent on the associated objective.

2 Preliminary and Related Work

Graph learning represents partially labeled nodes in graphs ranging from thousands to billions of nodes. GNNs aggregate local neighborhoods, whereas Graph Transformers aggregate globally but face scalability and overfitting challenges on massive graphs with scarce labels.

  • A graph contains N nodes, adjacency A, node features, and partially labeled nodes; representation learning produces embeddings for downstream tasks.
  • Graph Neural Networks: GNNs recursively aggregate embeddings from a node’s receptive field, typically its first-order neighbors, to integrate local information.
  • Graph Transformers: Graph Transformers generalize message passing to a densely connected graph by aggregating all node embeddings for each node update.
  • Graph Transformers: Global attention captures long-range interactions and unobserved potential links that may not appear in the input graph structure.
  • Graph Transformers: Deep multi-head Transformer architectures challenge scaling to billion-node graphs and may overfit when labeled nodes are scarce.

3 Simplifying and Empowering Transformers on Large Graphs

SGFormer combines a linear-complexity global attention mechanism with graph-based propagation in a lightweight architecture. It processes full graphs or sampled mini-batches and produces representations using both arbitrary-node interactions and local structural information.

  • SGFormer introduces simple global attention that captures implicit dependencies among nodes with linear complexity.
  • Input Layer: Input features are mapped into latent node embeddings before attention computation and propagation.
  • Simple Global Attention: One-layer global attention propagates information between arbitrary node pairs through a densely connected attention graph.
  • Simple Global Attention: The linear attention formulation combines all-pair attentive propagation with self-loop propagation, preserving centered-node information.
  • Simple Global Attention: O(N) attention computation is more efficient than the O(N^2) Softmax attention used in original Transformers.
  • Incorporation of Structural Information: SGFormer combines attention output with GNN-propagated embeddings before mapping the result to task-specific predictions.
  • Complexity Analysis: O(N + E) overall complexity follows from linear attention and the GNN module, enabling scaling on sparse graphs.
  • Scaling to Larger Graphs: Random mini-batch partitioning enables arbitrarily large graphs, while linear attention supports large batches for informative within-batch global interactions.

4 Comparison with Existing Models

Existing graph Transformers often add positional or edge embeddings, augmented losses, and deep multi-head attention, creating preprocessing and scalability burdens for large graphs.

  • Architectures: Existing models may use edge or positional embeddings, such as Laplacian features, degree centrality, or Weisfeiler–Lehman labeling, to capture graph information.These additions are architectural mechanisms for incorporating graph structure.
  • Architectures: Positional embeddings can require preprocessing with complexity up to O(N 3), creating time and memory costs on large graphs.
  • Architectures: Some models use augmented training losses, such as edge regularization, which can complicate optimization.
  • Architectures: Existing Transformers typically stack deep multi-head attention layers, while most have been optimized for graph classification on small graphs.

5 Empirical Evaluation

SGFormer is evaluated on medium- and large-sized node property prediction benchmarks, with comparisons spanning expressive GNNs, scalable GNNs, and graph Transformers. Results emphasize competitive accuracy, scalability to ogbn-papers100M, efficiency gains, and the limited benefit of deeper attention.

  • Experimental Setup: SGFormer is tested on graphs ranging from 2K to 0.1B nodes, including medium-sized citation and heterophilic graphs and large-scale networks.The evaluation includes datasets such as cora, pubmed, Amazon2M, pokec, ogbn-proteins, and ogbn-papers100M.
  • Experimental Setup: SGFormer is compared with expressive GNNs, scalable GNNs, neighbor-sampling models, and NodeFormer across node property prediction benchmarks.NodeFormer is identified as the main scalable graph Transformer competitor.
  • Results on Large-sized Graphs: SGFormer yields consistently superior results across five large-graph datasets and outperforms NodeFormer across all reported cases.The results are attributed to global attention learning implicit inter-dependencies beyond the input graph structures.
  • Efficiency and Scalability: SGFormer’s training time and GPU memory scale linearly with graph size, while Softmax attention runs out of memory at 40K nodes.The scalability test samples Amazon2M subsets from 10K to 100K nodes; SGFormer uses 1.5GB memory at 40K nodes.
  • Further Discussions: Adding attention layers provides no considerable performance boost, can reduce performance on actor, and increases training time.One-layer SGFormer attention consistently matches highly competitive multi-layer performance.

6 Theoretical Justifications

The analysis interprets attention propagation as optimization for adaptive graph-signal denoising. It shows that a multi-layer attention model can be replaced by one layer with equivalent global aggregation and potentially less redundancy.

  • Optimization interpretation: One Transformer layer is interpreted as an optimization step for a graph-signal denoising problem that adaptively smooths node features using global information.The denoised features combine initial and global information under a specified objective.
  • Illustration: Figure 5 illustrates the theoretical equivalence between multi-layer and one-layer attention models under the smoothness criteria.The illustration connects the equivalence result to potential redundancy savings.
  • Optimization interpretation: The attention update is equivalent to gradient descent on an objective balancing local and global smoothness.The parameter λ controls the trade-off between the two smoothness criteria.
  • Multi-layer equivalence: For any K-layer attention model, there exists a single-layer model that produces the same global aggregation behavior.The result is established for arbitrary positive integer K under the stated layer-wise updating rule.
  • Multi-layer equivalence: Multi-layer attention applies descent steps to layer-dependent objectives, whereas one layer pursues a fixed objective in one step.The one-layer model therefore avoids potential redundancy while retaining the same denoising effect.

7 Conclusions and Outlooks

SGFormer combines simple Transformer-style global attention with a vanilla GCN for scalable large-graph representation learning. It performs competitively across graph benchmarks, scales to a 0.1B-node graph, and leaves distribution-shift and task-specific applications for future work.

  • Conclusions: SGFormer combines a one-layer attention model with a vanilla GCN and achieves highly competitive performance across node property prediction benchmarks spanning thousands to billions of nodes.The evaluation covers diverse large-graph representation-learning settings.
  • Conclusions: 30x acceleration over state-of-the-art Transformers is reported on medium-sized graphs.The paper attributes this result to SGFormer’s simple and lightweight architecture.
  • Conclusions: SGFormer scales smoothly to an extremely large graph with 0.1B nodes.This supports the paper’s focus on scalability as a bottleneck in large-graph learning.
  • Limitations and outlook: The study assumes that training and testing data come from an identical distribution, leaving out-of-distribution generalization for future work.The authors identify distribution shifts as an under-explored area.
  • Limitations and outlook: The evaluation mainly uses benchmark settings, while domain-specific applications may require additional task-dependent designs.Recommender systems, circuit designs, and combinatorial optimization are left for future exploration.

A.1 Proof for Theorem 1

The proof represents each attention layer as a propagation matrix and composes these matrices across layers. It then constructs a single attention matrix whose one-step update reproduces the multi-layer output.

  • Optimization link: The layer update is connected to gradient descent with step size τ/(2λ) on the denoising cost function.The parameter λ weights local and global smoothness criteria.
  • Optimization link: Normalization of the attention scores establishes equivalence between the gradient-descent update and the Transformer layer update.The proof explicitly uses the row-normalization condition P_uv = 1.
  • Single-layer construction: The proof constructs C* from the composed propagation matrix P* and a step parameter τ*.This construction produces a single-layer attention representation matching the K-layer output.
  • Propagation representation: Each attention layer is represented by a propagation matrix, allowing the layer-wise update to be written in matrix form.The proof starts from the sequence of attention matrices and node embeddings generated by the K-layer model.
  • Propagation representation: Stacking K propagation layers yields the output embeddings as a product of propagation matrices applied to the initial embeddings.The composed product is denoted P* in the proof.

B Dataset Information

The experiments use 12 publicly available real-world datasets for node property prediction, covering citation, social, biological, web, and product graphs. Their sizes and task settings range from small citation networks to the 111-million-node ogbn-papers100M benchmark.

  • Dataset overview: The benchmark suite contains 12 publicly available real-world datasets for node property prediction.Dataset information is summarized in Table 5.
  • Citation and web graphs: Cora, Citeseer, and Pubmed are citation networks whose document nodes are classified by academic topic from bag-of-words features.The semi-supervised split uses 20 training instances per class, 500 validation instances, and 1,000 test instances.
  • Citation and web graphs: Actor, Squirrel, and Chameleon use Wikipedia-related networks with node classification tasks based on actor or page features.Actor has relatively low homophily, while Squirrel and Chameleon classify pages by average monthly traffic.
  • Social graphs: Deezer-Europe and Pokec are social networks used for gender prediction with 50%/25%/25% and 10%/10%/80% train/validation/test splits, respectively.Their nodes represent users connected by friendships or social-network links.
  • Large-scale benchmarks: The suite also includes protein-association, arXiv citation, and Amazon co-purchasing graphs for protein-function, subject-area, and product-category prediction.These datasets are represented by ogbn-proteins, ogbn-arxiv, and Amazon2M.
  • Large-scale benchmarks: ogbn-papers100M contains 111 million paper nodes, with approximately 1.5 million labeled papers classified into 172 arXiv subject areas.The dataset uses the public OGB split and is described as nearly the largest public graph benchmark.

C Implementation Details

SGFormer combines an input MLP, one-layer global attention, a graph convolutional network, and an output layer. Training and evaluation use full-graph or mini-batch schemes according to graph scale, with validation-based model selection and repeated trials.

  • Model architecture: SGFormer comprises input, global attention, GNN, and output modules.The input layer transforms features into latent node embeddings before global and graph-based processing.
  • Model architecture: The one-layer global attention maps initial embeddings to representations absorbing all-pair information among nodes.It receives Z(0) and outputs updated embeddings Z through the attention mechanism.
  • Model architecture: The graph convolutional network combines graph structure and initial embeddings with attention outputs as ZO = (1 − α)Z + αGN(Z(0), A).The resulting ZO is the final node representation used for prediction.
  • Model architecture: The output feed-forward layer maps node representations to predicted labels for classification or regression losses.The loss uses cross-entropy for classification and mean square error for regression.
  • Training and evaluation: Full-graph training is used for relatively small graphs, while larger graphs use mini-batch training.Full-graph inference computes attention across the whole graph; mini-batch training partitions shuffled nodes into batches and extracts corresponding subgraphs.
  • Training and evaluation: Evaluation selects the highest-validation-result model, runs fixed epoch budgets by graph scale, and reports mean and variance over five trials.Hyperparameters are selected using validation performance and grid search over shared and model-specific spaces.

D More Empirical Results

The supplementary experiments visualize SGFormer’s learned attention matrices across datasets. The resulting attention graphs differ substantially from the input graphs, suggesting that global attention captures informative data patterns for downstream prediction.

  • Attention patterns: SGFormer’s attention graphs are generally quite different from the corresponding input graphs across datasets.The comparison comes from visualized attention matrices shown in Figure 6.
  • Attention patterns: The observed differences imply that global attention captures informative patterns from data that help downstream prediction.This interpretation is reported as the motivation for examining the learned attention matrices.
Loading 2306.10759v5…