Topology-Aware Pooling via Graph Attention

Pooling operations have shown to be effective on various tasks in computer vision and natural language processing . One challenge of performing pooling operations on graph data is the lack of locality that is not well-defined on graphs . Previous studies used global ranking methods to sample some of the important nodes , but most of them are not able to incorporate graph topology information in computing ranking scores . In this work , we propose the topology-aware pooling ( TAP ) layer that uses attention operators to generate ranking scores for each node by attending each node to its neighboring nodes . The ranking scores are generated locally while the selection is performed globally , which enables the pooling operation to consider topology information . To encourage better graph connectivity in the sampled graph , we propose to add a graph connectivity term to the computation of ranking scores in the TAP layer . Based on our TAP layer , we develop a network on graph data , known as the topology-aware pooling network . Experimental results on graph classification tasks demonstrate that our methods achieve consistently better performance than previous models . 1 INTRODUCTION . Pooling operations have been widely applied in various fields such as computer vision ( He et al. , 2016 ; Huang et al. , 2017 ) , and natural language processing ( Zhang et al. , 2015 ) . Pooling operations can effectively reduce dimensional sizes ( Simonyan & Zisserman , 2015 ) and enlarge receptive fields ( Chen et al. , 2016 ) . The application of regular pooling operations depends on the well-defined spatial locality in grid-like data such as images and texts . However , it is still challenging to perform pooling operations on graph data . In particular , there is no spatial locality or order information among nodes ( Gao et al. , 2018 ) . Some works try to overcome this limitation with two kinds of methods ; those are node clustering ( Ying et al. , 2018 ) and primary nodes sampling ( Gao & Ji , 2019 ; Zhang et al. , 2018 ) . The node clustering methods create graphs with super-nodes by learning a nodes assignment matrix . These methods suffer from the over-fitting problem and need auxiliary link prediction tasks to stabilize the training ( Ying et al. , 2018 ) . The primary nodes sampling methods like top-k pooling ( Gao & Ji , 2019 ) rank the nodes in a graph and sample top-k nodes to form the sampled graph . It uses a small number of additional trainable parameters and is shown to be more powerful ( Gao & Ji , 2019 ) . However , the top-k pooling layer does not explicitly incorporate the topology information in a graph when computing ranking scores , which may cause performance loss . In this work , we propose a novel topology-aware pooling ( TAP ) layer that explicitly encodes the topology information when computing ranking scores . We use an attention operator to compute similarity scores between each node and its neighboring nodes . The average similarity score of a node is used as its ranking score in the selection process . To avoid isolated nodes problem in our TAP layer , we further propose a graph connectivity term for computing the ranking scores of nodes . The graph connectivity term uses degree information as a bias term to encourage the layer to select highly connected nodes to form the sampled graph . Based on the TAP layer , we develop topology-aware pooling networks for network embedding learning . Experimental results on graph classification tasks demonstrate that our proposed networks with TAP layer consistently outperform previous models . The comparison results between our TAP layer and other pooling layers based on the same network architecture demonstrate the effectiveness of our method compared to other pooling methods . 2 BACKGROUND AND RELATED WORK . In this section , we describe graph pooling operations and attention operators . 2.1 GRAPH POOLING OPERATIONS . The pooling operations on graph data mainly include two categories ; those are node clustering and node sampling . DIFFPOOL ( Ying et al. , 2018 ) realizes graph pooling operation by clustering nodes into super-nodes . By learning an assignment matrix , DIFFPOOL softly assigns each node to different clusters in the new graph with specified probabilities . The pooling operations under this category retain and encode all nodes information into the new graph . One challenge of methods in this category is that they may increase the risk of over-fitting by training another network to learn the assignment matrix . In addition , the new graph is mostly connected where each edge value represents the strength of connectivity between two nodes . The connectivity pattern in the new graph may greatly differ from that of the original graph . The node sampling methods mainly select a fixed number k of the most important nodes to form a new graph . In SortPool ( Zhang et al. , 2018 ) , the same feature of each node is used for ranking and k nodes with the largest values in this feature are selected to form the coarsened graph . Top-k pooling ( Gao & Ji , 2019 ) generates the ranking scores by using a trainable projection vector that projects feature vectors of nodes into scalar values . k nodes with the largest scalar values are selected to form the coarsened graph . These methods involve none or a very small number of extra trainable parameters , thereby avoiding the risk of over-fitting . However , these methods suffer from one limitation that they do not explicitly consider the topology information during pooling . Both SortPool and top-k pooling select nodes based on scalar values that do not explicitly incorporate topology information . In this work , we propose a pooling operation that explicitly encodes topology information in ranking scores , thereby leading to an improved operation . 2.2 ATTENTION OPERATORS . Attention operator has shown to be effective in challenging tasks in various fields such as computer vision ( Xu et al. , 2015b ; Lu et al. , 2016 ; Li et al. , 2018 ) and natural language processing ( Malinowski et al. , 2018 ; Bahdanau et al. , 2015 ; Vaswani et al. , 2017 ) . Attention operator is capable of capturing long-range relationships , thereby leading to better performances ( Wang et al. , 2018 ) . The inputs to an attention operator consist of three matrices ; those are a query matrix Q ∈ Rd×m , a key matrix K ∈ Rd×n , and a value matrix V ∈ Rp×n . The attention operator computes the response of each query vector in Q by attending it to all key vectors in K. It uses the resulting coefficient vector to take a weighted sum over value vectors in V . The layer-wise operation of an attention operation is defined as O = V softmax ( KTQ ) . When the attention operator is applied to graph , each node only attends to its neighboring nodes ( Veličković et al. , 2017 ) . Self-attention can also produce an attention mask to control information flow on selected nodes in pooling operation ( Lee et al. , 2019 ) . In our proposed pooling operation , we employ an attention operator to compute ranking scores that explicitly encode topology information . 3 TOPOLOGY-AWARE POOLING LAYERS AND NETWORKS . In this work , we propose the topology-aware pooling ( TAP ) layer that uses attention operators to encode topology information in ranking scores for node selection . We also propose a graph connectivity term in the computation of ranking scores , which encourages better graph connectivity in the coarsened graph . Based on our TAP layer , we propose the topology-aware pooling networks for network representation learning . 3.1 TOPOLOGY-AWARE POOLING LAYER . Pooling layers have shown to be important on grid-like data with regard to reducing feature map sizes and enlarging receptive fields ( Yu & Koltun , 2016 ; Carreira et al. , 2012 ) . On graph data , two kinds of pooling layers have been proposed ; those are node clustering ( Ying et al. , 2018 ) and primary nodes sampling ( Gao & Ji , 2019 ; Zhang et al. , 2018 ) . A primary nodes sampling method , known as top-k pooling ( Gao & Ji , 2019 ) , uses a projection vector to generate ranking scores for each node in the graph . The graph is created by choosing nodes with k-largest scores . However , the sampling process relies on the projection vector and does not explicitly consider the topology information in the graph . In this section , we propose the topology-aware pooling ( TAP ) layer that performs primary nodes sampling by considering the graph topology . In this layer , we generate the ranking scores based on local information . To this end , we employ an attention operator to compute the similarity scores between each node and its neighboring nodes . The ranking score for a node i is the mean value of the similarity scores with its neighboring nodes . The resulting ranking score for a node indicates the similarity between this node and its neighboring nodes . If a node has a high ranking score , it can highly represent a local graph that consists of it and its neighboring nodes . By choosing nodes with the highest ranking scores , we can retain the maximum information in the sampled graph . Suppose there are N nodes in a graph G , each of which contains C features . In layer ` , we use two matrices to represent the graph ; those are the adjacency matrix A ( ` ) ∈ RN×N and the feature matrix X ( ` ) ∈ RN×C . The non-zero entries in A ( ` ) represent edges in the graph . The ith row in X ( ` ) denotes the feature vector of node i . The layer-wise forward propagation rule of TAP in layer ` is defined as K = X ( ` ) W ( ` ) , ∈RN×C ( 1 ) E = X ( ` ) KT , ∈RN×N ( 2 ) Ẽ = E ◦A ( ` ) , ∈RN×N ( 3 ) d = N∑ j=1 A ( ` ) : j , ∈R N ( 4 ) s = sigmoid ( ∑N j=1 Ẽ : j d ) , ∈RN ( 5 ) idx = Rankingk ( s ) , ∈Rk ( 6 ) A ( ` +1 ) = A ( ` ) ( idx , idx ) , ∈Rk×k ( 7 ) X ( ` +1 ) = X ( ` ) ( idx , : ) diag ( s ( idx ) ) , ∈Rk×C ( 8 ) where W ( ` ) ∈ RC×C is a trainable weight matrix , A ( ` ) : j is the jth column of matrix A ( ` ) , ◦ denotes the element-wise matrix multiplication , Ẽ : j is the jth column of matrix Ẽ , k is the number of nodes selected in the sampled graph , and diag ( · ) constructs a diagonal matrix using input vector as diagonal elements . Rankingk operator ranks the scores and return the indices of k-largest values in s. To compute attention scores , we perform a linear transformation on feature matrix X ( ` ) in Eq . ( 1 ) , which results in the key matrix K. We use the input feature matrix as the query matrix . The similarity score matrix E is obtained by the matrix multiplication between X ( ` ) and K in Eq . ( 2 ) . Each value eij in E measures the similarity between node i and node j . Since E contains similarity scores for nodes that are not directly connected , we use the adjacency matrix A ( ` ) as a mask to set these entries in E to zeros in Eq . ( 3 ) , resulting in Ẽ . We compute the degree of each node in Eq . ( 4 ) . The ranking score of a node is computed in Eq . ( 5 ) by taking the average of similarity scores between this node and its neighboring nodes followed by a sigmoid operation . Here , we perform element-wise division between two vectors . The resulting score vector is s = [ s1 , s2 , . . . , sN ] T where si represents the ranking score of node i. Rankingk is an operator that selects k-largest values and returns the corresponding indices . In Eq . ( 6 ) , we use Rankingk to select the k-most important nodes with indices in idx . Using idx , we extract new adjacency matrix A ( ` +1 ) in Eq . ( 7 ) and new feature matrix X ( ` +1 ) in Eq . ( 8 ) . Here , we use the ranking scores s ( idx ) as gates to control information flow and enable the gradient back-propagation for trainable transformation matrix W ( ` ) ( Gao & Ji , 2019 ) . This method can be considered as a local-voting , global-ranking process . In our TAP layer , the ranking scores are derived from the similarity scores of each node with its neighboring nodes , thereby encoding the topology information of each node in its ranking score . This can be considered as a local voting process that each node gets its votes from the local neighborhood . When performing global ranking , the nodes that get the highest votes from local neighborhoods are selected such that maximum information in the graph can be retained . Figure 1 provides an illustration of our proposed TAP layer . Compared to the top-k pooling ( Gao & Ji , 2019 ) , our TAP layer considers topology information in the graph , thereby leading to a better coarsened graph .

Summary: This paper presented a new pooling method for learning graph-level embeddings. The key idea is to use the initial node attributes to compute all-pair attention scores for each node pair and then use these attention scores to formulate a new graph adjacency matrix beyond the original raw graph adjacency matrix. As the result, each node can average these attention score edges to compute the overall importance. Based on these scores, the method chooses top-k nodes to perform graph coarsening operation. In addition, a graph connectivity term is proposed to address the problems of isolated nodes. Experiments are performed to validate the effectiveness of the proposed method.