StructPool: Structured Graph Pooling via Conditional Random Fields

1 INTRODUCTION . Graph neural networks have achieved the state-of-the-art results for multiple graph tasks , such as node classification ( Veličković et al. , 2018 ; Gao & Ji , 2019b ; Gao et al. , 2018 ) and link prediction ( Zhang & Chen , 2018 ; Cai & Ji , 2020 ) . These results demonstrate the effectiveness of graph neural networks to learn node representations . However , graph classification tasks also require learning good graph-level representations . Since pooling operations are shown to be effective in many image and NLP tasks , it is natural to investigate pooling techniques for graph data ( Yu & Koltun , 2016 ; Springenberg et al. , 2014 ) . Recent work extends the global sum/average pooling operations to graph models by simply summing or averaging all node features ( Atwood & Towsley , 2016 ; Simonovsky & Komodakis , 2017 ) . However , these trivial global pooling operations may lose important features and ignore structural information . Furthermore , global pooling are not hierarchical so that we can not apply them where multiple pooling operations are required , such as Graph UNet ( Gao & Ji , 2019a ) . Several advanced graph pooling methods , such as SORTPOOL ( Zhang et al. , 2018 ) , TOPKPOOL ( Gao & Ji , 2019a ) , DIFFPOOL ( Ying et al. , 2018 ) , and SAGPOOL ( Lee et al. , 2019 ) , are recently proposed and achieve promising performance on graph classification tasks . However , none of them explicitly models the relationships among different nodes and thus may ignore important structural information . We argue that such information is important and should be explicitly captured in graph pooling . In this work , we propose a novel graph pooling technique , known as the STRUCTPOOL , that formulates graph pooling as a structured prediction problem . Following DIFFPOOL ( Ying et al. , 2018 ) , we consider graph pooling as a node clustering problem , and each cluster corresponds to a node in the new graph after pooling . Intuitively , two nodes with similar features should have a higher probability of being assigned to the same cluster . Hence , the assignment of a given node should depend on both the input node features and the assignments of other nodes . We formulate this as a structured prediction problem and employ conditional random fields ( CRFs ) ( Lafferty et al. , 2001 ) to capture such high-order structural relationships among the assignments of different nodes . In addition , we generalize our method by incorporating the graph topological information so that our method can control the clique set in our CRFs . We employ the mean field approximation to compute the assignments and describe how to incorporate it in graph networks . Then the networks can be trained in an end-to-end fashion . Experiments show that our proposed STRUCTPOOL outperforms existing methods significantly and consistently . We also show that STRUCTPOOL incurs acceptable computational cost given its superior performance . 2 BACKGROUND AND RELATED WORK . 2.1 GRAPH CONVOLUTIONAL NETWORKS . A graph can be represented by its adjacency matrix and node features . Formally , for a graph G consisting of n nodes , its topology information can be represented by an adjacency matrix A ∈ { 0 , 1 } n×n , and the node features can be represented as X ∈ Rn×c assuming each node has a c-dimensional feature vector . Deep graph neural networks ( GNNs ) learn feature representations for different nodes using these matrices ( Gilmer et al. , 2017 ) . Several approaches are proposed to investigate deep GNNs , and they generally follow a neighborhood information aggregation scheme ( Gilmer et al. , 2017 ; Xu et al. , 2019 ; Hamilton et al. , 2017 ; Kipf & Welling , 2017 ; Veličković et al. , 2018 ) . In each step , the representation of a node is updated by aggregating the representations of its neighbors . Graph Convolutional Networks ( GCNs ) are popular variants of GNNs and inspired by the first order graph Laplacian methods ( Kipf & Welling , 2017 ) . The graph convolution operation is formally defined as : Xi+1 = f ( D − 12 ÂD− 1 2XiPi ) , ( 1 ) where Â = A + I is used to add self-loops to the adjacency matrix , D denotes the diagonal node degree matrix to normalize Â , Xi ∈ Rn×ci are the node features after ith graph convolution layer , Pi ∈ Rci×ci+1 is a trainable matrix to perform feature transformation , and f ( · ) denotes a non-linear activation function . Then Xi ∈ Rn×ci is transformed to Xi+1 ∈ Rn×ci+1 where the number of nodes remains the same . A similar form of GCNs proposed in ( Zhang et al. , 2018 ) can be expressed as : Xi+1 = f ( D −1ÂXiPi ) . ( 2 ) It differs from the GCNs in Equation ( 1 ) by performing different normalization and is a theoretically closer approximation to the Weisfeiler-Lehman algorithm ( Weisfeiler & Lehman , 1968 ) . Hence , in our models , we use the latter version of GCNs in Equation ( 2 ) . 2.2 GRAPH POOLING . Several advanced pooling techniques are proposed recently for graph models , such as SORTPOOL , TOPKPOOL , DIFFPOOL , and SAGPOOL , and achieve great performance on multiple benchmark datasets . All of SORTPOOL ( Zhang et al. , 2018 ) , TOPKPOOL ( Gao & Ji , 2019a ) , and SAGPOOL ( Lee et al. , 2019 ) learn to select important nodes from the original graph and use these nodes to build a new graph . They share the similar idea to learn a sorting vector based on node representations using GCNs , which indicates the importance of different nodes . Then only the top k important nodes are selected to form a new graph while the other nodes are ignored . However , the ignored nodes may contain important features and this information is lost during pooling . DIFFPOOL ( Ying et al. , 2018 ) treats the graph pooling as a node clustering problem . A cluster of nodes from the original graph are merged to form a new node in the new graph . DIFFPOOL proposes to perform GCNs on node features to obtain node clustering assignment matrix . Intuitively , the cluster assignment of a given node should depend on the cluster assignments of other nodes . However , DIFFPOOL does not explicitly consider such high-order structural relationships , which we believe are important for graph pooling . In this work , we propose a novel structured graph pooling technique , known as the STRUCTPOOL , for effectively learning high-level graph representations . Different from existing methods , our method explicitly captures high-order structural relationships between different nodes via conditional random fields . In addition , our method is generalized by incorporating graph topological information A to control which node pairs are included in our CRFs . 2.3 INTEGRATING CRFS WITH GNNS . Recent work ( Gao et al. , 2019 ; Qu et al. , 2019 ; Ma et al. , 2019 ) investigates how to combine CRFs with GNNs . The CGNF ( Ma et al. , 2019 ) is a GNN architecture for graph node classification which explicitly models a joint probability of the entire set of node labels via CRFs and performs inference via dynamic programming . In addition , the GMNN ( Qu et al. , 2019 ) focuses on semi-supervised object classification tasks and models the joint distribution of object labels conditioned on object attributes using CRFs . It proposes a pseudolikelihood variational EM framework for model learning and inference . Recent work ( Gao et al. , 2019 ) integrates CRFs with GNNs by proposing a CRF layer to encourage similar nodes to have similar hidden features so that similarity information can be preserved explicitly . All these methods are proposed for node classification tasks and the CRFs are incorporated in different ways . Different from existing work , our STRUCTPOOL is proposed for graph pooling operation and the energy is optimized via mean field approximation . All operations in our STRUCTPOOL can be realized by GNN operations so that our STRUCTPOOL can be easily used in any GNNs and trained in an end-to-end fashion . 3 STRUCTURED GRAPH POOLING . 3.1 GRAPH POOLING VIA NODE CLUSTERING . Even though pooling techniques are shown to facilitate the training of deep models and improve their performance significantly in many image and NLP tasks ( Yu & Koltun , 2016 ; Springenberg et al. , 2014 ) , local pooling operations can not be directly applied to graph tasks . The reason is there is no spatial locality information among graph nodes . Global max/average pooling operations can be employed for graph tasks but they may lead to information loss , due to largely reducing the size of representations trivially . A graph G with n nodes can be represented by a feature matrix X ∈ Rn×c and an adjacent matrix A ∈ { 0 , 1 } n×n . Graph pooling operations aim at reducing the number of graph nodes and learning new representations . Suppose that graph pooling generates a new graph G̃ with k nodes . The representation matrices of G̃ are denoted as X̃ ∈ Rk×c̃ and Ã ∈ { 0 , 1 } k×k . The goal of graph pooling is to learn relationships between X , A and X̃ , Ã . In this work , we consider graph pooling via node clustering . In particular , the nodes of the original graph G are assigned to k different clusters . Then each cluster is transformed to a new node in the new graph G̃ . The clustering assignments can be represented as an assignment matrix M ∈ Rn×k . For hard assignments , mi , j ∈ { 0 , 1 } denotes if node i in graph G belongs to cluster j . For soft assignments , mi , j ∈ [ 0 , 1 ] denotes the probability that node i in graph G belongs to cluster j and ∑ j mi , j = 1 . Then the new graph G̃ can be computed as X̃ = MTX , Ã = g ( MTAM ) , ( 3 ) where g ( · ) is a function that g ( ãi , j ) = 1 if ãi , j > 0 and g ( ãi , j ) = 0 otherwise . 3.2 LEARNING CLUSTERING ASSIGNMENTS VIA CONDITIONAL RANDOM FIELDS . Intuitively , node features describe the properties of different nodes . Then nodes with similar features should have a higher chance to be assigned to the same cluster . That is , for any node in the original graph G , its cluster assignment should not only depend on node feature matrix X but also condition on the cluster assignments of the other nodes . We believe such high-order structural information is useful for graph pooling and should be explicitly captured while learning clustering assignments . To this end , we propose a novel structured graph pooling technique , known as STRUCTPOOL , which generates the assignment matrix by considering the feature matrix X and the relationships between the assignments of different nodes . We propose to formulate this as a conditional random field ( CRF ) problem . The CRFs model a set of random variables with a Markov Random Field ( MRF ) , conditioned on a global observation ( Lafferty et al. , 2001 ) . We formally define Y = { Y1 , · · · , Yn } as a random field where Yi ∈ { 1 , · · · , k } is a random variable . Each Yi indicates to which cluster the node i is assigned . Here the feature representation X is treated as global observation . We build a graphical model on Y , which is defined as G′ . Then the pair ( Y , X ) can be defined as a CRF , characterized by the Gibbs distribution as P ( Y |X ) = 1 Z ( X ) exp − ∑ c∈CG′ ψc ( Yc|X ) , ( 4 ) where c denotes a clique , CG′ is a set of cliques in G′ , Z ( X ) is the partition function , and ψc ( · ) is a potential function induced by c ( Krähenbühl & Koltun , 2011 ; Lafferty et al. , 2001 ) . Then the Gibbs energy function for an assignment y = { y1 , · · · , yn } for all variables can be written as E ( y|X ) = ∑ c∈CG′ ψc ( yc|X ) . ( 5 ) Finding the optimal assignment is equivalent to maximizing P ( Y |X ) , which can also be interpreted as minimizing the Gibbs energy .

Summary: The paper proposes a new graph pooling method by learning the node assignments from a CRF based structure prediction formulation. Different from existing pooling methods, the proposed model explicitly captures high-order structural relationships between nodes via a CRF, and the pairwise energy in the CRF is defined by a $l$-hop connection instead of the original neighborhoods. The paper is written clearly, and the experimental results are good. The sensitivity analysis is also complete.