Coarformer: Transformer for large graph via graph coarsening

1 INTRODUCTION . In recent years , the Transformer architecture ( Vaswani et al. , 2017 ) has been derived into several variants , e.g. , BERT ( Devlin et al. , 2019 ) and ViT ( Dosovitskiy et al. , 2021 ) , which achieve unprecedented successes in natural language processing ( NLP ) and computer vision ( CV ) , respectively . Some recent works ( Kreuzer et al. , 2021 ; Ye & Ji , 2021 ; Ying et al. , 2021a ) attempt to generalize Transformer for graph data by treating each node as a token and designing dedicated positional encoding for the nodes . These works ’ performance has surpassed that of graph neural networks ( GNN ) on an increasing number of graph-related tasks , particularly molecular property prediction ( Ying et al. , 2021b ) . These works , however , demonstrate their superiority on small graphs , such as molecules with tens of atoms ( i.e. , nodes ) . When applied to a large graph , these Transformer-based methods ( Dwivedi & Bresson , 2021 ; Zhang et al. , 2020 ) explicitly or implicitly restrict each node ’ s receptive field to its neighbors . Nevertheless , large graphs , such as Arxiv and Products in OGBN ( Hu et al. , 2020 ) , also demand the global receptive field and the powerful expressiveness of the Transformer architecture . These graphs have more than hundreds of thousands of nodes and graph diameters larger than twenty , where the needed receptive field should be large , and the correlations among the nodes ’ features are complex . So , what limits the applicability of Transformer on large graphs ? Here , we identify two main obstacles : ( 1 ) Transformer is built on a node-to-node attention mechanism . When applied to a large graph with numerous nodes , the massive distant nodes can divert a significant portion of the attention no matter whether they are indeed related . As a result , the target node can neglect its local neighborhood , which is indispensable for learning generalizable node representations . Hence , Transformer is prone to causing over-fitting on large graphs , which will be exacerbated for the semi-supervised node-level tasks where the number of labeled nodes is limited . We will provide some empirical evidences for this point in Sec . 5.1.1 and Appendix A.1 . ( 2 ) The global receptive field of Transformer is costly , where the pairwise interactions among tokens lead to quadratic computational complexity regarding the number of nodes . Although some Sparse Transformer methods ( Roy et al. , 2021 ; Kitaev et al. , 2019 ; Ren et al. , 2021 ) can improve the efficiency of the original Transformer , they have not exploited the unique characteristics of graph data and require a quadratic or at least sub-quadratic space complexity , which is still unaffordable in most practical cases . We argue that both local and global information is useful for encoding each node in a large graph . Meanwhile , we notice that GNN and Transformer are proficient at capturing the local and global information , respectively , but neither of them can easily extract both kinds of information . Thus , a straightforward strategy is to employ them together to play to their strengths simultaneously . To this end , we propose to wipe out the discussed obstacles of applying Transformer by letting it work on a down-sampled graph that preserves both the sketch of global graph structure and the aggregated node features yet has much fewer nodes . And the similar idea is used in graph pooling ( Ying et al. , 2018 ; Baek et al. , 2020 ) and hierarchical GNN ( Fang et al. , 2020 ; Sobolevsky , 2021 ) . In this way , we essentially trade some fine-grained long-range information for the efficiency of Transformer . To encode both local and global information for each node , we propose a two-view architecture Coarformer consisting of a fine-grained local view and a coarse global view . In the local view , we can apply a GNN-based module to the original input graph to encode each node by its local topological structures and the node features . In the global view , we apply a Transformer-based module to a coarse graph produced by an adopted graph coarsening algorithm . Such kinds of algorithms mimic a down-sampling to the original graph via grouping the nodes into a less number of super-nodes . The Transformer-based module makes pairwise interactions among these super-nodes , capturing the coarse but long-range dependencies of the original graph . Additionally , we design a cross-view propagation strategy for these two views in order to facilitate their interaction . As a result , Coarformer encodes each node by both its local and global information , combining GNN and Transformer ’ s merits . It is worth noting that the number of super-nodes balances the efficiency of the Transformer-based module and how coarse the global view is . In practice , a small number that reduces the computational complexity from quadratic to linear w.r.t . the number of original nodes can still provide helpful global information , making Coarformer applicable for large graphs . We conduct extensive experiments on real-world datasets to study the proposed Coarformer . At first , the combinations of a GNN model with various Transformer-based methods consistently surpass a single GNN model , which supports our idea of encoding nodes by local and global information . More importantly , Coarformer achieves the best performances with the least running time and GPU memory consumption , compared to the related Sparse Transformer-based methods . Meanwhile , we empirically show that the local and global information is complementary , and our coarse global view can boost the performances of various kinds of GNNs . 2 BACKGROUND AND MOTIVATIONS . In this section , we summarize GNN , Transformer for graphs , and graph coarsening , along with a discussion on their connections to motivate our method . Before that , we firstly present notations for graph data . Let G = ( V , E ) be a undirected unweighted graph , where V is the node set , E is the edge set . We denote the number of nodes ( i.e. , |V| ) by n and use X ∈ Rn×k0 to denote the feature matrix whose i-th row Xi , : is the feature vector of node vi . We use A ∈ { 0 , 1 } n×n to denote the adjacency matrix of G , where Aij = 1 if and only if ( vi , vj ) ∈ E . The degree of a node vi is denoted by di = ∑n j=1 Aij , and the degree matrix D is a diagonal matrix with di being its i-th diagonal entry . GNN . To represent nodes based on their features ( i.e. , X ) and topological structures ( i.e. , A ) , most GNN models stack the message passing layers to calculate node representations H ( l ) by recursively aggregating the representations of its neighbors . Taking GCN ( Kipf & Welling , 2017 ) for example , it defines this procedure as H ( 0 ) = X , H ( l ) = D̃− 1 2 ÃD̃ 1 2H ( l−1 ) W ( l ) , where D̃ = D + In , Ã = A + In , and W ( l ) ∈ Rk ( l−1 ) ×kl is a learnable parameter . In such a way , an L-layer GNN model calculates H ( L ) to represent the nodes , where the receptive field for each node is its L-hop neighborhood . To capture long-range information , we need to increase L , where a large depth L is prone to causing over-smoothing ( Li et al. , 2018 ) and bottleneck ( Alon & Yahav , 2021 ) issues . Transformer for graphs . Therefore , some recent works apply the Transformer architecture to capture useful long-range information in the graph , which relies solely on the attention mechanism , dispensing with the message passing used in GNN . Taking the scaled dot-product attention for example , the node representations are updated as the weighted sums of themselves : H ( l ) = Softmax ( QKT√ m ) V = Softmax ( ( H ( l−1 ) W ( l ) Q ) ( H ( l−1 ) W ( l ) K ) T √ m ) H ( l−1 ) W ( l ) , ( 1 ) where WQ , WK ∈ Rk ( l−1 ) ×m , and W ∈ Rk ( l−1 ) ×kl are learnable parameters . Comparing Eq . 1 with the updating equation used in GCN , Softmax ( QKT√ m ) is usually allowed to be a dense n × n matrix , while the ( i , j ) -th entry of D̃− 1 2 ÃD̃ 1 2 is non-zero only if Ãij 6= 0 . Thus , one attention layer is sufficient for each h ( l−1 ) j to attend to any h ( l ) i , no matter how far away vj is to vi . Although these pairwise interactions provide the global receptive field for Transformer to capture long-range information , they often distract the necessary attention of each node from its neighborhood , especially when n is large . Besides , these pairwise interactions requireO ( n2 ) computational complexity , which further hinders the application of Transformer to large graphs . Graph coarsening . To simplify a given graph ( mainly reducing the number of nodes ) while preserving its global information as much as possible , graph coarsening ( Ron et al. , 2011 ; Loukas , 2019 ) produces a partition P = { C1 , . . . , Cn′ } of V and regards each cluster Ci as a super-node . Consequently , we get a coarse graph G′ = ( V ′ , E ′ ) , where |V ′| = n′ and ( Ci , Cj ) ∈ E ′ if and only if ∃vk ∈ Ci , vl ∈ Cj , s.t . ( vk , vl ) ∈ E . Meanwhile , this partition can be characterized by a matrix P̂ ∈ { 0 , 1 } n×n′ , with P̂ij = 1 if and only if vi ∈ Cj . Then its normalized version can be defined by P = P̂C− 1 2 , where C is a n′×n′ diagonal matrix with |Ci| as its i-th diagonal entry . We can define the feature matrix and weighted adjacency matrix for G′ by X ′ = P TX and A′ = P TAP , which characterize a coarse global view of the original graph . We define the coarsening rate as c = n ′ n . 3 METHODOLOGY . In this section , we introduce a two-view architecture Coarformer composed of a fine-grained local view and a coarse global view . We present the overview of Coarformer in Figure 1 . At first , we preprocess an input graph G via graph coarsening to generate a coarse graph G′ . Then a GNN-based module and a Transformer-based module work on G and G′ , respectively , which can be interpreted as encoding the nodes from the fine-grained local view and the coarse yet global view . We further design a scheme for Coarformer to enable neural messages to propagate across these two views , so the resulting node representations synthesize both the local and the global information . In the fine-grained local view , any GNN model ( e.g. , GCN ) can be adopted , which recursively calculates the node representations H ( l ) , l = 1 , . . . , L. As for the coarse global view , the details about the Transformer-based module are presented in Sec . 3.1 . Then we elaborate on how to enhance the communication between these two views in Sec . 3.2 . Lastly , we present how Coarformer makes inference and how to optimize Coarformer in Sec . 3.3 .

Summary: This paper proposes to use the transformer architecture on the coarse graph obtained from the graph coarsening algorithm.  To be more specific, * Since directly using the transformer architecture to the large-scale graph is computationally prohibitive, the authors use the existing graph coarsening algorithms to the large graphs, and then use the transformer on the coarse graphs, which can capture the global information of the given graph in contrast to the GNNs capturing local information. * To simultaneously use both local and global structures of the graph, the authors first separately capture the local and global information of the given graph with GNNs on the original graph and transformers on the coarse graph, respectively, and then propagate information of local and global to each other with the cross-view propagation scheme. * The authors show that the proposed Coarformer outperforms baselines on node classification tasks, and also show that the proposed model is highly efficient against the global transformer models for graphs.