Hierarchical Graph Transformer with Adaptive
Node Sampling

Recently, Transformer [33] has shown its superiority in an increasing number of domains [7, 8, 40], e.g. Bert [7] in NLP and ViT [8] in CV. Existing works attempting to generalize Transformer to graph data mainly focus on two problems: (1) How to design dedicated positional encoding for the nodes; (2) How to alleviate the quadratic computational complexity of the vanilla Transformer and scale the Graph Transformer to large graphs. As for the positional encoding, GT [9] firstly uses Laplacian eigenvectors to enhance node features. Graph-Bert [41] studies employing Weisfeiler-Lehman code to encode structural information. Graphormer [38] utilizes centrality encoding to enhance node features while incorporating edge information with spatial (SPD-indexed attention bias) and edge encoding. SAN [21] further replaces the static Laplacian eigenvectors with learnable positional encodings and designs an attention mechanism that distinguishes local connectivity. For the scalability issue, one immediate idea is to restrict the number of attending nodes. For example, GAT [34] and GT-Sparse [9] only consider the 1-hop neighboring nodes; Gophormer [46] uses GraphSAGE [11] sampling to uniformly sample ego-graphs with pre-defined maximum depth; Graph-Bert [41] restricts the receptive field of each node to the nodes with top-k intimacy scores (e.g., Katz and PPR). However, these fixed node sampling strategies sacrifice the advantage of the Transformer architecture. SAC [22] tries to use an LSTM edge predictor to predict edges for self-attention operations. However, the fact that LSTM can hardly be parallelized reduces the computational efficiency of the Transformer.

In parallel, many efforts have been devoted to reducing the computational complexity of the Transformer in the field of NLP [23] and CV [32]. In the domain of NLP, Longformer [2] applies block-wise or strode patterns while only fixing on fixed neighbors. Reformer [19] replaces dot-product attention by using approximate attention computation based on locality-sensitive hashing. Routing Transformer [30] employs online k-means clustering on the tokens. Linformer [35] demonstrates that the self-attention mechanism can be approximated by a low-rank matrix and reduces the complexity from $\mathcal{O}(n^{2})$ to $\mathcal{O}(n)$. As for vision transformers, Swin Transformer [24] proposes the shifted windowing scheme which brings greater efficiency by limiting self-attention computation to non-overlapping local windows while also allowing for cross-window connection. Focal Transformer [37] presents a new mechanism incorporating both fine-grained local and coarse-grained global attention to capture short- and long-range visual dependencies efficiently. However, these sparse transformers do not take the unique graph properties into consideration.

Graph neural networks (GNNs) [18, 11, 12, 44, 43, 31, 45, 42] follow a message-passing schema that iteratively updates the representation of a node by aggregating representations from neighboring nodes. When generalizing to large graphs, Graph Neural Networks face a similar scalability issue. This is mainly due to the uncontrollable neighborhood expansion in the aggregation stage of GNN. Several node sampling algorithms have been proposed to limit the neighborhood expansion, which mainly falls into node-wise sampling methods and layer-wise sampling methods. In node-wise sampling, each node samples $k$ neighbors from its sampling distribution, then the total number of nodes in the $l$-th layer becomes $\mathcal{O}(k^{l})$. GraphSage [11] is one of the most well-known node-wise sampling methods with the uniform sampling distribution. GCN-BS [25] introduces a variance reduced sampler based on multi-armed bandits. To alleviate the exponential neighbor expansion $\mathcal{O}(k^{l})$ of the node-wise samplers, layer-wise samplers define the sampling distribution as a probability of sampling nodes given a set of nodes in the upper layer [4, 16, 49]. From another perspective, these sampling methods can also be categorized into fixed sampling strategies [11, 4, 49] and adaptive strategies [25, 39]. However, none of the above sampling methods in GNNs can be directly applied in Graph Transformer as Graph Transformer does not follow the message passing schema.

Let $G=(A,X)$ denote the unweighted graph where $A\in\mathbb{R}^{n\times n}$ represents the symmetric adjacency matrix with $n$ nodes, and $X\in\mathbb{R}^{n\times p}$ is the attribute matrix of $p$ attributes per node. The element $A_{ij}$ in the adjacency matrix equals to 1 if there exists an edge between node $v_{i}$ and node $v_{j}$, otherwise $A_{ij}=0$. The label of node $v_{i}$ is $y_{i}$. In the node classification problem, the classifier has the knowledge of the labels of a subset of nodes $V_{L}$. The goal of semi-supervised node classification is to infer the labels of nodes in $V\backslash V_{L}$ by learning a classification function.

The Transformer architecture consists of a series of Transformer layers [33]. Each Transformer layer has two parts: a multi-head self-attention (MHA) module and a position-wise feed-forward network (FFN). Let $\mathbf{H}=\left[\bm{h}_{1},\cdots,\bm{h}_{m}\right]^{\top}\in\mathbb{R}^{m\times d}$ denote the input to the self-attention module where $d$ is the hidden dimension, $\bm{h}_{i}\in\mathbb{R}^{d\times 1}$ is the hidden representation at position $i$, and $m$ is the number of positions. The MHA module firstly projects the input $\mathbf{H}$ to query-, key-, value-spaces, denoted as $\mathbf{Q},\mathbf{K},\mathbf{V}$, using three matrices $\mathbf{W}_{Q}\in\mathbb{R}^{d\times d_{K}},\mathbf{W}_{K}\in\mathbb{R}^{d\times d_{K}}$ and $\mathbf{W}_{V}\in\mathbb{R}^{d\times d_{V}}$:

Then, in each head $i\in\{1,2,\dots,B\}$ ($B$ is the total number of heads), the scaled dot-product attention mechanism is applied to the corresponding $\{\mathbf{Q}_{i},\mathbf{K}_{i},\mathbf{V}_{i}\}$ :

Finally, the outputs from different heads are further concatenated and transformed to obtain the final output of MHA:

where $\mathbf{W}_{O}\in\mathbb{R}^{d\times d}$. In this work, we employ $d_{K}=d_{V}=d/B$.

The goal of Graph Coarsening [29, 26, 17] is to reduce the number of nodes in a graph by clustering them into super-nodes while preserving the global information of the graph as much as possible. Given a graph $G=(V,E)$ ($V$ is the node set and $E$ is the edge set), the coarse graph is a smaller weighted graph $G^{\prime}=(V^{\prime},E^{\prime})$. $G^{\prime}$ is obtained from the original graph by first computing a partition $\{C_{1},C_{2},\cdots,C_{|V^{\prime}|}\}$ of $V$, i.e., the clusters $C_{1}\cdots C_{|V^{\prime}|}$ are disjoint and cover all the nodes in $V$. Each cluster $C_{i}$ corresponds to a super-node in $G^{\prime}$ The partition can also be characterized by a matrix $\hat{P}\in\{0,1\}^{|V|\times|V^{\prime}|}$, with $\hat{P}_{ij}=1$ if and only if node $v_{i}$ in $G$ belongs to cluster $C_{j}$. Its normalized version can be defined by $P\triangleq\hat{P}D^{-\frac{1}{2}}$, where $D$ is a $|V^{\prime}|\times|V^{\prime}|$ diagonal matrix with $|C_{i}|$ as its $i$-th diagonal entry. The feature matrix and weighted adjacency matrix of $G^{\prime}$ are defined by $X^{\prime}\triangleq P^{T}X$ and $A^{\prime}\triangleq P^{T}AP$. After Graph Coarsening, the number of nodes/edges in $G^{\prime}$ is significantly smaller than that of $G$. The coarsening rate can be defined as $c=\frac{|V^{\prime}|}{|V|}$.

Our Adaptive Node Sampling module aims to adaptively choose the batch of most informative nodes by a multi-armed bandit mechanism. In our setting, it is intuitive that the contributions of nodes to the learning performance can be time-sensitive. Meanwhile, the rewards which are associated with the model training process are not independent random variables across the iterations. The above situation can satisfy the adversarial setting in the multi-armed bandit problem [1]. To adaptively choose the most informative nodes with the designed sampling strategies, we adjust the method ALBL proposed in [14], which modifies the EXP4.P method [3]. EXP4.P possesses a strong theoretical guarantee for the adversarial setting.

Formally, let $w^{t}=(w_{1}^{t},\cdots,w_{K}^{t})$ be the adaptive weight vector in iteration $t$, where the $k$-th non-negative element $w_{k}^{t}$ is the weight corresponding to the $k$-th node sampling strategy. The weight vector $w^{t}$ is then scaled to a probability vector $p^{t}=(p_{1}^{t},\cdots,p_{K}^{t})$ where $p_{k}^{t}\in[p_{min},1]$ with $p_{min}>0$. Our ANS-GT adaptively sample nodes based on the probability vector and then obtains the reward of the action.

For each center node, we consider the sampling probability matrix $Q^{t}\in\mathbb{R}^{K\times n}$, where $K$ is the number of sampling heuristics and $n$ is the number of nodes in the graph. Specifically, $Q_{k,j}^{t}$ denotes the $k$-th sampling strategy’s preference on selecting node $j$ in iteration $t$ and $Q^{t}$ is normalized to satisfy $\sum_{j=1}^{n}Q_{k,j}^{t}=1$. Note that our ANS-GT is a general framework and is not restricted to a certain set of sampling heuristics. In our work, we adopt four representative sampling heuristics:

1-/2-hop neighbors: We adopt the normalized adjacency matrix $\widetilde{A}=\hat{D}^{-\frac{1}{2}}\hat{A}\hat{D}^{-\frac{1}{2}}$ for 1-hop neighbors and $\widetilde{A}^{2}$ for 2-hop neighbors, where $\hat{A}=A+I$ is the adjacency matrix of the graph $G$ with self connections added and $\hat{D}$ is a diagonal matrix with $\hat{D}_{ii}=\sum_{j}\hat{A}_{ij}$.

KNN: We adopts the cosine similarity of node attributes to measure the similarities of nodes. Mathematically, the similarity score $S_{ij}$ between node $i$ and $j$ is calculated as $S_{ij}=x_{i}\cdot x_{j}/(|x_{i}|\cdot|x_{j}|)$ where $x_{i}$ is the feature vector of node $v_{i}$.

PPR: The Personalized PageRank [27] matrix $S$ is calculated as: $S=c(I-(1-c)\overline{A})$, where factor $c\in[0,1]$ (set to 0.15 in our experiments). $\overline{A}$ denotes the column-normalized adjacency matrix.

Finally, given the probability vector $p^{t}$ and the node sampling matrices $Q^{t}$, the final node sampling probability is:

We introduce a novel reward scheme based on the attention weights which is intrinsic in Transformer. Formally, given an attention matrix $\mathcal{A}=\operatorname{Softmax}(QK^{T})/\sqrt{d}$, we use the first row of the attention matrix, i.e., $\mathcal{A}_{1,i}$ multiplying the magnitude of corresponding value $V_{i}$ to represent the significance of each node to the center node: $s_{i}=\mathcal{A}_{1,i}\times\|\mathcal{V}_{i}\|$. In the situation of multiple heads and layers in the Transformer, we average the significance scores in multiple attention matrices. The reward to the $k$-th sampling strategy is: $r_{k}=\sum_{i=1}^{N}\frac{s_{i}Q^{t}_{ki}}{\psi^{t}_{i}}$, where $N$ is the number of sampled nodes for each center node. $r_{k}$ can be interpreted as the dot product between the significance score vector and the normalized sampling probability vector. The intuition behind the reward design is that the reward to a certain sampling heuristic is higher if the sampling probability distribution and the node significance score distribution are closer. Thus, exploiting such a sampling heuristic can help graph transformer sample more informative nodes. Finally, we update $w^{t}$ with the reward. In experiments, for the efficiency and stability of training, we update the sampling weights and resample nodes every $T$ epochs. The pseudo-code of ANS-GT is listed in Algorithm 1.

We argue that most node sampling strategies (e.g., 1-hop neighbors) focus on local information and neglect long-range dependencies or global contexts. Therefore, to efficiently capture both the local and global information in the graph, we propose a novel hierarchical graph attention scheme including fine-grained attention and coarse-grained attention. Specifically, we use the proposed adaptive node sampling for local fined-grained attention. On the other hand, we adopt the graph coarsening algorithm [26] to generate the coarsened graph $G^{\prime}$. The sampled $n_{s}$ super-nodes from $G^{\prime}$ are used to capture long-range dependencies. Similar to [46], we also use $n_{g}$ global nodes with learnable features to store global context information. Finally, the hierarchical nodes are concatenated with the center nodes as the input sequences.

For the positional encoding, we use the proximity encoding in [46]: $\Phi_{m}(v_{i},v_{j})=\widetilde{A}^{m}[i,j],m\in\{0,\cdots,M-1\}$, where $\widetilde{A}$ denotes the normalized adjacency matrices with self-loop. Note that our framework is agnostic to the positional encoding scheme and we left other positional encoding methods such as Laplacian eigenvectors [9] for future exploration. We follow the Graphormer framework to obtain the output of the $l$-th transformer layer, $\mathbf{H}^{(l)}$:

We apply the layer normalization (LN) before the multi-head self-attention (MHA) and the feed-forward network (FFN).

In the training and inference, we sample $\mathcal{S}$ input sequences for each center node and use the center node representation from the final Transformer layer $z_{c}^{(s)}$ for prediction. Note that the computational complexity is controllable by choosing suitable number of sampled nodes. A MLP (Multi-Layer Perceptron) is used to predict the node class:

where $\widetilde{\bm{y}}_{c}\in\mathbb{R}^{C\times 1}$ stands for the classification result, $C$ stands for the number of classes. In the training process, we optimize the average cross entropy loss of labeled training nodes $V_{L}$:

where $\bm{y}_{i}\in\mathbb{R}^{C\times 1}$ is the ground truth label of center node $v_{i}$. In the inference stage, we take a bagging aggregation to improve accuracy and reduce variance:

Compared with existing Graph Transformers, ANS-GT requires extra computational costs on graph coarsening and sampling weights update. Here we want to show that the overall computational complexity of ANS-GT is linear with the number of nodes $n$. Hence, ANS-GT is scalable to large graphs. First, the computational complexity of graph coarsening is linear with $n$ [26] and we only need to do it once before training. Second, the computational cost of self-attention calculation in one epoch is $\mathcal{O}(n\mathcal{S}(N+n_{s}+n_{g})^{2})$ where the number of sampled nodes $N$, the number of sampled super-nodes $n_{s}$, the number of global nodes $n_{g}$, and the number of augmentations $\mathcal{S}$ are specified constants. Finally, the cost of updating sampling weights is linear to $n$, which is mainly attributed to the computation of rewards. Empirical efficiency analyses of ANS-GT are shown in the Appendix C.

Datasets. To comprehensively evaluate the effectiveness of ANS-GT, we conduct experiments on the six benchmark datasets including citation graphs Cora, CiteSeer, and PubMed [18]; Wikipedia graphs Chameleon, Squirrel; the Actor co-occurrence graph [5]; and WebKB datasets [28] including Cornell, Texas, and Wisconsin. We set the train-validation-test split as 60%/20%/20%. The statistics of datasets are shown in the Appendix C.

Baselines. To evaluate the effectiveness of ANS-GT on graph representation learning, we compare it with 12 baseline methods, including 8 popular GNN methods, i.e. GCN [18], GAT [34], GraphSAGE [11], JKNet [36], APPNP [20], Geom-GCN [28], H₂GCN [48], and GPRGNN [6] along with four state-of-the-art Graph Transformers, i.e. GT [9], SAN [21], Graphormer [38], and Gophormer [46]. We use node classification accuracy as the evaluation metric.

Implementation Details. We adopt AdamW as the optimizer and set the hyper-parameter $\epsilon$ to 1e-8 and ($\beta 1,\beta 2$) to (0.99,0.999). The peak learning rate is set to 2e-4 with a 100 epochs warm-up stage followed by a linear decay learning rate scheduler. We adopt the Variational Neighborhoods [26] with a coarsening rate of 0.01 as the default coarsening method. All models were trained on one NVIDIA Tesla V100 GPU.

Parameter Settings. In the default setting, the dropout rate is set to 0.5, the end learning rate is set to 1e-9, the hidden dimension $d$ is set to 128, the number of training epochs is set to 1,000, update period $T$ is set to 100, $N$ is set to 20, $M$ is set to 10, and the number of attention head $H$ is set as 8. We tune other hyper-parameters on each dataset based on by grid search. The searching space of batch size, number of data augmentation $\mathcal{S}$, the number of layers $L$, number of sampled nodes, number of sampled super-nodes, number global nodes are $\{8,16,32\}$, $\{4,8,16,32\}$, $\{2,3,4,5,6\}$, $\{10,15,20,25\}$, $\{0,3,6,9\}$, $\{1,2,3\}$ respectively.

Node Classification Performance. The node classification results are shown in Table 1. We apply 3 independent runs on random data splitting and report the means and standard deviations. We have the following obervations: (1) Generally, we observe that ANS-GT overperforms all Graph Transformer baselines and achieves state-of-the-art results on all 6 datasets, which demonstrates the effectiveness of our proposed model. (2) We note that some Graph Transformer baselines achieve poor performance on node classification (e.g., GT only obtains $25.68\%$ on Squirrel) compared with graph neural network models. This is probably due to the full graph attention mechanisms or the fixed node sampling schemes of existing Graph Transformers. For instance, ANS-GT achieves an accuracy of 45.88$\%$ on Squirrel while the best baseline has 43.15$\%$.

Effectiveness of Adaptive Node Sampling. Our proposed adaptive node sampling module can adjust the weights for sampling based on the rewards as the training progresses. To evaluate its effectiveness and give more insights into the ANS module, we show the normalized sampling weights as a function of training epochs on four datasets in Figure 3. Generally, we observe that the sampling weights of different sampling strategies are time-sensitive and gradually stabilize with the increase of the number of epochs. Interestingly, we find PPR and 1-hop neighbors achieves high weights on Cora while 2-hop neighbors dominate other sampling strategies on Squirrel. This may be explained by the fact that Cora and Squirrel are strong homophily/heterophily dataset respectively. For Citeseer and Actor, the weights of KNN firstly goes up and gradually decreases. This is probably due to the reason that nodes with similar attributes are most useful for the training at the beginning stage; local neighbors such as nodes with high PPR scores are more useful at the fine-tuning stage. Furthermore, Figure 4 shows the ablation studies of the adaptive node sampling module. We can observe that ANS-GT has a large advantage compared to its variant without the adaptive sampling module denoted as HGT (e.g., On Chameleon, ANS-GT achieves 65.42$\%$ while HGT only has 57.20$\%$).

Graph Coarsening Methods. In ANS-GT, we apply a hierarchical attention mechanism where the interactions between the center node and super-nodes generated by graph coarsening are considered. Here, we evaluate ANS-GT with different coarsening algorithms and different coarsening rates. Specifically, the considered coarsening algorithms include Variation Neighborhoods (VN) [26], Variation Edges (VE) [26], and Algebraic (JC) [29]. We vary the coarsening rate from 0.01 to 0.50. In Table 2, we can observe that there is no significant difference between different coarsening algorithms, indicating the robustness of ANS-GT w.r.t. them. As for the coarsening rate, the results indicate that the coarsening rate of 0.01 to 0.10 has the best performance.