Deformable Graph Transformer

Graph Neural Networks (GNNs). Consider an undirected graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ with a set of $N$ nodes $\mathcal{V}=\{v_{1},v_{2},\dots,v_{N}\}$ and a set of edges $\mathcal{E}=\{(v_{i},v_{j})\ |\ v_{i},v_{j}\in\mathcal{V}\}$ where the nodes $v_{i},v_{j}\in\mathcal{V}$ are connected. Each node $v_{i}\in\mathcal{V}$ has a feature vector $\mathbf{x}_{i}\in\mathbb{R}^{F}$, where $F$ is the dimensionality of the node feature, and a set of neighborhoods $\mathcal{N}{(i)}=\{v_{j}\in\mathcal{V}|(v_{i},v_{j})\in\mathcal{E}\}$.

Given a graph $\mathcal{G}$ and a set of node features $\left\{\mathbf{x}_{i}\right\}_{i=1}^{N}$, Graph Neural Networks (GNNs) aim to learn each node representation by an iterative aggregation of transformed representations of the node itself and its neighborhoods as follows:

$${\textbf{h}}^{(l)}_{i}=\sigma\left(\textbf{W}^{(l)}\left(c_{ii}^{(l)}\textbf{h}_{i}^{(l-1)}+\sum_{v_{j}\in\mathcal{N}{(i)}}{c_{ij}^{(l)}\textbf{h}_{j}^{(l-1)}}\right)\right),$$

(1)

where ${\mathbf{h}}^{(l)}_{i}\in\mathbb{R}^{d^{(l)}}$ is a hidden representation of node $v_{i}$ in the $l$-th GNN layer, ${\mathbf{h}}^{(0)}_{i}=\mathbf{x}_{i}$, $\mathbf{W}^{(l)}\in\mathbb{R}^{d^{(l)}\times d^{(l-1)}}$ is a learnable weight matrix at the $l$-th GNN layer, $\sigma$ is a non-linear activation function, $c_{ij}^{(l)}$ and $c_{ii}^{(l)}$ represent weights for aggregation characterized by each GNN. For example, GCN (Kipf & Welling, 2017) can be represented as a form of (1) if $c_{ij}=(\text{deg}(i)\text{deg}(j))^{-1/2}$ and $c_{ii}=(\text{deg}(i))^{-1}$, where $\text{deg}(i)$ is the degree of node $v_{i}$, and GAT (Veličković et al., 2018) learns $c_{ij}^{(l)}$ and $c_{ii}^{(l)}$ based on the attention mechanism.

Transformer-based Graph Models. Most Transformer-based graph models (Dwivedi & Bresson, 2020; Kreuzer et al., 2021) learn each node representation with self-attention mechanism as Transformers (Vaswani et al., 2017) learn each token with self-attention. Given a node $v_{q}$ with a corresponding node embedding $\mathbf{z}_{q}\in\mathbb{R}^{c}$ and a set of key/value vectors $\mathcal{F}=\left\{\mathbf{f}_{k}\right\}_{k=1}^{N}$, the Multi-Head Attention (MHA) for Transformer-based graph models is formulated as follows:

$$\text{MHA}\left(\mathbf{z}_{q},\mathcal{F}\right)=\sum_{m=1}^{M}\mathbf{W}_{m}\left[\sum_{k=1}^{N}\mathbf{A}_{mqk}\cdot\mathbf{W}^{\prime}_{m}\mathbf{f}_{k}\right],$$

(2)

where $m$ is the index of $M$ attention heads, $\mathbf{W}_{m}\in\mathbb{R}^{c\times c_{v}}$ and $\mathbf{W}_{m}^{\prime}\in\mathbb{R}^{c_{v}\times c}$ are weight matrix parameters. The attention weight $\mathbf{A}_{mqk}$ between the query $\mathbf{z}_{q}\in\mathbb{R}^{c}$ and the key $\mathbf{f}_{k}\in\mathbb{R}^{c}$, which is generally calculated as $\mathbf{A}_{mqk}=\frac{\exp\left[{\mathbf{z}_{q}^{\top}\mathbf{U}^{\top}_{m}\mathbf{V}_{m}\mathbf{f}_{k}}/{\sqrt{c_{v}}}\right]}{Z}$, where $Z\in\mathbb{R}$ is a normalization factor to achieve $\sum_{k=1}^{N}\mathbf{A}_{mqk}=1$ and $\mathbf{U}_{m},\mathbf{V}_{m}$ are weight matrices to compute queries and keys.

By harnessing the power of the multi-head self attention, Transformer-based graph models have shown superior performance on graph-related tasks (Dwivedi & Bresson, 2020; Kreuzer et al., 2021). But, they have difficulty in learning representations on large-scale graphs. Since the large-scale graph has a vast number of nodes, the scalability issue of the multi-head attention is dramatically exacerbated. Also, the enormous number of keys to attend per query node increases the risk of the aggregation of noisy information from irrelevant key nodes. In this paper, we design a sparse Transformers named Deformable Graph Transformer (DGT) that attends to only a small number of relevant keys through deformable sampling (Zhu et al., 2021). By reducing the number of key nodes, it is possible to make transformer-based architecture effective and efficient on large-scale graphs.

In this section, we present Deformable Graph Attention (DGA), a key module in our proposed Deformable Graph Transformer (DGT). The overview of the deformable graph attention is illustrated in Figure 1. The general idea of deformable graph attention is to design the deformable attention mechanism for graph-structured data, and therefore DGT performs the attention mechanism with partially selected relevant key nodes, which makes the transformer-based architecture efficient and effective on large-scale graphs. However, extending the deformable attention mechanism to graph representation learning is non-trivial since the offset-based interpolation in the deformable attention only works in Euclidean space whereas graphs are non-Euclidean data. Also, even when embedding nodes in a graph into a low-dimensional space via graph embedding methods (e.g., Node2vec (Grover & Leskovec, 2016)), nodes are irregularly distributed in the low-dimensional space, which makes difficult to effectively deform an attention map.

To address this challenge, we propose a NodeSort module that converts a graph into a sorted sequence of nodes in a regular space. We define a base node $v_{b}$ as the first node in a sorted sequence which is similar to an ego node in an ego graph. NodeSort differentially sorts nodes depending on the base node. In other words, NodeSort provides a relative ordering that varies across base nodes whereas conventional topological sort yields a single (absolute) ordering for a graph. Specifically, given a base node, NodeSort sorts nodes and returns a sequence of their features as follows:

$$S_{\pi b}=\text{NodeSort}_{\pi}(\mathcal{G},v_{b},\left\{\mathbf{x}_{i}\right\}_{i=1}^{N})=[\mathbf{x}_{\sigma_{\pi b}^{-1}(i)}]_{i=1}^{N},$$

(3)

where $\pi$ denotes a specific criterion for sorting nodes and $\sigma_{\pi b}$ is a bijective mapping from $\mathcal{V}$ to $\mathcal{V}$ depending on a base node $v_{b}$. To consider both structural and semantic proximity, we generate multiple sorted node sequences for each node $v_{b}$, $\{S_{\pi b}\}_{\pi\in\Pi}$, from a set of diverse criteria $\Pi$, such as Personalized PageRank (PPR) score, BFS, and feature similarity. Note that $\sigma_{\pi b}$ redefines neighbors of each node $v_{b}$ based on various aspects $\pi\in\Pi$ beyond 1-hop neighbors in the existing GNNs.

Now, we introduce our Deformable Graph Attention (DGA), which is a sparse attention by dynamically sampling key/value pairs from the set of sorted sequences of node features. To benefit from various properties of the graphs, the deformable graph attention module is designed to deal with diverse node sequences. Different from existing methods based on deformable sampling modules in computer vision, which only considers spatial proximity on a grid, DGA captures both structural and semantic proximity. Given the set of sorted sequences $\{S_{\pi q}\}_{\pi\in\Pi}$ and features of query node $\mathbf{z}_{q}$, DGA is defined as

$$\text{DGA}\left(\mathbf{z}_{q},\{S_{\pi q}\}_{\pi\in\Pi}\right)=\sum_{\pi\in\Pi}\sum_{m=1}^{M}\mathbf{W}_{\pi m}\left[\sum_{k=1}^{K}\mathbf{A}_{\pi mqk}\cdot\mathbf{W}^{\prime}_{\pi m}\tilde{S}_{\pi q}\left(\Delta\mathbf{p}_{\pi mqk}\right)\right],$$

(4)

where $\tilde{S}_{\pi q}\left(\Delta\mathbf{p}_{\pi mqk}\right)$ denotes the representation of the $k$-th key, $K$ denotes the number of keys, $\mathbf{W}_{\pi m}\in\mathbb{R}^{c\times c_{v}}$ and $\mathbf{W}^{\prime}_{\pi m}\in\mathbb{R}^{c_{v}\times c}$ are the learnable weight matrices for each criteria $\pi\in\Pi$ and the $m$-th attention head, $\mathbf{A}_{\pi mqk}=\theta_{\pi mk}^{\text{att}}(\mathbf{z}_{q})$ denotes the attention weight from a linear function $\theta^{\text{att}}$ between the $q$-th query and the $k$-th key of the $m$-th attention head and criterion $\pi$, which is normalized by $\sum_{k=1}^{K}\mathbf{A}_{\pi mqk}=1$ and the attention weight needs to be in a interval $[0,1]$.

$\Delta\mathbf{p}_{\pi mqk}$ is a sampling offset of the $k$-th key for criterion $\pi$ and $m$-th attention head, which is generated by $\theta_{\pi mk}^{\text{off}}(\mathbf{z}_{q})$, where $\theta^{\text{off}}$ is a linear function with an activation function. As the offset $\Delta\mathbf{p}_{\pi mqk}$ is fractional, we compute $\tilde{S}_{\pi q}\left(\mathbf{p}\right)$ by kernel-based interpolation:

$$\tilde{S}_{\pi q}\left(\textbf{p}\right)=\sum_{i}g(\mathbf{p},i)\cdot S_{\pi q}[i],\quad g(a,b)=\begin{cases}\exp\left(-\frac{\left(a-b\right)^{2}}{\gamma}\right),&\mbox{ if }\lvert a-b\rvert<\epsilon\\ 0,&\mbox{ otherwise}\end{cases}$$

(5)

where $\gamma\in\mathbb{R}^{++}$ denotes the bandwidth of the kernel, $\epsilon$ is a hyper-parameter for truncating the kernel, and $S_{\pi q}[i]$ is the node feature at a $i$-th index of the sequence $S_{\pi q}$.

A standard way of computing $\tilde{S}_{\pi q}\left(\mathbf{p}\right)$ is calculating Eq. (5) by substituting $g$ with $g\left(a,b\right)=\max\left(0,1-\left\lvert a-b\right\rvert\right)$. But, the existing way considers only two nodes whose coordinates are ceiling and floor values of a point $\mathbf{p}$. It might not be sufficient for the offsets $\Delta\mathbf{p}$ to move to get relevant information by looking at only two nodes. As the scale of the graph increases, a target node requires much more nodes to learn its representation. So, we apply the Radial Basis Function kernel for computing $\tilde{S}_{\pi q}$ as in Eq. (5) to look at a wide range of nodes for representing query nodes.

Positional Encoding is a crucial component in Transformer to reflect domain-specific positional information into its attention mechanism. A major issue with positional encoding on graphs is the absence of absolute positions of nodes, unlike other domains. One remedy to this issue is to encode relations between nodes. Here, we propose learnable positional encodings, Katz PE, for Transformer-based graph models based on connectivity between nodes. To be specific, inspired by the matrix of Katz indices (Katz, 1953) which counts all paths between nodes with the decaying weight $\beta$ to reflect the preference for shorter paths, i.e., $\hat{\mathbf{A}}=\sum_{k=1}^{\infty}\beta^{k-1}{\mathbf{A}^{k}}$, our method learns positional embeddings, $\text{PE}_{i}$, of each node $i$ by the nonlinear transform of $\hat{\mathbf{A}}$ as follows:

$$\text{Katz PE}(v_{i})=\text{MLP}(\hat{\mathbf{A}}[v_{i}]^{T}),$$

(6)

where MLP is a Multi-Layer Perceptron, and $\hat{\mathbf{A}}[v_{i}]$ is the row vector of node $v_{i}$ in $\hat{\mathbf{A}}$. We limit the maximum $k$ in $\hat{\mathbf{A}}$ to $K$, i.e., $\hat{\mathbf{A}}=\sum_{k=1}^{K}\beta^{k-1}{\mathbf{A}^{k}}$ for the efficient calculation. In addition, when $N$ is large, then we sample $N^{\prime}$ anchor nodes with a high degree and utilize the submatrix of Katz indices $\hat{\mathbf{A}}^{\prime}\in\mathbb{R}^{N\times N^{\prime}}$. Our learnable Katz PE is simple yet more effective for both our DGT and vanilla Transformer than existing pre-computed positional encodings. See Section 4.3, for more details.

Finally, we introduce our Deformable Graph Transformer (DGT) built upon our proposed Graph Deformable Attention and Positional Encoding. Deformable Graph Transformer first encodes node feature $\mathbf{x}_{i}$ with the learnable function $f_{\theta}$, which can be MLP, and combines with positional embeddings from Eq. (6) as

$$\mathbf{z}_{i}^{(0)}=f_{\theta}(\mathbf{x}_{i})+\text{Katz PE}(v_{i}).$$

(7)

Then, given a set of sorted sequences $\{S_{q}^{\pi}\}_{\pi\in\Pi}$ from Eq. (4), each $l$-th Deformable Graph Attention layer in DGT performs the attention mechanism with a sampled set of informative keys and applies skip-connection and MLP to update node representations as follows:

$$\mathbf{\hat{z}}_{i}^{(l)}=\mbox{DGA}\left(\mathbf{z}_{i}^{(l)},\{S_{\pi i}\}_{\pi\in\Pi},\mathbf{Z}^{(l-1)}\right)+\mathbf{z}_{i}^{(l-1)},\quad\mathbf{z}_{i}^{(l)}=\mbox{MLP}(\mathbf{\hat{z}}_{i}^{(l)})+\mathbf{\hat{z}}_{i}^{(l)}.$$

(8)

After the stack of $L$ Deformable Graph Attention blocks, each node representation $\mathbf{z}_{i}^{(L)}$ is used for node classification on top and MLP followed by a softmax layer is used as $\hat{y}_{i}=\text{Softmax}(\text{MLP}(\mathbf{z}^{(L)}_{i}))$. Our loss function is a cross-entropy on nodes that have ground truth labels.

We provide the comparison of computational complexity between the self-attention in most Transformer-based graph models and the Deformable Graph Attention (DGA) in our DGT. As the number of nodes $N$ increases, our DGA with a linear complexity of $N$ is more efficient than the self-attention with a quadratic complexity of $N$.

Self-Attention in most Transformer-based graph models (Vaswani et al., 2017; Ying et al., 2021; Dwivedi & Bresson, 2020). Suppose that $N$ is the number of nodes, $C$ is the dimensionality of hidden representations. The self-attention operation requires a huge computation cost with the complexity of $\mathcal{O}\left(N^{2}C+NC^{2}\right)$, where $C$ is the dimensionality of hidden representations.

Deformable Graph Attention. Consider $K$ is the number of keys, $T$ is the number of critera, and $W$ is the number of nonzero values of $g$ in Eq. (5). Then, the computational complexity of Deformable Graph Attention is $\mathcal{O}(N\cdot(C^{2}T+WKCT))$ ,which is a linear complexity with respect to the number of nodes $N\gg W,K,C,T$. The details are in Sec A.

Datasets. We validate the effectiveness of our model on node classification using four heterophilic graph datasets and four homophilic graph datasets, which are distinguished by the edge-based homophily ratio (Zhu et al., 2020) defined as $h=\frac{\left\lvert\left\{(v_{i},v_{j}):(v_{i},v_{j})\in\mathcal{E}\land y_{i}=y_{j}\right\}\right\rvert}{\left\lvert\mathcal{E}\right\rvert}$. Each dataset has an edge-based homophily ratio ranged from $h=0.22$ (very heterophilic) to $h=0.81$ (very homophilic). For large-scale graphs, we evaluate our method on twitch-gamers, obgn-arxiv, and Reddit datasets. More details about the datasets are in Sec C.1.

Baselines. To demonstrate the effectiveness of our Deformable Graph Transformer (DGT), we compare DGT with following baselines: (1) six standard GNNs: GCN (Kipf & Welling, 2017), GAT (Veličković et al., 2018), GraphSAGE (Hamilton et al., 2017), JKNet (Xu et al., 2018), SGC (Wu et al., 2019), GATv2 (Brody et al., 2022); (2) four GNNs designed for heterophilic settings: MixHop (Abu-El-Haija et al., 2019), Geom-GCN (Pei et al., 2020), H2GCN (Zhu et al., 2020), DeformableGCN (Park et al., 2022); (3) four Transformer-based architectures for graphs: Transformer (Vaswani et al., 2017), Graphormer (Ying et al., 2021), GT-full (Dwivedi & Bresson, 2020), GT-sparse (Dwivedi & Bresson, 2020). (4) two variants of our proposed DGT: DGT-light using a single sorting criterion ($|\Pi|=1$) and DGT using multiple sorting criteria ($|\Pi|=3$).

Implementation details. The best model on the validation split is used for reporting the performance. We adopt the splits (48%/ 32%/ 20%) of nodes per class for (train/ validation/ test) following (Pei et al., 2020; Zhu et al., 2020) and the experiments are repated 30 times on Actor, Squirrel, Chameleon, Cora, and Citeseer datasets. For twitch-gamers, ogbn-arxiv, and Reddit, the experiments are conducted with the splits provided by (Lim et al., 2021), (Hu et al., 2020), and (Hamilton et al., 2017), respectively and repeated 10 times. More implementation details are in Sec C.2.

Table 1 summarizes node classification results of our proposed DGT and DGT-light, GNN-based models, and Transformer-based graph models. DGT consistently shows superior performance across all eight datasets, which achieves the state-of-the-art performance on seven datasets and competitive performance on one dataset. Especially, our DGT significantly outperforms transformer-based baselines by a large margin up to 106%. Surprisingly, existing Transformer-based graph models show poor performance compared to GNN baselines except for Actor. This implies that Transformer-based graph models failed to filter out irrelevant messages and focus on useful nodes. In addition, they are not applicable to large-scale graphs such as ogbn-arxiv, twitch-gamers, and Reddit due to their huge computational costs.

On the other hand, our DGT consistently outperforms both GNNs and Transformer-based graph models on almost all datasets and efficiently handles large-scale graphs by utilizing a small set of relevant nodes. GNNs generally perform well in homophilic graphs such as Cora and Citeseer, but show relatively inferior performance in heterophilic graphs such as Actor, Squirrel, Chameleon, and twitch-gamers. This is because most GNNs utilize directly connected nodes for aggregation even in heterophilic graphs. Instead, our DGT and DGT-light show larger performance gain in heterophilic graph datasets since it captures long-range dependency, which is important for learning on heterophilic graphs.

We also compare our DGT and DGT-light with other Transformer-based architectures to validate the efficiency (FLOPs) of our DGT in Table 2. The table shows that our DGT and DGT-light improves not only the performance but also the efficiency with deformable graph attention on various datasets. As the number of nodes increases, the gap between Transformer and variants of DGT with respect to FLOPS becomes bigger. In particular, on the twitch-gamers, DGT-light (8.05G) reduces computational costs by $\times$ 449 compared to Transformer (3622G).

Here, we provide quantitative results of additional experiments to demonstrate the contribution of each component of our Deformable Graph Transformer. We first provide the ablation study of the NodeSort module and examine the effectiveness of our positional encoding comparing with popular positional encoding methods.

NodeSort module. We conduct an ablation study of the NodeSort module to study the effect of the relative ordering that varies depending on base nodes and various sorting criteria. We compare several constructions: absolute ordering with random permutation (AR), absolute ordering with BFS (AB), absolute ordering with multiple criteria (AM), relative ordering with BFS (RB), and relative ordering with multiple criteria (RM). As shown in Table 3, the absolute ordering approaches (AR, AB, AM) show poor performance, even worse than MLP in Citeseer. The absolute ordering with BFS (AB) shows no performance gain compared to a randomly permuted absolute ordering (AR) on three datasets. This shows that a single absolute node sequence is not sufficient to learn complex relationships between nodes in the graph. On the other hand, the relative ordering with BFS (RB) shows a significant performance improvement up to 5.07(%). Further, the relative ordering with multiple criteria (RM) consistently shows superior performance compared to the relative ordering with BFS (RB).

Positional encoding. To validate the effectiveness of our positional encoding, we compare our proposed PE methods with models without positional encodings (w/o PE) and various encoding methods such as Node2Vec (Grover & Leskovec, 2016) and Laplacian Eigvecs used in GT (Dwivedi & Bresson, 2020) on four datasets. We use Transformer and DGT for the base models. Table 4 demonstrates that our positional encoding is effective on both base models. In particular, it improves the performance by 17.19(%) compared to DGT without positional encoding on Squirrel.

We provide qualitative analysis to understand why DGT is effective. We visualize the top 20 key nodes with high attention scores for a given query node $v_{q}$ in Graphormer and DGT. In DGT, we compute attention scores of each node, $v_{i}$ by $w_{i}=\sum_{k}\mathbf{A}_{\pi mqk}\cdot g(\Delta\mathbf{p}_{\pi mqk},i)$. As shown in Figure 2, a within 1-hop neighborhood of the given query node $v_{q}$, 6 out of 7 neighbors have different labels. So, both DGT and Graphormer aggregate messages (dashed lines) from remote nodes beyond 1-hop neighbors through the attention mechanism. 17 of the top 20 nodes in DGT are nodes with the same label whereas only 4 out of 20 nodes have the same label in Graphormer. Also, the ratio of the attention scores for the nodes with the same label, i.e., ${\sum_{\left\{v_{i}:v_{i}\in\mathcal{V}\land y_{q}=y_{i}\right\}}w_{i}}/{\sum_{v_{i}\in\mathcal{V}}w_{i}}$ is 0.97 for DGT and 0.21 for Graphormer. This evidences that our DGT effectively performs the sparse attention by focusing on a small set of relevant key nodes compared to Graphormer.