SlowGuess/NeurIPS_2023_OCR huggingface dataset:
SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations
Qitian Wu $^{1}$ , Wentao Zhao $^{1}$ , Chenxiao Yang $^{1}$ , Hengrui Zhang $^{2}$ , Fan Nie $^{1}$ , Haitian Jiang $^{3}$ , Yatao Bian $^{4}$ , Junchi Yan $^{1*}$
<sup>1</sup> Dept. of Computer Science and Engineering & MoE Key Lab of AI, Shanghai Jiao Tong University
$^{2}$  Dept. of Computer Science, University of Illinois at Chicago
$^{3}$  Courant Institute, New York University
4 Tencent AI Lab
{echo740,permanent,chr26195,youluo2001,yanjunchi}@sjtu.edu.cn,
hzhan55@uic.edu, hatian.jiang@nyu.edu, yatao.bian@gmail.com
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 quadratic 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 at the minimal cost of one propagation layer and linear complexity w.r.t. node numbers. Moreover, SGFormer requires none of positional encodings, feature/graph pre-processing or extra 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. The codes are publicly available at https://github.com/qitianwu/SGFormer.
Learning on large graphs that connect interdependent data points is a fundamental problem in machine learning, with diverse applications in social and natural sciences [45; 15; 48; 43; 14]. One key challenge is to obtain effective node representations, especially under limited computation budget (e.g., time and space), that can be efficiently utilized for various downstream tasks.
Recently, Transformers have emerged as a popular class of foundation encoders for graph-structured data, by treating nodes in the graph as input tokens, and show highly competitive performance on
*Corresponding author. This work was partly supported by National Key Research and Development Program of China (2020AAA0107600), NSFC (62222607) and STCSM (22511105100).
graph-level tasks [10; 66; 58; 63; 39] and node-level tasks [57; 54]. The global attention mechanism in Transformers [50] can capture implicit inter-dependencies among nodes that are not embodied by input graph structures, but could potentially make a difference in data generation (e.g., the undetermined structures of proteins that lack known tertiary structures [37; 64]). This advantage provides Transformers with the desired expressivity and leads to superior performance over graph neural networks in small-graph-based applications [64; 3; 11; 31; 36].
A concerning trend in current architectures is their tendency to automatically adopt the design philosophy of Transformers used in vision and language tasks [8; 2; 9]. This involves stacking deep multi-head attention layers, which results in large model sizes and a data-hungry nature. However, this design approach poses a significant challenge for Transformers to scale to large graphs.
- Due to the global all-pair attention mechanism, the time and space complexity of Transformers often scales quadratically with respect to the number of nodes, and the computation graph grows exponentially as the number of layers increases. Thereby, training deep Transformers for large graphs with node numbers in the millions can be extremely resource-intensive and may require delicate techniques for partitioning the inter-connected nodes into smaller mini-batches in order to mitigate computational overhead [57; 54; 35; 5].
- In small-graph-based tasks such as graph-level prediction for molecular property [18], where each instance is a graph and there are typically abundant labeled graph instances in a dataset, large Transformers may have sufficient supervision for generalization. However, in large-graph-based tasks such as node-level prediction for protein functions [19], where there is usually only a single graph and each node is an instance, labeled nodes may be relatively limited. This exacerbates the vulnerability of large Transformers to overfitting in such cases.
Our Contributions. This paper presents our initial attempt to reassess the need for deep attention when learning representations on large graphs. Critically, we demonstrate that a single-layer, single-head attention model can perform surprisingly competitive across twelve graph benchmarks, with node numbers ranging from thousands to billions. Our model, framed as Simplified Graph Transformer (SGFormer), is equipped with a simple global attention that scales linearly w.r.t. the number of nodes. Despite its simplicity, this model retains the necessary expressivity to learn all-pair interactions without the need for any approximation. Furthermore, SGFormer does not require any positional encodings, feature or graph pre-processing, extra loss functions, or edge embeddings.
Experiments show that SGFormer achieves highly competitive performance in an extensive range of node property prediction datasets, which are used as common benchmarks for model evaluation w.r.t. the fundamental challenge of representation learning on graphs, when compared to powerful GNNs and state-of-the-art graph Transformers. Furthermore, SGFormer achieves up to  $37\mathrm{x} / 141\mathrm{x}$  speedup in terms of training/inference time costs over scalable Transformers on medium-sized graphs. Notably, SGFormer can scale smoothly to the web-scale graph ogbn-papers100M with 0.1B nodes, two orders-of-magnitude larger than the largest demonstration of existing works on graph Transformers.
Additionally, as a separate contribution, we provide insights into why the one-layer attention model can be a powerful learner for learning (node) representations on large graphs. Specifically, by connecting the Transformer layer to a well-established signal denoising problem with a particular optimization objective, we demonstrate that a one-layer attention model can produce the same denoising effect as multi-layer attentions. Furthermore, a single attention layer can contribute to a steepest descent on the optimization objective. These results suggest that a one-layer attention model is expressive enough to learn global information from all-pair interactions.
We denote a graph as  $\mathcal{G} = (\mathcal{V},\mathcal{E})$  where the node set  $\mathcal{V}$  comprises  $N$  nodes, and the edge set  $\mathcal{E} = \{(u,v)\mid a_{uv} = 1\}$  is defined by a symmetric (and usually sparse) adjacency matrix  $\mathbf{A} = [a_{uv}]_{N\times N}$ , where  $a_{uv} = 1$  if nodes  $u$  and  $v$  are connected, and 0 otherwise. Each node has an input feature  $\mathbf{x}_u\in \mathbb{R}^D$  and a label  $y_{u}$ . The nodes in the graph are only partially labeled, forming a node set denoted as  $\mathcal{V}_{tr}\subset \mathcal{V}$  (wherein  $|\mathcal{V}_{tr}|\ll N$ ). Learning representations on graphs aims to produce node embeddings  $\mathbf{z}_u\in \mathbb{R}^d$  that are useful for downstream tasks. The size of the graph, as measured by the number of nodes  $N$ , can be arbitrarily large, usually ranging from thousands to billions.
Graph Neural Networks. The layer-wise updating rule for GNNs [44; 23] can be defined as recursively aggregating the embeddings of neighboring nodes to compute the node representation:
$$
\mathbf {z} _ {u} ^ {(k + 1)} = \eta^ {(k)} \left(\overline {{\mathbf {z}}} _ {u} ^ {(k + 1)}\right), \quad \overline {{\mathbf {z}}} _ {u} ^ {(k + 1)} = \operatorname {A g g} \left(\left\{\mathbf {z} _ {v} ^ {(k)} \mid v \in \mathcal {R} (u) \right\}\right), \tag {1}
$$
where  $\mathbf{z}_u^{(k)}$  denotes the embedding at the  $k$ -th layer,  $\eta$  denotes (parametric) feature transformation, and Agg is an aggregation function over the embeddings of nodes in  $\mathcal{R}(u)$  which is the receptive field of node  $u$  determined by  $\mathcal{G}$ . Common GNNs, such as GCN [23] and GAT [51], along with their numerous successors, typically assume that  $\mathcal{R}(u)$  is the set of first-order neighboring nodes in  $\mathcal{G}$ . By stacking multiple layers (e.g.,  $K$ ) of the updating rule in (1), the model can integrate information from the local neighborhood into the representation. Because GNNs involve neighboring nodes in the computation graph, which exponentially increase w.r.t.  $K$ , training them on large graphs (e.g., with a million nodes) in a full-batch manner can be challenging. To reduce the overhead, GNNs require subtle techniques like neighbor sampling [65], graph partition [7], or historical embeddings [12]. Some works adopt knowledge distillation for inference on sparser graphs [61; 67] to accelerate inference, and more recently, MLP architectures are used to replace GNNs to accelerate training [60; 17].
Graph Transformers. Beyond message passing within local neighborhood, Transformers have recently gained attention as powerful graph encoders [10; 66; 63; 34; 25; 58; 4; 20; 39; 22]. These models leverage global all-pair attention, which aggregates all node embeddings to update the representation of each node. Global attention can be seen as a generalization of GNNs' message passing to a densely connected graph where  $\mathcal{R}(u) = \mathcal{V}$ , and equips the model with the ability to capture implicit dependencies such as long-range interactions or unobserved potential links in the graph. However, the all-pair attention incurs  $O(N^2)$  complexity and becomes a computation bottleneck that limits most Transformers to handling only small-sized graphs (with up to hundreds of nodes). For larger graphs, recent efforts have resorted to strategies such as sampling a small (relative to  $N$ ) subset of nodes for attention computation [16] or using ego-graph features as input tokens [68]. However, these strategies sacrifice the expressivity needed to capture all-pair interactions among arbitrary nodes. Another line of recent works design new attention mechanisms that can efficiently achieve all-pair message passing within linear complexity. One of the pioneering work [57] proposes kernelized Gumbel-Softmax message passing that utilizes random feature maps to approximate the all-pair attention with linear complexity. Another work [54] leverages the principled graph diffusion equation and resorts to the non-local diffusion operator to derive a linear attention network that models all-pair interactions without any approximation.
Another observation is that nearly all of the Transformer models mentioned above tend to stack deep multi-head attention layers, in line with the design of large models used in vision and language tasks [8; 2; 9]. However, this architecture presents challenges for scaling to industry-scale graphs, where  $N$  can reach billions. Moreover, the model can become vulnerable to overfitting when the number of labeled nodes  $|\mathcal{V}_{tr}|$  is much smaller than  $N$ . This is a common issue in extremely large graphs where node labels are scarce [19].
3 Simplifying and Empowering Transformers on Large Graphs
In this section, we introduce our approach, which we refer to as Simplified Graph Transformers (SGFormer). The key component of SGFormer is a simple global attention model that captures the implicit dependencies among nodes with linear complexity.
Input Layer. We first use a neural layer to map input features  $\mathbf{X} = [\mathbf{x}_u]_{u=1}^N$  to node embeddings in the latent space, i.e.,  $\mathbf{Z}^{(0)} = f_I(\mathbf{X})$  where  $f_I$  can be a shallow (e.g., one-layer) MLP. The node embeddings  $\mathbf{Z}^{(0)} \in \mathbb{R}^{N \times d}$  will be used for subsequent attention computation and propagation.
Simple Global Attention. The global interactions in our model are captured by an all-pair attention unit that computes pair-wise influence between arbitrary node pairs. Unlike other Transformers, which often require multiple attention layers for desired capacity, we found that a single-layer global attention is sufficient. This is because through one-layer propagation over a densely connected attention graph, the information of each node can be adaptively propagated to arbitrary nodes within the batch. Therefore, despite its simplicity, our model has sufficient expressivity to capture implicit dependencies among arbitrary node pairs while significantly reducing computational overhead. In
Figure 1: Illustration of the proposed model SGFormer and its data flow. The input graph data entails node features  $\mathbf{X}$  and graph adjacency  $\mathbf{A}$ . For large graphs, we need to use mini-batch sampling that randomly partitions the input graph into mini-batches with smaller sizes. Each mini-batch is composed of the features of the nodes within this mini-batch  $\mathbf{X}_m$  and the local graph adjacency  $\mathbf{A}_m$  (one can also use neighbor sampling as an alternative). The mini-batch data  $(\mathbf{X}_m, \mathbf{A}_m)$  (for large graphs) or the whole graph data  $(\mathbf{X}, \mathbf{A})$  (for small graphs) will be fed into the SGFormer model that is implemented with a one-layer global attention and a GNN network. The model outputs the node representations for final prediction.
specific, we adopt a linear attention function defined as follows:
$$
\mathbf {Q} = f _ {Q} \left(\mathbf {Z} ^ {(0)}\right), \quad \tilde {\mathbf {Q}} = \frac {\mathbf {Q}}{\| \mathbf {Q} \| _ {\mathcal {F}}}, \quad \mathbf {K} = f _ {K} \left(\mathbf {Z} ^ {(0)}\right), \quad \tilde {\mathbf {K}} = \frac {\mathbf {K}}{\| \mathbf {K} \| _ {\mathcal {F}}}, \quad \mathbf {V} = f _ {V} \left(\mathbf {Z} ^ {(0)}\right), \tag {2}
$$
$$
\mathbf {D} = \operatorname {d i a g} \left(\mathbf {1} + \frac {1}{N} \tilde {\mathbf {Q}} (\tilde {\mathbf {K}} ^ {\top} \mathbf {1})\right), \quad \mathbf {Z} = \beta \mathbf {D} ^ {- 1} \left[ \mathbf {V} + \frac {1}{N} \tilde {\mathbf {Q}} (\tilde {\mathbf {K}} ^ {\top} \mathbf {V}) \right] + (1 - \beta) \mathbf {Z} ^ {(0)}, \tag {3}
$$
where  $f_{Q}, f_{K}, f_{V}$  are all shallow neural layers (e.g., a linear feed-forward layer in our implementation),  $\| \cdot \|$  denotes the Frobenius norm,  $\mathbf{1}$  is an  $N$ -dimensional all-one column vector, the diag operation changes the  $N$ -dimensional column vector into a  $N \times N$  diagonal matrix, and  $\beta$  is a hyper-parameter for residual link. In Eqn. 3,  $\mathbf{Z}$  combines the all-pair attentive propagation over  $N$  nodes and the self-loop propagation. The former allows the model to capture the influence from other nodes, while the latter preserves the information of the centered nodes. The computation of Eqn. (3) can be achieved in  $\mathcal{O}(N)$  time complexity, which is much more efficient than the Softmax attention in original Transformers [50] that requires  $\mathcal{O}(N^2)$ . While the Softmax attention possesses provable expressivity [1], its quadratic complexity hinders the scalability for large graphs. Our adopted design reduces the quadratic complexity to  $\mathcal{O}(N)$  and in the meanwhile guarantee the expressivity for learning all-pair interactions. Therefore, it can successfully be applied for learning representations on graphs with a large number of nodes at the cost of moderate GPU memory.
Incorporation of Structural Information. For accommodating the prior information of the input graph  $\mathcal{G}$ , existing models tend to use positional encodings [39], edge regularization loss [57] or augmenting the Transformer layers with GNNs [58]. Here we resort to a simple-yet-effective scheme that combines  $\mathbf{Z}$  with the propagated embeddings by GNNs at the output layer:
$$
\mathbf {Z} _ {O} = (1 - \alpha) \mathbf {Z} + \alpha \operatorname {G N} \left(\mathbf {Z} ^ {(0)}, \mathbf {A}\right), \quad \hat {Y} = f _ {O} \left(\mathbf {Z} _ {O}\right), \tag {4}
$$
where  $\alpha$  is a weight hyper-parameter, and the GN module can be a simple GNN architecture (e.g., GCN [23]) that possesses good scalability for large graphs. The output function  $f_{O}$  maps the final representation  $\mathbf{Z}_{O}$  to prediction depending on specific downstream tasks, and in our implementation  $f_{O}$  is a linear feed-forward layer.
Complexity Analysis. The overall computational complexity of our model is  $\mathcal{O}(N + E)$ , where  $E = |\mathcal{E}|$ , as the GN module requires  $\mathcal{O}(E)$ . Due to the typical sparsity of graphs (i.e.,  $E \ll N^2$ ), our model can scale linearly w.r.t. graph sizes. Furthermore, with only one-layer global attention and simple GNN architectures, our model is fairly lightweight, enabling efficient training and inference.
Scaling to Larger Graphs. For larger graphs that even GCN cannot be trained on using full-batch processing with a single GPU, we can use the random mini-batch partitioning method utilized by [57; 54]. This incurs only negligible additional costs during training and allows the model to scale to arbitrarily large graphs. Moreover, due to the linear complexity of our all-pair attention mechanism, we can employ large batch sizes, which facilitate the model in capturing informative global interactions among nodes within each mini-batch. Our model is also compatible with advanced techniques such as neighbor sampling [65], graph clustering [7], and historical embeddings [12], which we leave for future exploration along this orthogonal research direction. Fig. 1 presents the data flow of the proposed model.
Table 1: Comparison of (typical) graph Transformers w.r.t. required components (PE for positional embeddings, MA for multiple attention heads, AL for augmented training loss and EE for edge embeddings), all-pair expressivity and algorithmic complexity w.r.t. node number  $N$  and edge number  $E$  (often  $E \ll  {N}^{2}$  ). The largest demonstration means the largest graph size used by the papers.
Model	Model Components	All-pair
Expressivity	Algorithmic Complexity	Largest
Demo.
PE	MA	AL	EE	Pre-processing	Training
GraphTransformer [10]	R	R	-	R	Yes	O(N3)	O(N2)	0.2K
Graphorper [63]	R	R	-	R	Yes	O(N3)	O(N2)	0.3K
GraphTrans [58]	-	R	-	-	Yes	-	O(N2)	0.3K
SAT [4]	R	R	-	-	Yes	O(N3)	O(N2)	0.2K
EGT [20]	R	R	R	R	Yes	O(N3)	O(N2)	0.5K
GraphGPS [39]	R	R	-	R	Yes	O(N3)	O(N + E)	1.0K
Gophormer [68]	R	R	R	-	No	-	O(Nsm2)	20K
NodeFormer [57]	R	R	R	-	Yes	-	O(N + E)	2.0M
DIFFormer [54]	-	R	-	-	Yes	-	O(N + E)	1.6M
SGFormer (ours)	-	-	-	-	Yes	-	O(N + E)	0.1B
4 Comparison with Existing Models
In this section, we provide a more in-depth discussion comparing our model with prior art and illuminating its potential in wide application scenarios. Table 1 presents a head-to-head comparison of current graph Transformers in terms of their architectures, expressivity, and scalability. Most existing Transformers have been developed and optimized for graph classification tasks on small graphs, while some recent works have focused on Transformers for node classification, where the challenge of scalability arises due to large graph sizes.
- Architectures. Regarding model architectures, some existing models incorporate edge/positional embeddings (e.g., Laplacian decomposition features [10], degree centrality [63], Weisfeiler-Lehman labeling [66]) or utilize augmented training loss (e.g., edge regularization [20; 57]) to capture graph information. However, the positional embeddings require an additional pre-processing procedure with a complexity of up to  $O(N^3)$ , which can be time- and memory-consuming for large graphs, while the augmented loss may complicate the optimization process. Moreover, existing models typically adopt a default design of stacking deep multi-head attention layers for competitive performance. In contrast, SGFormer does not require any of positional embeddings, augmented loss or pre-processing, and only uses a single-layer, single-head global attention, making it both efficient and lightweight.
- Expressivity. There are some recently proposed graph Transformers for node classification [68; 5] that limit the attention computation to a subset of nodes, such as neighboring nodes or sampled nodes from the graph. This approach allows linear scaling w.r.t. graph sizes, but sacrifices the expressivity for learning all-pair interactions. In contrast, NodeFormer [57] and SGFormer maintain attention computation over all  $N$  nodes in each layer while still achieving  $O(N)$  complexity. However, unlike NodeFormer which relies on random features for approximation, SGFormer does not require any approximation or stochastic components and is more stable during training.
- Scalability. In terms of algorithmic complexity, most existing graph Transformers have  $\mathcal{O}(N^2)$  complexity due to global all-pair attention, which is a critical computational bottleneck that hinders their scalability even for medium-sized graphs with thousands of nodes. While neighbor sampling can serve as a plausible remedy, it often sacrifices performance due to the significantly reduced receptive field [35]. SGFormer scales linearly w.r.t.  $N$  and supports full-batch training on large graphs with up to 0.1M nodes. For further larger graphs, SGFormer is compatible with mini-batch training using large batch sizes, which allows the model to capture informative global information while having a negligible impact on performance. Notably, due to the linear complexity and simple architecture, SGFormer can scale to the web-scale graph ogbn-papers100M (with 0.1B nodes) when trained on a single GPU, two orders-of-magnitude larger than the largest demonstration of NodeFormer [57] and DIFFormer [54].
We apply SGFormer to various datasets of node property prediction tasks, which are used as common benchmarks for evaluating the model's efficacy for learning effective representations and scalability to large graphs. In Sec. 5.1, we test SGFormer on medium-sized graphs (from 2K to 30K nodes) and compare it with an extensive set of expressive GNNs and Transformers. In Sec. 5.2, we scale SGFormer to large-sized graphs (from 0.1M to 0.1B nodes) where its superiority is demonstrated over