Multi-hop Attention Graph Neural Networks

1 Introduction

Self-attention [Bahdanau et al., 2015; Vaswani et al., 2017] has pushed the state-of-the-art in many domains including graph presentation learning [Devlin et al., 2019]. Graph Attention Network (GAT) [Veliˇckovi´c et al., 2018] and related models [Li et al., 2018; Wang et al., 2019a; Liu et al., 2019;

∗Contact Author: xjtuwgt@gmail.com

D

!",$

A

C

%&,' = )(+&, +')

%.,/ = 0

%′&,' = 2 [%.,', %&,. , [%&,'])

%′.,/ = 2([%',/, %.,'])

Figure 1: Multi-hop attention diffusion. Consider making a prediction at nodes A and D. Left: A single GAT layer computes attention scores α between directly connected pairs of nodes (i.e., edges) and thus αD,C = 0. Furthermore, the attention αA,B between A and B only depends on node representations of A and B. Right: A single MAGNA layer: (1) captures the information of D's two-hop neighbor node C via multi-hop attention α′D,C; and (2) enhances graph structure learning by considering all paths between nodes via diffused attention, which is based on powers of graph adjacency matrix. MAGNA makes use of node D's features for attention computation between A and B. This means that two-hop attention in MAGNA is context (node D) dependent.

Oono and Suzuki, 2020] developed attention mechanism for Graph Neural Networks (GNNs), which compute attention scores between nodes connected by an edge, allowing the model to attend to messages of node's neighbors.

However, such attention computation on pairs of nodes connected by edges implies that a node can only attend to its immediate neighbors to compute its (next layer) representation. This implies that receptive field of a single GNN layer is restricted to one-hop network neighborhoods. Although stacking multiple GAT layers could in principle enlarge the receptive field and learn non-neighboring interactions, such deep GAT architectures suffer from the oversmoothing problem [Wang et al., 2019a; Liu et al., 2019; Oono and Suzuki, 2020] and do not perform well. Furthermore, edge attention in a single GAT layer is based solely on representations of the two nodes at the edge endpoints, and does not depend on their graph neighborhood context. In other words, the one-hop attention mechanism in GATs limits their ability to explore the relationship between the broader graph structure. While previous works [Xu et al., 2018; Klicpera et al., 2019b] have shown advantages in performing multi-hop message-passing in a single layer, these ap-

B

proaches are not graph-attention based. Therefore, incorporating multi-hop neighboring context into the attention computation in graph neural networks remains to be explored.

Here we present Multi-hop Attention Graph Neural Network (MAGNA), an effective multi-hop self-attention mechanism for graph structured data. MAGNA uses a novel graph attention diffusion layer (Figure 1), where we first compute attention weights on edges (represented by solid arrows), and then compute self-attention weights (dotted arrows) between disconnected pairs of nodes through an attention diffusion process using the attention weights on the edges.

Our model has two main advantages: (1) MAGNA captures long-range interactions between nodes that are not directly connected but may be multiple hops away. Thus the model enables effective long-range message passing, from important nodes multiple hops away. (2) The attention computation in MAGNA is context-dependent. The attention value in GATs [Veliˇckovi´c et al., 2018] only depends on node representations of the previous layer, and is zero between nonconnected pairs of nodes. In contrast, for any pair of nodes within a chosen multi-hop neighborhood, MAGNA computes attention by aggregating the attention scores over all the possible paths (length ≥ 1) connecting the two nodes.

Mathematically we show that MAGNA places a Personalized Page Rank (PPR) prior on the attention values. We further apply spectral graph analysis to show that MAGNA emphasizes on large-scale graph structure and lowering highfrequency noise in graphs. Specifically, MAGNA enlarges the lower Laplacian eigen-values, which correspond to the large-scale structure in the graph, and suppresses the higher Laplacian eigen-values which correspond to more noisy and fine-grained information in the graph.

We experiment on standard datasets for semi-supervised node classification as well as knowledge graph completion. Experiments show that MAGNA achieves state-of-the-art results: MAGNA achieves up to 5.7% relative error reduction over previous state-of-the-art on Cora, Citeseer, and Pubmed. MAGNA also obtains better performance on a large-scale Open Graph Benchmark dataset. On knowledge graph completion, MAGNA advances state-of-the-art on WN18RR and FB15k-237 across four metrics, with the largest gain of 7.1% in the metric of Hit at 1.

Furthermore, we show that MAGNA with just 3 layers and 6 hop wide attention per layer significantly out-performs GAT with 18 layers, even though both architectures have the same receptive field. Moreover, our ablation study reveals the synergistic effect of the essential components of MAGNA, including layer normalization and multi-hop diffused attention. We further observe that compared to GAT, the attention values learned by MAGNA have higher diversity, indicating the ability to better pay attention to important nodes.

2 Multi-hop Attention Graph Neural Network (MAGNA)

We first discuss the background and explain the novel multihop attention diffusion module and the MAGNA architecture.

2.1 Preliminaries

Let G = (V, E) be a given graph, where V is the set of Nn nodes, E ⊆ V×V is the set of Ne edges connecting M pairs of nodes in V. Each node v ∈ V and each edge e ∈ E are associated with their type mapping functions: φ : V → T and ψ : E →R. Here T and R denote the sets of node types (labels) and edge/relation types. Our framework supports learning on heterogeneous graphs with multiple elements in R.

A general Graph Neural Network (GNN) approach learns an embedding that maps nodes and/or edge types into a continuous vector space. Let X ∈ RNn×dn and R ∈ RNr×dr be the node embedding and edge/relation type embedding, where Nn = |V|, Nr = |R|, dn and dr represent the embedding dimension of node and edge/relation types, each row xi = X[i :] represents the embedding of node vi (1 ≤ i ≤ Nn), and rj = R[j :] represents the embedding of relation rj (1 ≤ j ≤ Nr). MAGNA builds on GNNs, while bringing together the benefits of Graph Attention and Diffusion techniques.

2.2 Multi-hop Attention Diffusion

We first introduce attention diffusion to compute the multihop attention directly, which operates on the MAGNA's attention scores at each layer. The input to the attention diffusion operator is a set of triples (vi, rk, vj), where vi, vj are nodes and rk is the edge type. MAGNA first computes the attention scores on all edges. The attention diffusion module then computes the attention values between pairs of nodes that are not directly connected by an edge, based on the edge attention scores, via a diffusion process. The attention diffusion module can then be used as a component in MAGNA architecture, which we will further elaborate in Section 2.3.

Edge Attention Computation At each layer l, a vector message is computed for each triple (vi, rk, vj). To compute the representation of vj at layer l+1, all messages from triples incident to vj are aggregated into a single message, which is then used to update vl+1 j .

In the first stage, the attention score s of an edge (vi, rk, vj ) is computed by the following:

where δ = LeakyReLU, W (l) h , W (l) t ∈ Rd(l)×d(l) , W (l) r ∈ Rd(l)×dr and v(l) a ∈R1×3d(l) are the trainable weights shared by l-th layer. h(l) i ∈ Rd(l) represents the embedding of node i at l-th layer, and h(0) i = xi. rk is the trainable relation embedding of the k-th relation type rk (1 ≤ k ≤ Nr), and a‖b denotes concatenation of embedding vectors a and b. For graphs with no relation type, we treat as a degenerate categorical distribution with 1 category 1 .

Applying Eq. 1 on each edge of the graph G, we obtain an attention score matrix S(l) :

1 In this case, we can view that there is only one "pseudo" relation type (category), i.e., Nr = 1

Subsequently we obtain the attention matrix A(l) by performing row-wised softmax over the score matrix S(l) : A(l) = softmax(S(l) ). A(l) ij denotes the attention value at layer l when aggregating message from node j to node i.

Attention Diffusion for Multi-hop Neighbors In the second stage, we further enable attention between nodes that are not directly connected in the network. We achieve this via the following attention diffusion procedure. The procedure computes the attention scores of multi-hop neighbors via graph diffusion based on the powers of the 1-hop attention matrix A:

where θi is the attention decay factor and θi > θi+1. The powers of attention matrix, Ai , give us the number of relation paths from node h to node t of length up to i, increasing the receptive field of the attention (Figure 1). Importantly, the mechanism allows the attention between two nodes to not only depend on their previous layer representations, but also taking into account of the paths between the nodes, effectively creating attention shortcuts between nodes that are not directly connected (Figure 1). Attention through each path is also weighted differently, depending on θ and the path length.

In our implementation we utilize the geometric distribution θi = α(1 - α)i , where α ∈ (0, 1]. The choice is based on the inductive bias that nodes further away should be weighted less in message aggregation, and nodes with different relation path lengths to the target node are sequentially weighted in an independent manner. In addition, notice that if we define θ0 = α ∈ (0, 1], A0 = I, then Eq. 3 gives the Personalized Page Rank (PPR) procedure on the graph with the attention matrix A and teleport probability α. Hence the diffused attention weights, Aij, can be seen as the influence of node j to node i. We further elaborate the significance of this observation in Section 4.3.

We can also view Aij as the attention value of node j to i since ∑Nn j=1 Aij = 1. 2 We then define the graph attention diffusion based feature aggregation as

where Θ is the set of parameters for computing attention. Thanks to the diffusion process defined in Eq. 3, MAGNA uses the same number of parameters as if we were only computing attention between nodes connected via edges. This ensures runtime efficiency (refer to Appendix 3 A for complexity analysis) and model generalization.

Approximate Computation for Attention Diffusion For large graphs computing the exact attention diffusion matrix A using Eq. 3 may be prohibitively expensive, due to computing the powers of the attention matrix [Klicpera et al., 2019a]. To resolve this bottleneck, we proceed as follows: Let H(l) be the input entity embedding of the l-th layer (H(0) = X)

2 Obtained by the definition A(l) = softmax(S(l) ) and Eq. 3.

3 Appendix can be found via https://arxiv.org/abs/2009.14332

and θi = α(1 - α)i . Since MAGNA only requires aggregation via AH(l) , we can approximate AH(l) by defining a sequence Z (K) which converges to the true value of AH(l) (Proposition 1) as K → ∞:

where 0 ≤ k < K.

In the Appendix we give the proof which relies on the expansion of Eq. 5.

Using the above approximation, the complexity of attention computation with diffusion is still O(|E|), with a constant factor corresponding to the number of hops K. In practice, we find that choosing the values of K such that 3 ≤ K ≤ 10 results in good model performance. Many real-world graphs exhibit small-world property, in which case even a smaller value of K is sufficient. For graphs with larger diameter, we choose larger K, and lower the value of α.

2.3 Multi-hop Attention based GNNArchitecture

Figure 2 provides an architecture overview of the MAGNA Block that can be stacked multiple times.

Multi-headGraphAttentionDiffusionLayer Multi-head attention [Vaswani et al., 2017; Veliˇckovi´c et al., 2018] is used to allow the model to jointly attend to information from different representation sub-spaces at different viewpoints. In Eq. 6, the attention diffusion for each head i is computed separately with Eq. 4, and aggregated:

where ‖ denotes concatenation and Θi are the parameters in Eq. 1 for the i-th head (1 ≤ i ≤ M), Wo represents a parameter matrix, and LN = LayerNorm. Since we calculate the attention diffusion in a recursive way in Eq. 5, we add layer normalization which helpful to stabilize the recurrent computation procedure [Ba et al., 2016].

Deep Aggregation Moreover our MAGNA block contains a fully connected feed-forward sub-layer, which consists of a two-layer feed-forward network. We also add the layer normalization and residual connection in both sub-layers, allowing for a more expressive aggregation step for each block [Xiong et al., 2020]:

MAGNA generalizes GAT MAGNA extends GAT via the diffusion process. The feature aggregation in GAT is H(l+1) = σ(AH(l)W (l) ), where σ represents the activation function. We can divide GAT layer into two components as follows:

In component (1), MAGNA removes the restriction of attending to direct neighbors, without requiring additional parameters as A is induced from A. For component (2) MAGNA uses layer normalization and deep aggregation to achieve higher expressive power compared to elu nonlinearity in GAT.

3 Analysis of Graph Attention Diffusion

In this section, we investigate the benefits of MAGNA from the viewpoint of discrete signal processing on graphs [Sandryhaila and Moura, 2013]. Our first result demonstrates that MAGNA can better capture large-scale structural information. Our second result explores the relation between MAGNA and Personalized PageRank (PPR).

3.1 Spectral Properties of Graph Attention Diffusion

We view the attention matrix A of GAT, and A of MAGNA as weighted adjacency matrices, and apply Graph Fourier transform and spectral analysis (details in Appendix) to show the effect of MAGNA as a graph low-pass filter, being able to more effectively capture large-scale structure in graphs. By Eq. 3, the sum of each row of either A or A is 1. Hence the normalized graph Laplacians are ˆLsym = I - A and Lsym = I -A for A and A respectively. We can get the following proposition:

Proposition 2. Let ˆλgi and λgi be the i-th eigeinvalues of ˆLsym and Lsym.

Refer to Appendix for the proof. We additionally have λgi ∈ [0, 2] (proved by [Ng et al., 2002]). Eq. 9 shows that when λgi is small such that α 1-α + λgi < 1, then ˆλgi > λgi , whereas for large λgi , ˆλgi < λgi . This relation indicates that the use of A increases smaller eigenvalues and decreases larger eigenvalues 4 . See Section 4.3 for its empirical evidence. The low-pass effect increases with smaller α.

The eigenvalues of the low-frequency signals describe the large-scale structure in the graph [Ng et al., 2002] and have been shown to be crucial in graph tasks [Klicpera et al., 2019b]. As λgi ∈ [0, 2] [Ng et al., 2002] and α 1-α > 0, the reciprocal format in Eq. 9 will amplify the ratio of lower eigenvalues to the sum of all eigenvalues. In contrast, high eigenvalues corresponding to noise are suppressed.

3.2 Personalized PageRankMeets Graph Attention Diffusion

We can also view the attention matrix A as a random walk matrix on graph G since ∑Nn j=1 Ai,j = 1 and Ai,j > 0. If we perform Personalized PageRank (PPR) with parameter α ∈ (0, 1] on G with transition matrix A, the fully Personalized PageRank [Lofgren, 2015] is defined as:

Using the power series expansion for the matrix inverse, we obtain

Comparing to the diffusion Equation 3 with θi = α(1-α)i , we have the following proposition.

Proposition 3. Graph attention diffusion defines a personalized page rank with parameter α ∈ (0, 1] on G with transition matrix A, i.e., A =Appr.

The parameter α in MAGNA is equivalent to the teleport probability of PPR. PPR provides a good relevance score between nodes in a weighted graph (the weights from the attention matrix A). In summary, MAGNA places a PPR prior over node pairwise attention scores: the diffused attention between node i and j depends on the attention scores on the edges of all paths between i and j.

4 Experiments

We evaluate MAGNA on two classical tasks 5 : (1) on node classification we achieve an average of 5.7% relative error reduction; (2) on knowledge graph completion we achieve 7.1% relative improvement in the Hit@1 metric. 6

4.1 Task 1: Node Classification

Datasets We employ four benchmark datasets for node classification: (1) standard citation network benchmarks Cora, Citeseer and Pubmed [Sen et al., 2008; Kipf and Welling, 2016]; and (2) a benchmark dataset ogbn-arxiv on

4 The eigenvalues of A and A correspond to the same eigenvectors, as shown in Proposition 2 in Appendix.

5 All datasets are public; code will be released after publication.

6 Please see the definitions of these two tasks in Appendix.

Table 1: Node classification accuracy on Cora, Citeseer, Pubmed. MAGNA achieves state-of-the-art.

170k nodes and 1.2m edges from the Open Graph Benchmark [Weihua Hu, 2020]. We follow the standard data splits for all datasets. Further information about these datasets is summarized in the Appendix.

Baselines We compare against a comprehensive suite of state-of-the-art GNN methods including: GCNs [Kipf and Welling, 2016], Chebyshev filter based GCNs [Defferrard et al., 2016], DualGCN [Zhuang and Ma, 2018], JKNet [Xu et al., 2018], LGCN [Gao et al., 2018], Diffusion-GCN (Diff-GCN) [Klicpera et al., 2019b], APPNP [Klicpera et al., 2019a], Graph U-Nets (g-U-Nets) [Gao and Ji, 2019], and GAT [Veliˇckovi´c et al., 2018].

Experimental Setup For datasets Cora, Citeseer and Pubmed, we use 6 MAGNA blocks with hidden dimension 512 and 8 attention heads. For the large-scale ogbn-arxiv dataset, we use 2 MAGNA blocks with hidden dimension 128 and 8 attention heads. Refer to Appendix for detailed description of all hyper-parameters and evaluation settings.

Results MAGNA achieves the best on all datasets (Tables 1 and 2) 7 , out-performing mutlihop baselines such as Diffusion GCN, APPNP and JKNet. The baseline performance and their embedding dimensions are from the previous papers. Appendix Table 6 further demonstrates that large 512 dimension embedding only benefits the expressive MAGNA, whereas GAT and Diffusion GCN performance degrades.

Ablation study We report (Table 1) the model performance after removing each component of MAGNA (layer normalization, attention diffusion and feed forward layers) from every MAGNA layer. Note that the model is equivalent to GAT without these three components. We observe that diffusion and layer normalization play a crucial role in improving the node classification performance for all datasets. Since MAGNA computes attention diffusion in a recursive manner, layer normalization is crucial in ensuring training stability [Ba et al., 2016]. Meanwhile, comparing to GAT (see the next-to-last row of Table 1), attention diffusion allows multihop attention in every layer to benefit node classification.

7 We also compared to GAT and Diffusion-GCN (with LayerNorm and feed-forward Layer) over random splits in Appendix.

4.2 Task 2: Knowledge Graph Completion

Datasets We evaluate MAGNA on standard benchmark knowledge graphs: WN18RR [Dettmers et al., 2018] and FB15K-237 [Toutanova and Chen, 2015]. See the statistics of these KGs in Appendix.

Baselines We compareMAGNAwith state-of-the-art baselines, including (1) translational distance based models: TransE [Bordes et al., 2013] and its latest extension RotatE [Sun et al., 2019], OTE [Tang et al., 2020] and ROTH [Chami et al., 2020]; (2) semantic matching based models: ComplEx [Trouillon et al., 2016], QuatE [Zhang et al., 2019], CoKE [Wang et al., 2019b], ConvE [Dettmers et al., 2018], DistMult [Yang et al., 2015], TuckER [Balazevic et al., 2019] and AutoSF [Zhang et al., 2020]; (3) GNN-based models: R-GCN [Schlichtkrull et al., 2018], SACN [Shang et al., 2019] and A2N [Bansal et al., 2019].

Experimental Setup We use the multi-layer MAGNA as encoder for both FB15k-237 and WN18RR. We randomly initialize the entity embedding and relation embedding as the input of the encoders, and set the dimensionality of the initialized entity/relation vector as 100 used in DistMult [Yang et al., 2015]. We select other MAGNA model hype-parameters, including number of layers, hidden dimension, head number, top-k, learning rate, hop number, teleport probability α and dropout ratios (see the settings of these parameter in Appendix), by a random search during the training.

Training procedure We use the standard training procedure used in previous KG embedding models [Balazevic et al., 2019; Dettmers et al., 2018] (Appendix for details). We follow an encoder-decoder framework: The encoder applies the proposed MAGNA model to compute the entity embeddings. The decoder makes link prediction given the embeddings. To show the power of MAGNA, we employ a simple decoder DistMult [Yang et al., 2015].

Evaluation We use the standard split for the benchmarks, and the standard testing procedure of predicting tail (head) entity given the head (tail) entity and relation type. We exactly follow the evaluation used by all previous works, namely the Mean Reciprocal Rank (MRR), Mean Rank (MR), and hit rate at K (H@K). See Appendix for a detailed description of this standard setup.

Results MAGNA achieves new state-of-the-art in knowledge graph completion on all four metrics (Table 3). MAGNA compares favourably to both the most recent shallow embedding methods (QuatE), and deep embedding methods (SACN). Note that with the same decoder (DistMult), MAGNA using its own embeddings achieves drastic improvements over using the corresponding DistMult embeddings.

4.3 MAGNAModel Analysis

Here we present (1) spectral analysis results, (2) robustness to hyper-parameter changes, and (3) attention distribution analysis to show the strengths of MAGNA.

Spectral Analysis: Why MAGNA works for node classification? We compute the eigenvalues of the graph Laplacian of the attention matrix A, ˆλgi , and compare to that of the diffused matrix A, λgi . Figure 3 (a) shows the ratio ˆλgi /λgi on the Cora dataset. Low eigenvalues corresponding to large-scale structure in the graph are amplified (up to a factor of 8), while high eigenvalues corresponding to eigenvectors with noisy information are suppressed [Klicpera et al., 2019b].

MAGNA Model Depth Here we conduct experiments by varying the number of GCN, Diffusion-GCN (PPR based) GAT and our MAGNA layers to be 3, 6, 12, 18 and 24 for node classification on Cora. Results in Fig. 3 (b) show that deep GCN, Diffusion-GCN and GAT (even with residual connection) suffer from degrading performance, due to the over-smoothing problem [Li et al., 2018; Wang et al., 2019a]. In contrast, the MAGNA model achieves consistent best results even with 18 layers, making deep MAGNA model robust and expressive. Notice that GAT with 18 layers cannot out-perform MAGNA with 3 layers and K=6 hops, although they have the same receptive field.

Effect of K and α Figs. 3 (c) and (d) report the effect of hop number K and teleport probability α on model performance. We observe significant increase in performance when considering multi-hop neighbors information (K > 1). However, increasing the hop number K has a diminishing returns, for K ≥ 6. Moreover, we find that the optimal K is correlated with the largest node average shortest path distance (e.g., 5.27 for Cora). This provides a guideline for choosing the best K.

We also observe that the accuracy drops significantly for larger α > 0.25. This is because small α increases the lowpass effect (Fig. 3 (a)). However, α being too small causes the model to only focus on the most large-scale graph structure and have lower performance.

Attention Distribution Last we also analyze the learned attention scores of GAT and MAGNA. We first define a discrepancy metric over the attention matrix A for node vi as ∆i = ‖A[i,:]-Ui‖ degree(vi) [Shanthamallu et al., 2020], where Ui is the uniform distribution score for the node vi. ∆i gives a measure of how much the learnt attention deviates from an uninformative uniform distribution. Large ∆i indicates more meaningful attention scores. Fig. 4 shows the distribution of the discrepancy metric of the attention matrix of the 1st head w.r.t. the first layer of MAGNA and GAT. Observe that attention scores learned in MAGNA have much larger discrepancy. This shows that MAGNA is more powerful than GAT in distinguishing important nodes and assigns attention scores accordingly.

5 RelatedWork

Our proposed MAGNA belongs to the family of Graph Neural Network (GNN) models [Battaglia et al., 2018; Wu et al., 2020; Kipf and Welling, 2016; Hamilton et al., 2017], while taking advantage of graph attention and diffusion techniques.

Graph Attention Neural Networks (GATs) generalize attention operation to graph data. GATs allow for assigning different importance to nodes of the same neighborhood at the feature aggregation step [Veliˇckovi´c et al., 2018]. Based on such framework, different attention-based GNNs have been proposed, including GaAN [Zhang et al., 2018], AGNN [Thekumparampil et al., 2018], GeniePath [Liu et al., 2019]. However, these models only consider direct neighbors for each layer of feature aggregation, and suffer from oversmoothing when they go deep [Wang et al., 2019a].

Diffusion based Graph Neural Network Recently Graph Diffusion Convolution [Klicpera et al., 2019b; Klicpera et al., 2019a] proposes to aggregate information from a larger (multi-hop) neighborhood at each layer, by sparsifying a generalized form of graph diffusion. This idea was also explored in [Liao et al., 2019; Luan et al., 2019; Xhonneux et al., 2020; Klicpera et al., 2019a] for multi-scale Graph Convolutional Networks. However, these methods do not incorporate attention mechanism which is crucial to model performance, and do not make use of edge embeddings (e.g., Knowledge graph) [Klicpera et al., 2019b]. Our approach defines a novel multi-hop context-dependent self-attention GNN which resolves the over-smoothing issue of GAT architectures [Wang et al., 2019a]. [Isufi et al., 2020; Cucurull et al., 2018; Feng et al., 2019] also extends attention mechanism for multihop information aggregation, but they require different set of parameters to compute the attention to neighbors of different hops, making these approaches much more expensive compared to MAGNA, and were not extended to the knowledge graph settings.

6 Conclusion

We proposed Multi-hop Attention Graph Neural Network (MAGNA), which brings together benefits of graph attention and diffusion techniques in a single layer through attention diffusion, layer normalization and deep aggregation. MAGNA enables context-dependent attention between any pair of nodes in the graph in a single layer, enhances largescale structural information, and learns more informative attention distribution. MAGNA improves over all state-of-theart methods on the standard tasks of node classification and knowledge graph completion.