KDLGT: A Linear Graph Transformer Framework via Kernel Decomposition Approach

1 Introduction

Recent years have seen significant advancements in the field of graph learning, with notable successes across a variety of domains including social networks [Li et al., 2021; Zhong et al., 2020], molecular graphs [Huang et al., 2020; Wang et al., 2021], and knowledge graphs [Liu et al., 2021; Yasunaga et al., 2021]. One of the key approaches in this field is Graph Neural Networks (GNNs), which have been widely adopted as a powerful embedding approach for various graph learning tasks [Kipf and Welling, 2017; Veliˇckovi´c et al., 2018; Hamilton et al., 2017]. The foundation of GNNs is the local sparse message-passing mechanism, which enables the nodes on the graphs to iteratively exchange messages through the edges connecting them. However, the limitations of the message-passing mechanism have become increasingly apparent in recent research [Xu et al., 2019; Morris et al., 2019; Maron et al., 2019], leading to a series of works [Ying et al., 2021; Zhang et al., 2020; Chen et al., 2022] that have turned to Transformer architectures[Vaswani et al., 2017] in pursuit of new breakthroughs.

∗Corresponding author

Graph Transformers (GTs) represent a successful endeavor to deploy Transformer architectures to graph data. By enabling nodes to attend to all other nodes within the graphs, GTs encode graph structures as a soft inductive bias, rather than the hard-coded message-passing approach. In contrast to GNNs, GTs utilize absolute and relative positional encodings (APEs and RPEs) to characterize graph topological structures, viewing the graphs as complete entities and enabling long-range interactions for nodes. This overcomes limitations inherent in the message-passing paradigm, such as limited expressiveness [Xu et al., 2019], over-smoothing [Alon and Yahav, 2021], and over-squashing [Alon and Yahav, 2021] issues. However, despite these achievements, there remain a plethora of challenges to be addressed in this area of research.

One of the most significant challenges faced by GTs is the poor scalability. This is due to the global attention mechanism, which results in quadratic time and memory complexity with respect to the number of nodes in the graph. This problem is particularly pronounced when utilizing GTs on datasets consisting of larger graphs, such as citation and social network graphs, as limited GPU memory and excessive running time impede their performance. Consequently, in the application scenarios, GTs do not demonstrate a significant advantage over GNNs.

In this work, we present the Kernel Decomposition Linear Graph Transformer (KDLGT), an accelerating framework for building scalable and powerful GTs. Unlike previous model-specific approaches, KDLGT is model-agnostic and aims to provide a general solution for accelerating all GT models. To achieve this, we employ the kernel decomposition approach, which rearranges the matrix multiplication order of self-attention by designing a kernel function decomposition for the softmax function. Under this approach, the way to deal with RPE matrices becomes the key to reducing time complexity. Inspired by the fact that the RPE matrices of undirected graphs are symmetric and can be decomposed into products of low-dimensional matrices, we categorize RPE matrices into three types, which cover the most commonly used RPEs. Using KDLGT, we select two representative RPEs (such as Shortest Path Distance) and design our models. Additionally, we provide a theoretical analysis of the difference between the KDLGT framework and the traditional self-attention method, and prove that this gap can be effectively bounded.

We conduct experiments on 10 real-world datasets to demonstrate the superior performance of our proposed KDLGT framework. Additionally, we evaluate KDLGT on a series of synthetic graphs with varying scales to verify its efficiency. The experimental results show that KDLGT not only improves learning efficiency significantly but also preserves precision performance effectively and achieves stateof-the-art results on a variety of datasets. Furthermore, as anticipated, the acceleration ratio increases as the graph size increases, indicating the strong scalability of KDLGT. In particular, KDLGT can achieve an acceleration ratio of approximately 10 on graphs of certain sizes. The contributions of this paper are listed as follows:

• We propose the KDLGT framework, which successfully reduces the quadratic complexity of GTs to linear and improves the scalability of GTs greatly.
• We provide a tight upper bound of the difference gap between KDLGT and GTs theoretically to illustrate KDLGT is a well-defined approximation of GTs.
• We conduct experiments on both real-world and synthetic datasets. The experimental results indicate that KDLGT can significantly improve the learning efficiency while preserving precision performance of GTs well.

2 RelatedWorks

2.1 Fast Transformers

There have been a lot of works attempting to improve the efficiency of Transformer models. During the earliest period, researchers tend to restrict the context of self-attention to predefined, fixed patterns, thus limiting the size of the attention matrix and computational complexity. One example of this approach is the chunking paradigm, which involves dividing the input sequence into fixed blocks and considering the local receptive field of each block [Parmar et al., 2018; Qiu et al., 2020]. Another approach is to limit attention to certain fixed intervals. Models such as the Sparse Transformer [Child et al., 2019] and Longformer [Beltagy et al., 2020] utilize stridden and dilated context windows for attention.

In parallel, another line, the low-rank method, focuses on optimizing the self-attention architecture by approximating the self-attention matrix [Wang et al., 2020]. The primary objective is to reduce the computational complexity of matrix multiplication from N 2 to kN, where N is the number of tokens and k is a constant dependent on the specific model. One notable example is Linformer [Wang et al., 2020], which shrinks the length dimension of the keys and values to a lower-dimensional representation. Besides, by treating selfattention as kernel functions, various low-rank methods can be developed through kernelization approaches [Katharopoulos et al., 2020; Choromanski et al., 2020; Peng et al., 2020], which adopt an efficient kernelized reconstruction of the selfattention matrix, thereby avoiding computing N 2 matrices.

2.2 Graph Transformers

With the development of Transformer architectures, it is found that it is effective to adopt the global receptive field of the Transformers on the graphs [Ying et al., 2021; Kreuzer et al., 2021; Chen et al., 2022; Zhang et al., 2020; Mialon et al., 2021]. Different from the Transformer architectures applied to sequences, the design of graph Transformers emphasizes the use of positional encodings to capture the topological signal of the graphs. Without such encodings, the Transformer can only operate on the fully-connected graphs.

SAN [Kreuzer et al., 2021] adopts Laplacian positional encodings for the nodes and combines two types of attention mechanism, one for virtual fully-connected graphs while another for real graph edges. Graphormer [Ying et al., 2021] incorporates three different positional encodings, namely centrality, spatial and edge encoding respectively, combined with dense attention architecture. SAT [Chen et al., 2022] proposes subtree and subgraph extractors to extract structural features on the graphs, and then uses the similarity scores between features to define positional encodings.

Among them, there are several works focused on improving the scalability of graph Transformers. GraphGPS [Ramp´aˇsek et al., 2022] represents the first attempt to accelerate graph Transformers by combining Fast Transformers with message passing neural networks. GKAT [Choromanski et al., 2022] integrates multiple masking mechanisms into Transformer architectures to improve scalability. Further, NodeFormer [Wu et al., 2022] developes a kernelized gumbel-softmax operator, which reduces the algorithmic complexity to linearity and enables the learning of latent graph structures from large graphs in a differentiable manner. It represents the first Transformer model capable of scaling all-pair message passing to large graphs.

3 Preliminaries

In this section, we recap the preliminaries in self-attention and graph Transformers.

3.1 Vaswani Self-Attention

The self-attention module is the key component of the Transformer architectures. It can be represented as the following formulation:

where Q,K,V ∈ RN×d , N and d denote the length of sequence and embedding dimension, respectively.

Further, after organizing the above formulation, we have:

where 1N ∈ RN is an all-one vector. In the following discussions, we will ignore the constant √d, since we can simply set Q′ = Q/√d to replace Q. It can be noticed that the multiplication complexity between D-1 and AV is O(Nd), while the computation complexity of D,A and AV is O(N 2 ), O(N 2d) and O(N 2d), respectively.

3.2 Graph Transformers

Compared with sequence data, graph data has rich structure features. In GNNs, the edge-based message passing methods are usually used to describe the graph structures. While for graph Transformers, positional encodings are usually used to construct an RPE matrix to describe structures. In general, it is formulated as:

where B ∈ RN×N is RPE matrix. (For preventing ambiguity, in the sequel, N represents the number of nodes in the graph.)

4 Methodology

In this section, we introduce KDLGT acceleration framework and our models in detail. First, we begin by summarizing the effective accelerating approaches for Transformers on sequential data (Section 4.1).

4.1 Overview of Effective Transformer Approach

Most existing methods for accelerating self-attention adopt the kernel decomposition approach[Choromanski et al., 2020]. The core idea is to construct a kernel function ϕ to rearrange the order of matrix multiplication and nonlinear function exp in (3), which can be formulated as:

where Q′ ,K′ ∈ RN×r , r is the new embedding dimension. In this way, the computation of attention score matrix A can be avoided. Instead, we can first compute the multiplication of (K′ )T and V, and then compute the result between Q′ and (K′ )T V. The time complexity of the two steps is both O(Nrd). For (4), similarly, we can reverse the matrix multiplication order by letting (K′ )T 1N first. In specific, we can approximate Vaswani self-attention as followings:

The time complexity is reduced from O(N 2d) to O(Nrd).

Let κ(x, y) = exp(xT y) ≈ ϕ(x)T ϕ(y), where x, y ∈ Rd . It can be noticed that the most important component in the kernel decomposition approach is the design of kernel function ϕ. In this work, we adopt the following well-defined function ϕ [Choromanski et al., 2020]:

where wi ∼ N(0, Id), r ≤ d is a sampling number as well as embedding dimension. We do not discuss here the validity and stability of function ϕ and refer readers to [Choromanski et al., 2020] for further details.

4.2 Kernel Decomposition Linear Graph Transformer

Further, when we turn our perspective to GT, the RPE matrix B becomes the most significant difference between GTs and Transformers. Therefore, we focus on the RPE matrix B and provide a detailed analysis.

Here, we only consider undirected graphs. In this scenario, as the graph structure is symmetric, B should also be a symmetric matrix. Suppose that the rank of B is d′ ≤ N, and its d′ nonzero eigenvalues are λi, i = 1, . . . , d′ , then we can decompose B as follows:

Let Bq = Udiag([λ1, . . . , λd′ , 0, . . . , 0])[:, 0 : d′ ], Bk = U diag([1, . . . , 1, 0, . . . , 0])[:, 0 : d′ ], where Bq, Bk ∈ RL×d′ ,[: , i : j] represents the i to j rows of the second dimension, then we have B = BqBTk .

Inspired by this, we categorize RPE matrices into the following three types for discussion and propose the Kernel Decomposition Linear Graph Transformer (KDLGT) accelerating framework.

Multiplication Decomposition

Therefore, QKT +B = [Q,Bq][K,Bk]T , where [, ] represents concatenation operator. If we view [Q, Bq] and [K, Bk] as Q and K in (6), then it can be accelerated by kernel decomposition approach, and the complexity is O(N(d + d′ )r).

Addition Decomposition

Thus QKT + B = [Q, b, 1N ][K, 1N , b]T . Same as above, the complexity is O(N(d+2)r) using the kernel decomposition approach.

Ineffective Decomposition

Suppose the rank of B is d′ = O(N), in this case, there is no direct effective acceleration method, such as the shortest path distance (SPD). However, we can still design approximate schemes to replace this type of RPE matrix (See subsection 4.3).

4.3 Our Models

As shown in Figure 1, we propose two types of models suitable for different decomposition scenarios under the KDLGT framework.

Linear Structure-Aware Transformer

There are lots of designs of RPE that match multiplication decomposition. We take Structure-Aware Transformer (SAT) [Chen et al., 2022] as an example, which summarizes a series of multiplication-decomposition RPEs. Specifically, the RPE matrix can be represented as:

where Wq,Wk ∈ Rd×d are parameter matrices and φ(v, G) is a structure extractor that extracts vector representations of some subgraph centered at v in the graph G. The structure extractor here includes k-subtree GNN extractor and ksubgraph GNN extractor, which describe the structural information of different granularities, respectively.

It can be easily noticed that:

where S ∈ RN×d , Si = φ(vi,G). Then, let BSAT q = SWq, BSAT k = SWk/√d, we have:

It is clear that the SAT model can be accelerated under the multiplication decomposition framework. Additionally, other similar RPE designs can also be implemented within this framework.

Shortest Anchor Path Distance Graph Transformer

Shortest Path Distance (SPD) is a commonly used type of RPE in GTs. However, unfortunately, in most cases, the rank of the SPD RPE matrix is quite high, which makes it difficult and meaningless to decompose (ineffective decomposition).

To address this issue, we propose a new type of distance, Shortest Anchor Path Distance (SAPD) to approximate SPD distance, which can be easily decomposed.

The general idea is that we only focus on the distances between anchor nodes and nodes rather than recording the distances between all pairs of nodes on the graph, thereby accelerating the computation. Anchor nodes, as a kind of coordinates, can re-characterize the distance relationship between nodes on the graph. Take arbitrary node a ∈ V as anchor node, we define SAPD induced by a as followings:

where vi, vj ∈ V , 1 ≤ i, j ≤ N, d(·, ·) denotes SPD.

In order to ensure rationality, we adopt linear transformation f instead of the embedding method similar to SPD for encoding RPE:

In experiments, we set f(x) = cx, where c ∈ R is a learnable parameter. Considering that the farther the distance is, the lower the weight should be, we set c < 0.

Therefore, BSAPD(a) can do addition decomposition and thus accelerate. In practice, we select K ≪ N nodes on the graph as anchor nodes and pool the embedding about different anchor nodes to get node features. The overall time complexity is O(NK(d + 2)r), which is still linear.

In order to ensure that the anchor nodes are evenly distributed on the graph, we adopt a greedy algorithm[Pavan and Pelillo, 2006] with time complexity of O(N) to solve the kdominant set on the graph as the anchor nodes set, where k is a receptive field hyperparameter. In this way, we can ensure that there is at least one anchor node in the k-hop neighborhood of each node and the absolute error between SPD and SAPD is within k, that is:

5 Theoretical Analysis

In order to ensure the effectiveness of KDLGT, we analyze the difference gap in the attention matrix distribution of Vaswani self-attention and the kernel decomposition method, draw the following conclusion and make a detailed proof.

Lemma 1. Let

Then we have:

Proof. For simplicity, here we note κ′ (Qi,Kj) = ϕ(Qi)T ϕ(Kj). According to Lemma 2 in [Choromanski et al., 2020], we have:

where MSEκ,κ′ (Qi, Kj ) = E[(κ(Qi, Kj ) - κ′ (Qi, Kj ))2 ]. Since for any random variable X, E(X2 ) = E2 (X)+Var(X), E[X] ≤ √E[X2 ] holds, we have:

Expanding and rearranging (19), we can derive:

Let cmax i =max{cij}Nj=1 . Here, we assume that cij ∈ [0, 1). Moreover, we obtain:

Finally, after further scaling, we have:

It can be easily verified that P ′ (Qi,Kj) = P(Qi,Kj) for ∀i, j ∈ {1, 2, . . . , N} if cmax = 0, which means that the distance of Vaswani and approximated self-attention matrix distributions is upper-bounded by the constant cmax . Besides, this upper bound increases monotonically with cmax . It is worth noting that we have assumed that the range of cij is

[0, 1) in the proof. This assumption can be validated from an experimental point of view. In practice, after normalization layers which are stacked following self-attention blocks, the mean and variance of Qi and Kj are 0 and 1, respectively. Therefore, E[cij] = 0, and this also implies that cmax has a large probability distribution around 0, which we also verify in Section 6.2. In conclusion, the difference between the two matrix distributions can be effectively controlled theoretically. Therefore, our KDLGT framework is a well-defined linear approximation of quadratic GTs.

6 Experiments

In this section, we present an evaluation of the precision and efficiency of our proposed Linear Structure-Aware Transformer (LSAT) and Shortest Anchor Path Distance Graph Transformer (SAPDGT) in comparison to state-of-the-art models on several graph benchmark datasets.

6.1 Experimental Setup

Datasets

We investigate the performance of LSAT and SAPDGT on both real-world datasets and synthetic datasets. The dataset statistical details are presented in Table 1.

For real-world datasets, 6 graph-level datasets and 4 node-level datasets are adopted. The benchmarking-GNN [Dwivedi et al., 2020] (ZINC), OGB [Hu et al., 2020] (OGBG-MOLHIV) and TUD [Morris et al., 2020] (MUTAG, COX2 MD, PROTEINS, NCI1) are all popular graph-level benchmark datasets, where each graph represents a molecule, and nodes represent atoms in the molecules. The Cora, Citeseer and PubMed [Yang et al., 2016] are popular citation datasets, whose nodes represent academic papers and node features are the word bag of papers. LastFM-Asia [Rozemberczki and Sarkar, 2020] is a social network that was collected from the public API in March 2020. Their nodes are LastFM users from Asian countries and edges are mutual follower relationships between them. The node features are extracted based on the artists liked by the users. In the experiments, for the datasets without public splits, we use random split with the ratio of training/validation/test sets being 7/1.5/1.5.

For synthetic datasets, we generate a series of graphs for efficiency experiments. The size of the synthetic graphs increases from 211 to 214 in proportion to √2. (When the size is a non-integer, it will be rounded down.) Besides, we adopt 6-regular graphs here to limit the density of the graph,

Baselines

In the experiments, in addition to comparing with SAT and Graphormer which are quadratic-complexity graph Transformers and used as precision performance upper bound, we also select the following effective graph Transformers as strong baselines.

• GPS [Ramp´aˇsek et al., 2022] GPS is an effective GT architecture designed through GNN + Performer paradigm. In experiments, we use its Laplacian eigenvectors encodings (LapPE) and random-walk structural encoding (RWSE) as position encodings. In order to

avoid the influence of the GNN encoder in the GPS model, we set None (no GNN encoder) as the comparison experiment.

• GKAT [Choromanski et al., 2022] GKAT is a GT model which applies low-rank masked attention via Random Walks Graph-Nodes Kernel (RWGNK).
• DGT [Park et al., 2022] DGT reduces self-attention quadratic time complexity by performing sparse attention with dynamically sampled key and value pairs.

Settings

Our models are implemented in PyTorch. We use Adam [Kingma and Ba, 2015] as the optimizer and set hyperparameter ϵ to 1e-7 and (β1, β2) to (0.99, 0.999), respectively. Besides, the initial learning rate is set to 1e-4 with a linear decay learning rate scheduler. The training and inference batch sizes are both set to 128. All models are trained and evaluated on 3 NVIDIA RTX 3090 GPUs for the fairness of efficiency comparison.

6.2 Experimental Results

Graph and Node Representation Experiments

Table 2 summarizes the performance of LSAT and SAPDGT on graph-level and node-level datasets. First of all, in general it can be observed that the precision performance degradation of LSAT and SAPDGT is not significant compared to SAT and Graphormer, respectively. Additionally, it is noteworthy that SAPDGT achieves better results than its upper bound Graphormer on the MUTAG, PROTEINS, and NCI1 datasets, and LSAT also outperforms its upper bound SAT on the Cora, Citeseer, and LastFM-Asia datasets. This highlights the exceptional generalization capability of our proposed KDLGT.

Specifically, on graph-level tasks, LSAT and SAPDGT achieve state-of-the-art results on the OGBG-MOLHIV, COX2 MD, PROTEINS, and NCI1 datasets. Besides, there is a significant performance margin between LSAT, SAPDGT and baselines on the NCI1 dataset. Furthermore, by comparing the results of GPS with or without GNN encoders, it can be observed that the contribution of GNN encoders to the performance of GPS is substantial. If the blessing of GNN encoders is lost, the performance of GPS will drop sharply. We think that this method of using GNN encoders to enhance performance is not specific to GPS and is applicable to various GT models, including our proposed models. Therefore, it is fairer and more reasonable to compare the performance of our model with the GPS model without GNN encoders. Under these conditions, it can be observed that LSAT and SAPDGT achieve state-of-the-art results on all graph-level datasets.

In regards to node-level tasks, unfortunately, it can be noticed that SAT and Graphormer do not perform well on node classification datasets with high edge homogeneity such as the Cora, Citeseer, and PubMed datasets. As a result, LSAT and SAPDGT do not exhibit significant advantages over the baselines on these datasets. Conversely, on the node classification datasets with low edge homogeneity such as LastFMAsia, LSAT outperforms the baselines by a significant margin. This illustrates that LSAT is more suitable for learning on large-scale, low-homogeneity graph data.

Efficiency Experiments

Here, we take the multiplication decomposition type, specifically LSAT, as an example, and use the accelerating ratio as a metric to analyze the performance of our proposed KDLGT framework under different graph sizes and model settings of embedding dimension d. The number of attention heads is fixed at 4 and the number of label classes of synthetic data is set to 5. The results of evaluating and training are recorded and presented in Figure 2.

First of all, it can be noticed that LSAT exhibits significant acceleration efficiency compared to SAT on large-scale graph data, achieving a remarkable result of 9.63x under certain data scales. Since LSAT is linear complexity while SAT is quadratic, it is expected that the acceleration ratio should increase linearly with respect to graph size N in theory, which is also supported by the results presented in Figure 2. It should be noted that here the horizontal axis is log2(N), thus the curve actually grows linearly.

Moreover, the embedding dimension d also plays an important role. Generally, as the ratio of N/d increases, the acceleration ratio becomes more pronounced. In extreme cases, such as when N ≤ 211.5 , the acceleration ratio may be less than 1 due to the relatively small size of the graph, where d becomes the dominant factor affecting efficiency. In this scenario, the time saved by the KDLGT framework is less than the time required for other additional operations (such as sampling), resulting in an overall increase in time. This indicates that the KDLGT framework is more suitable for accelerating large-scale data, which aligns with the goal of this research.

Lastly, it is not surprising that the KDLGT framework also optimizes the space complexity of the Vaswani self-attention. In the experiments, we notice that when N > 214 , an out-ofmemory (OOM) problem occurs for SAT while LSAT can still operate normally. In conclusion, the KDLGT framework achieves obvious advantages in both time complexity and space complexity when applied to large-scale graph data.

The Distribution of the cmax

We conduct the cmax distribution experiment on the NCI1 dataset as ablation studies. We adopt a three-layer LSAT model and fix the parameters after training for experiments. For visualization, we randomly sample 1, 000 graphs and use the frequency of occurrence to approximate the probability distribution. In particular, we directly obtain the maximum value of the cmax of different attention heads for simplicity. The experimental results are shown in Figure 3.

Overall, it can be observed that the distribution of the cmax is concentrated around 0, which supports our assumption in Section 5. Therefore, it is reasonable to assert that cij ∈ [0, 1) for ∀i, j ∈ {1, 2, . . . , N}.

Taking the bottom layer as an example, we can verify the distribution gap between SAT and LSAT attention matrix exactly. The expectation and maximum values of cmax are around 0.1 and 0.3, corresponding to the upper bound value 1 + cmax /1 - cmax of 1.22 and 1.86, respectively. Considering the extremely infrequent occurrence of maximum value, we think this numerical behavior is acceptable.

7 Conclusion

In this work, we present the Kernel Decomposition Linear Graph Transformer (KDLGT), an accelerating framework for building scalable and powerful GTs. Under KDLGT framework, we select two representative GTs and design our models LSAT and SAPDGT. On one hand, a rigorous theoretical analysis is conducted to ensure performance guarantees. On the other hand, a series of experiments are carried out to evaluate the KDLGT in terms of precision and efficiency. Both the theoretical analysis and experimental results demonstrate that the KDLGT not only significantly improves learning efficiency but also preserves the precision performance of the GTs and achieves state-of-the-art results on various datasets.

Acknowledgements

This work was supported by National Key R&D Program of China (Grant 2022ZD0160302), the major key project of PCL, China (Grant No. PCL2021A12), the National Natural Science Foundation of China (Grant No. 62276006 and No. 62276004) and Qualcomm.