FAST GRAPH ATTENTION NETWORKS USING EFFECTIVE RESISTANCE BASED GRAPH SPARSIFICATION

1 INTRODUCTION . Graphs are efficient representations of pairwise relations , with many real-world applications including product co-purchasing network ( ( McAuley et al. , 2015 ) ) , co-author network ( ( Hamilton et al. , 2017b ) ) , etc . Graph neural networks ( GNN ) have become popular as a tool for inference from graph based data . By leveraging the geometric structure of the graph , GNNs learn improved representations of the graph nodes and edges that can lead to better performance in various inference tasks ( ( Kipf & Welling , 2016 ; Hamilton et al. , 2017a ; Veličković et al. , 2018 ) ) . More recently , the attention mechanism has demonstrated superior performance for inference over nodes in GNNs ( ( Veličković et al. , 2018 ; Xinyi & Chen , 2019 ; Thekumparampil et al. , 2018 ; Lee et al. , 2020 ; Bianchi et al. , 2019 ; Knyazev et al. , 2019 ) ) . However , attention based GNNs suffer from huge computational cost . This may hinder the applicability of the attention mechanism to large graphs . GNNs generally rely on graph convolution operations . For a graph G with N nodes , graph convolution with a kernel gw : R ! R is defined as gw ? h = Ugw ( ⇤ ) U > h ( 1 ) where U is the matrix of eigenvectors and ⇤ is the diagonal matrix of the eigenvalues of the normalized graph Laplacian matrix defined as Lnorm = I D 1/2AD 1/2 , ( 2 ) with D and A being the degree matrix and the adjacency matrix of the graph , and gw is applied elementwise . Since computing U and ⇤ can be very expensive ( O ( N3 ) ) , most GNNs use an approximation of the graph convolution operator . For example , in graph convolution networks ( GCN ) ( Kipf & Welling , 2016 ) , node features are updated by computing averages as a first order approximation of Eq.equation 1 over the neighbors of the nodes . A single neural network layer is defined as : H ( l+1 ) GCN = ⇣ eD 1/2 eA eD 1/2H ( l ) W ( l ) ⌘ , ( 3 ) where H ( l ) and W ( l ) are the activations and the weight matrix at the lth layer respectively and eA = A+ I and eD is the degree matrix of eA . Attention based GNNs add another layer of complexity : they compute pairwise attention coefficients between all connected nodes . This process can significantly increase the computational burden , especially on large graphs . Approaches to speed up GNNs were proposed in ( Chen et al. , 2018 ; Hamilton et al. , 2017a ) . However , these sampling and aggregation based methods were designed for simple GCNs and are not applicable to attention based GNNs . There has also been works in inducing sparsity in attention based GNNs ( Ye & Ji , 2019 ; Zheng et al. , 2020 ) , but they focus on addressing potential overfitting of attention based models rather than scalability . In this paper , we propose Fast Graph Attention neTwork ( FastGAT ) , an edge-sampling based method that leverages effective resistances of edges to make attention based GNNs lightweight . The effective resistance measures importance of the edges in terms of preserving the graph connectivity . FastGAT uses this measure to prune the input graph and generate a randomized subgraph with far fewer edges . Such a procedure preserves the spectral features of a graph , hence retaining the information that the attention based GNNs need . At the same time , the graph is amenable to more complex but computationally intensive models such as attention GNNs . With the sampled subgraph as their inputs , the attention based GNNs enjoy much smaller computational complexity . Note that FastGAT is applicable to all attention based GNNs . In this paper , we mostly focus on the Graph Attention NeTwork model ( GAT ) proposed by ( Veličković et al. , 2018 ) . However we also show FastGAT is generalizable to two other attention based GNNs , namely the cosine similarity based approach ( Thekumparampil et al. , 2018 ) and Gated Attention Networks ( Zhang et al. , 2018 ) . We note that Graph Attention Networks can be re-interpreted as convolution based GNNs . We show this explicitly in the Appendix . Based on this re-interpretation , we theoretically prove that spectral sparsification preserves the feature representations computed by the GAT model . We believe this interpretation also opens up interesting connections between sparsifying state transition matrices of random walks and speeding up computations in GNNs . The contributions of our paper are as outlined below : • We propose FastGAT , a method that uses effective resistance based spectral graph sparsification to accelerate attention GNNs in both inductive and transductive learning tasks . The rapid subsampling and the spectrum preserving property of FastGAT help attention GNNs retain their accuracy advantages and become computationally light . • We provide a theoretical justification for using spectral sparsification in the context of attention based GNNs by proving that spectral sparsification preserves the features computed by GNNs . • FastGAT outperforms state-of-the-art algorithms across a variety of datasets under both transductive and inductive settings in terms of computation , achieving a speedup of up to 10x in training and inference time . On larger datasets such as Reddit , the standard GAT model runs out of memory , whereas FastGAT achieves an F1 score 0.93 with 7.73 second per epoch time in training . • Further , FastGAT is generalizable to other attention based GNNs such as the cosine similarity based attention ( Thekumparampil et al. , 2018 ) and the Gated Attention Network ( Zhang et al. , 2018 ) . 2 RELATED WORK . Accelerating graph based inference has drawn increasing interest . Two methods proposed in ( Chen et al. , 2018 ) ( FastGCN ) and ( Huang et al. , 2018 ) speed up GCNs by using importance sampling to sample a subset of nodes per layer during training . Similarly , GraphSAGE ( Hamilton et al. , 2017a ) also proposes an edge sampling and aggregation based method for inductive learning based tasks . All of the above works use simple aggregation and target simple GCNs , while our work focus on more recent attention based GNNs such as ( Veličković et al. , 2018 ) . We are able to take advantage of the attention mechanism , while still being computationally efficient . Graph sparsification aims to approximate a given graph by a graph with fewer edges for efficient computation . Depending on final goals , there are cut-sparsifiers ( ( Benczúr & Karger , 1996 ) ) , pairwise distance preserving sparsifiers ( ( Althöfer et al. , 1993 ) ) and spectral sparsifiers ( ( Spielman & Teng , 2004 ; Spielman & Srivastava , 2011 ) ) , among others ( ( Zhao , 2015 ; Calandriello et al. , 2018 ; Hübler et al. , 2008 ; Eden et al. , 2018 ; Sadhanala et al. , 2016 ) ) . In this work , we use spectral sparsification to choose a randomized subgraph while preserving spectral properties . Apart form providing the strongest guarantees in preserving graph structure ( ( Chu et al. , 2018 ) ) , they align well with GNNs due to their connection to spectral graph convolutions . Graph sparsification on neural networks have been studied recently ( ( Ye & Ji , 2019 ; Zheng et al. , 2020 ; Ioannidis et al. , 2020 ; Louizos et al. , 2017 ) ) . However , their main goal is to alleviate overfitting in GNNs not reducing the training time . They still require learning attention coefficients and binary gate values for all edges in the graph , hence not leading to any computational or memory benefit . In contrast , FastGAT uses a fast subsampling procedure , thus resulting in a drastic improvement in training and inference time . It is also highly stable in terms of training and inference . 3 FASTGAT : ACCELERATING GRAPH ATTENTION NETWORKS VIA EDGE SAMPLING . 3.1 THE FASTGAT ALGORITHM . Let G ( E , V ) be a graph with N nodes and M edges . An attention based GNN computes attention coefficients ↵i , j for every pair of connected nodes i , j 2 V in every layer ` . The ↵i , j ’ s are then used as averaging weights to compute the layer-wise feature updates . In the original GAT formulation , the attention coefficients are ↵ij = exp LeakyReLU ( a > [ Whi||Whj ] ) P j2Ni exp ( LeakyReLU ( a > [ Whi||Whj ] ) ) , ( 4 ) where hi ’ s are the input node features to the layer , W and a are linear mappings that are learnt , Ni denotes the set of neighbors of node i , and || denotes concatenation . With the ↵ij ’ s as defined above , the node-i output embedding of a GAT layer is h 0 i = 0 @ X j2Ni ↵ijWhj 1 A . ( 5 ) For multi-head attention , the coefficients are computed independently in each attention head with head-dependent matrices W and attention vector a . Note that the computational burden in GATs arises directly from computing the ↵i , j ’ s in every layer , every attention head and every forward pass during training . Goal : Our objective is to achieve performance equivalent to that of full graph attention networks ( GAT ) , but with only a fraction of the original computational complexity . This computational saving is achieved by reducing the number of attention computations . Idea : We propose to use edge-sampling functions that sparsify graphs by removing nonessential edges . This leads to direct reduction in the number of attention coefficients to be computed , hence reducing the burden . Choosing the sampling function is crucial for retaining the graph connectivity . Let EdgeSample ( E , A , q ) denote a randomized sampling function that , given an edge set E , adjacency matrix A and a number of edges to be sampled q , returns a subset of the original edge set Es ⇢ E with |Es| = q . Our algorithm then uses this function to sparsify the graph in every layer and attention head . Following this , the attention coefficients are computed only for the remaining edges . A more detailed description is given in Algorithm 1 . In every layer and attention head , a randomized subgraph with q ⌧ M edges is generated and the attention coeffients are computed only for this subset of edges . We use a specialized distribution that depends on the contribution of each edge to the graph connectivity . We provide further details in Section 3.2 . Note that in the general description below , the attention coefficients themselves are used as weights for sparsification and the reweighted attention coefficients are used to compute the feature update . Doing so helps in theoretical analysis of the algorithm . However in practice , we replace this expensive procedure with a one-time sampling of the graph with the original edge weights and compute the attention coefficients for only the remaining edges . In particular , we use two simpler variations of FastGAT include : i ) FastGAT-const , where the subgraph g is kept constant in all the layers and attention heads and ii ) FastGAT-layer , where the subgraph is different in each layer ( drawn stochastically from the original edge weights ) , but the same across all the attention heads within a layer . Algorithm 1 : The FastGAT Algorithm Input : Graph G ( V , E ) , Num . layers = L , Num . Attention heads K ( ` ) , ` = 1 , · · · , L Initial Weight matrices W ( ` ) , Non-linearity , Feature matrix H 2 RN⇥D Randomized edge sampling function EdgeSample ( · ) , Attention function ( · ) Num . edges sampled q for each layer ` do for each attention head k 2 { 1 , 2 , · · · , K ( ` ) } do Compute attention matrix ( ` ) k 2 RN⇥N , with ( ` ) k ( i , j ) = ✓k ( h ( ` ) i , h ( ` ) j ) Sample a graph ̂ ( ` ) k = EdgeSample ( ( ` ) k , A , q ) Compute H ( ` +1 ) k = ⇣ ̂ ( ` ) k H ( ` ) k W ( ` ) ⌘ H ( ` +1 ) = || k H ( ` ) k // Concatenate the output of attention heads Compute loss and update W ’ s // gradient based weight update

Summary: This paper proposes a paradigm which speeds up the training/inference time of GATs while not compromising too much performance. The method adopts a layerwise sampling procedure. In particular. The authors propose to sample a sub-portion of edges for each layer based on their effective resistance. Such sampling keeps the spectral similar to the original results theoretically and gives a guarantee to the performance drop.