Mincut Pooling in Graph Neural Networks

1 INTRODUCTION . A fundamental component in deep convolutional neural networks is the pooling operation , which replaces the output of convolutions with local summaries of nearby points and is usually implemented by maximum or average operations ( Lee et al. , 2016 ) . State-of-the-art architectures alternate convolutions , which extrapolate local patterns irrespective of the specific location on the input signal , and pooling , which lets the ensuing convolutions capture aggregated patterns . Pooling allows to learn abstract representations in deeper layers of the network by discarding information that is superfluous for the task , and keeps model complexity under control by limiting the growth of intermediate features . Graph Neural Networks ( GNNs ) extend the convolution operation from regular domains , such as images or time series , to data with arbitrary topologies and unordered structures described by graphs ( Battaglia et al. , 2018 ) . The development of pooling strategies for GNNs , however , has lagged behind the design of newer and more effective message-passing ( MP ) operations ( Gilmer et al. , 2017 ) , such as graph convolutions , mainly due to the difficulty of defining an aggregated version of the original graph that supports the pooled signal . A naı̈ve pooling strategy in GNNs is to average all nodes features ( Li et al. , 2016 ) , but it has limited flexibility since it does not extract local summaries of the graph structure , and no further MP operations can be applied afterwards . An alternative approach consists in pre-computing coarsened versions of the original graph and then fit the data to these deterministic structures ( Bruna et al. , 2013 ) . While this aggregation accounts for the connectivity of the graph , it ignores task-specific objectives as well as the node features . In this paper , we propose a differentiable pooling operation implemented as a neural network layer , which can be seamlessly combined with other MP layers ( see Fig . 1 ) . The parameters in the pooling layer are learned by combining the task-specific loss with an unsupervised regularization term , which optimizes a continuous relaxation of the normalized minCUT objective . The minCUT identifies dense graph components , where the nodes features become locally homogeneous after the message-passing . By gradually aggregating these components , the GNN learns to distil global properties from the graph . The proposed minCUT pooling operator ( minCUTpool ) yields partitions that 1 ) cluster together nodes which have similar features and are strongly connected on the graph , and 2 ) take into account the objective of the downstream task . Message-passing minCUT Pooling Message-passing Figure 1 : A deep GNN architecture where message-passing is followed by minCUT pooling . 2 BACKGROUND . 2.1 MINCUT AND SPECTRAL CLUSTERING . Given a graph G = { V , E } , |V| = N , and the associated adjacency matrix A ∈ RN×N , the K-way normalized minCUT ( simply referred to as minCUT ) aims at partitioning V in K disjoint subsets by removing the minimum volume of edges . The problem is equivalent to maximizing 1 K K∑ k=1 links ( Vk ) degree ( Vk ) = 1 K K∑ k=1 ∑ i , j∈Vk Ei , j∑ i∈Vk , j∈V\Vk Ei , j , ( 1 ) where the numerator counts the edge volume within each cluster , and the denominator counts the edges between the nodes in a cluster and the rest of the graph ( Shi & Malik , 2000 ) . Let C ∈ RN×K be a cluster assignment matrix , so that Ci , j = 1 if node i belongs to cluster j , and 0 otherwise . The minCUT problem can be expressed as maximize 1 K K∑ k=1 CTk ACk CTk DCk , s.t . C ∈ { 0 , 1 } N×K , C1K = 1N , ( 2 ) where D = diag ( A1N ) is the degree matrix ( Dhillon et al. , 2004 ) . Since problem ( 2 ) is NP-hard , it is usually recast in a relaxed formulation that can be solved in polynomial time and guarantees a near-optimal solution ( Yu & Shi , 2003 ) : arg max Q∈RN×K 1 K K∑ k=1 QTk AQk , s.t . Q = C ( C TDC ) − 1 2 , QTQ = IK . ( 3 ) While the optimization problem ( 3 ) is still non-convex , there exists an optimal solution Q∗ = UKO , where UK ∈ RN×K contains the eigenvectors of A corresponding to the K largest eigenvalues , and O ∈ RK×K is an orthogonal transformation ( Ikebe et al. , 1987 ) . Since the elements of Q∗ are real values rather than binary cluster indicators , the spectral clustering ( SC ) approach can be used to find discrete cluster assignments . In SC , the rows of Q∗ are treated as node representations embedded in the eigenspace of the Laplacian , and are clustered together with standard algorithms such as k-means ( Von Luxburg , 2007 ) . One of the main limitations of SC lies in the computation of the spectrum of A , which has a memory complexity of O ( N2 ) and a computational complexity of O ( N3 ) . This prevents its applicability to large datasets . To deal with such scalability issues , the constrained optimization in ( 3 ) can be solved by gradient descent algorithms that refine the solution by iterating operations whose individual complexity is O ( N2 ) , or even O ( N ) ( Han & Filippone , 2017 ) . Those algorithms search the solution on the manifold induced by the orthogonality constraint on the columns of Q , by performing gradient updates along the geodesics ( Wen & Yin , 2013 ; Collins et al. , 2014 ) . Alternative approaches rely on the QR factorization to constrain the space of feasible solutions ( Damle et al. , 2016 ) , and alleviate the cost O ( N3 ) of the factorization by ensuring that orthogonality holds only on one minibatch at a time ( Shaham et al. , 2018 ) . Other works based on neural networks include an autoencoder trained to map the ith row of the Laplacian to the ith components of the first K eigenvectors , to avoid the spectral decomposition ( Tian et al. , 2014 ) . Yi et al . ( 2017 ) use a soft orthogonality constraint to learn spectral embeddings as a volumetric reparametrization of a precomputed Laplacian eigenbase . Shaham et al . ( 2018 ) ; Kampffmeyer et al . ( 2019 ) propose differentiable loss functions to partition generic data and process out-of-sample data at inference time . Nazi et al . ( 2019 ) generate balanced node partitions with a GNN , but adopt an optimization that does not encourage cluster assignments to be orthogonal . 2.2 GRAPH NEURAL NETWORKS . Many approaches have been proposed to process graphs with neural networks , including recurrent architectures ( Scarselli et al. , 2009 ; Li et al. , 2016 ) or convolutional operations inspired by filters used in graph signal processing ( Defferrard et al. , 2016 ; Bianchi et al. , 2019 ) . Since our focus is on graph pooling , we base our GNN implementation on a simple MP operation , which combines the features of each node with its 1st-order neighbors . To account for the initial node features , it is possible to introduce self-loops by adding a ( scaled ) identity matrix to the diagonal of A ( Kipf & Welling , 2017 ) . Since our pooling will modify the structure of the adjacency matrix , we prefer a MP implementation that leaves the original A unaltered and accounts for the initial node features by means of skip connections . Let Ã = D− 1 2 AD− 1 2 ∈ RN×N be the symmetrically normalized adjacency matrix and X ∈ RN×F the matrix containing the node features . The output of the MP layer is X ( t+1 ) = MP ( X ( t ) , Ã ) = ReLU ( ÃX ( t ) Wm + X ( t ) Ws ) , ( 4 ) where ΘMP = { Wm , Ws } are the trainable weights relative to the mixing and skip component of the layer , respectively . 3 PROPOSED METHOD . The minCUT pooling strategy computes a cluster assignment matrix S ∈ RN×K by means of a multi-layer perceptron , which maps each node feature xi into the ith row of S : S = softmax ( ReLU ( XW1 ) W2 ) , ( 5 ) where ΘPool = { W1 ∈ RF×H , W2 ∈ RH×K } are trainable parameters . The softmax function guarantees that si , j ∈ [ 0 , 1 ] and enforces the constraints S1K = 1N inherited from the optimization problem in ( 2 ) . The parameters ΘMP and ΘPool are jointly optimized by minimizing the usual task-specific loss , as well as an unsupervised loss Lu , which is composed of two terms Lu = Lc + Lo = − Tr ( ST ÃS ) Tr ( ST D̃S ) ︸ ︷︷ ︸ Lc + ∥∥∥∥ STS‖STS‖F − IK√K ∥∥∥∥ F︸ ︷︷ ︸ Lo , ( 6 ) where ‖ · ‖F indicates the Frobenius norm . The cut loss term , Lc , evaluates the minCUT given by the cluster assignment S , and is bounded by −1 ≤ Lc ≤ 0 . Minimizing Lc encourages strongly connected nodes to be clustered together , since the inner product 〈si , sj〉 increases when ãi , j is large . Lc has a single maximum , reached when the numerator Tr ( ST ÃS ) = 1K ∑K k=1 S T k ÃSk = 0 . This occurs if , for each pair of connected nodes ( i.e. , ãi , j > 0 ) , the cluster assignments are orthogonal ( i.e. , 〈si , sj〉 = 0 ) . Lc reaches its minimum , −1 , when Tr ( ST ÃS ) = Tr ( ST D̃S ) . This occurs when in a graph with K disconnected components the cluster assignments are equal for all the nodes in the same component and orthogonal to the cluster assignments of nodes in different components . However , Lc is a non-convex function and its minimization can lead to local minima or degenerate solutions . For example , given a connected graph , a trivial optimal solution is the one that assigns all nodes to the same cluster . As a consequence of the continuous relaxation , another degenerate minimum occurs when the cluster assignments are all uniform , that is , all nodes are equally assigned to all clusters . This problem is exacerbated by prior message-passing operations , which make the node features more uniform . The orthogonality loss term , Lo , penalizes the degenerate minima of Lc by encouraging the cluster assignments to be orthogonal and the clusters to be of similar size . Since the two matrices in Lo have unitary norm it is easy to see that 0 ≤ Lo ≤ 2 . Therefore , Lo does not dominate over Lc and the two terms can be safely summed directly ( see Fig . 4 for an example ) . IK can be interpreted as a ( rescaled ) clustering matrix IK = ŜT Ŝ , where Ŝ assigns exactly N/K points to each cluster . The value of the Frobenius norm between clustering matrices is not dominated by the performance on the largest clusters ( Law et al. , 2017 ) and , thus , can be used to optimize intra-cluster variance . Contrarily to SC methods that search for feasible solutions only within the space of orthogonal matrices , Lo only introduces a soft constraint that could be violated during the learning procedure . Since Lc is non-convex , the violation compromises the theoretical guarantee of convergence to the optimum of ( 3 ) . However , we note that : 1. the cluster assignments S are well initialized : after the MP operation , the features of the connected vertices become similar and , since the MLP is a smooth function ( Nelles , 2013 ) , it yields similar cluster assignments for those vertices ; 2. in the GNN architecture , the minCUT objective is a regularization term and , therefore , a solution which is sub-optimal for ( 3 ) could instead be adequate for the specific objective of the downstream task ; 3. optimizing the task-specific loss helps the GNN to avoid the degenerate minima of Lc .

Summary: This paper proposes a solution to the important problem of pooling in graph neural networks. The method relies on minimizing a surrogate function inside the standard SGD loop and in conjunction with the optimization of the model parameters - such loss function aiming at optimizing the minCut on the graph. By that it aims to effective achieve a soft clustering of nodes that are both well connected and that have similar embeddings. This in an elegant choice, somewhat resembling the DiffPool method since it's also end-to-end trainable. However it adds the local graph connectivity information due to the minCut loss (and related orthogonality penalty to achieve non trivial solutions on the relaxed minCut continuous problem). Such local graph connectivity is indeed important information to consider when carrying out pooling.