Graph Information Bottleneck for Subgraph Recognition

1 INTRODUCTION . Classifying the underlying labels or properties of graphs is a fundamental problem in deep graph learning with applications across many fields , such as biochemistry and social network analysis . However , real world graphs are likely to contain redundant even noisy information ( Franceschi et al. , 2019 ; Yu et al. , 2019 ) , which poses a huge negative impact for graph classification . This triggers an interesting problem of recognizing an informative yet compressed subgraph from the original graph . For example , in drug discovery , when viewing molecules as graphs with atoms as nodes and chemical bonds as edges , biochemists are interested in identifying the subgraphs that mostly represent certain properties of the molecules , namely the functional groups ( Jin et al. , 2020b ; Gilmer et al. , 2017 ) . In graph representation learning , the predictive subgraph highlights the vital substructure for graph classification , and provides an alternative way for yielding graph representation besides mean/sum aggregation ( Kipf & Welling , 2017 ; Velickovic et al. , 2017 ; Xu et al. , 2019 ) and pooling aggregation ( Ying et al. , 2018 ; Lee et al. , 2019 ; Bianchi et al. , 2020 ) . In graph attack and defense , it is vital to purify a perturbed graph and mine the robust structures for classification ( Jin et al. , 2020a ) . Recently , the mechanism of self-attentive aggregation ( Li et al. , 2019 ) somehow discovers a vital substructure at node level with a well-selected threshold . However , this method only identifies isolated important nodes but ignores the topological information at subgraph level . Consequently , it ∗This work was done when Junchi Yu was a research intern at Tencent AI LAB . †Corresponding Author leads to a novel challenge as subgraph recognition : How can we recognize a compressed subgraph with minimum information loss in terms of predicting the graph labels/properties ? Recalling the above challenge , there is a similar problem setting in information theory called information bottleneck ( IB ) principle ( Tishby et al. , 1999 ) , which aims to juice out a compressed data from the original data that keeps most predictive information of labels or properties . Enhanced with deep learning , IB can learn informative representation from regular data in the fields of computer vision ( Peng et al. , 2019 ; Alemi et al. , 2017 ; Luo et al. , 2019 ) , reinforcement learning ( Goyal et al. , 2019 ; Igl et al. , 2019 ) and natural language precessing ( Wang et al. , 2020 ) . However , current IB methods , like VIB ( Alemi et al. , 2017 ) , is still incapable for irregular graph data . It is still challenging for IB to compress irregular graph data , like a subgraph from an original graph , with a minimum information loss . Hence , we advance the IB principle for irregular graph data to resolve the proposed subgraph recognition problem , which leads to a novel principle , Graph Information Bottleneck ( GIB ) . Different from prior researches in IB that aims to learn an optimal representation of the input data in the hidden space , GIB directly reveals the vital substructure in the subgraph level . We first i ) leverage the mutual information estimator from Deep Variational Information Bottleneck ( VIB ) ( Alemi et al. , 2017 ) for irregular graph data as the GIB objective . However , VIB is intractable to compute the mutual information without knowing the distribution forms , especially on graph data . To tackle this issue , ii ) we adopt a bi-level optimization scheme to maximize the GIB objective . Meanwhile , the continuous relaxation that we adopt to approach the discrete selection of subgraph will lead to unstable optimization process . To further stabilize the training process and encourage a compact subgraph , iii ) we propose a novel connectivity loss to assist GIB to effectively discover the maximally informative but compressed subgraph , which is defined as IB-subgraph . By optimizing the above GIB objective and connectivity loss , one can recognize the IB-subgraph without any explicit subgraph annotation . On the other hand , iv ) GIB is model-agnostic and can be easily plugged into various Graph Neural Networks ( GNNs ) . We evaluate the properties of the IB-subgraph in three application scenarios : improvement of graph classification , graph interpretation , and graph denoising . Extensive experiments on both synthetic and real world datasets demonstrate that the information-theoretic IB-subgraph enjoys superior graph properties compared to the subgraphs found by SOTA baselines . 2 RELATED WORK . Graph Classification . In recent literature , there is a surge of interest in adopting graph neural networks ( GNN ) in graph classification . The core idea is to aggregate all the node information for graph representation . A typical implementation is the mean/sum aggregation ( Kipf & Welling , 2017 ; Xu et al. , 2019 ) , which is to average or sum up the node embeddings . An alternative way is to leverage the hierarchical structure of graphs , which leads to the pooling aggregation ( Ying et al. , 2018 ; Zhang et al. , 2018 ; Lee et al. , 2019 ; Bianchi et al. , 2020 ) . When tackling with the redundant and noisy graphs , these approaches will likely to result in sub-optimal graph representation . Recently , InfoGraph ( Sun et al. , 2019 ) maximize the mutual information between graph representations and multi-level local representations to obtain more informative global representations . Information Bottleneck . Information bottleneck ( IB ) , originally proposed for signal processing , attempts to find a short code of the input signal but preserve maximum information of the code ( Tishby et al. , 1999 ) . ( Alemi et al. , 2017 ) firstly bridges the gap between IB and the deep learning , and proposed variational information bottleneck ( VIB ) . Nowadays , IB and VIB have been wildly employed in computer vision ( Peng et al. , 2019 ; Luo et al. , 2019 ) , reinforcement learning ( Goyal et al. , 2019 ; Igl et al. , 2019 ) , natural language processing ( Wang et al. , 2020 ) and speech and acoustics ( Qian et al. , 2020 ) due to the capability of learning compact and meaningful representations . However , IB is less researched on irregular graphs due to the intractability of mutual information . Subgraph Discovery . Traditional subgraph discovery includes dense subgraph discovery and frequent subgraph mining . Dense subgraph discovery aims to find the subgraph with the highest density ( e.g . the number of edges over the number of nodes ( Fang et al. , 2019 ; Gionis & Tsourakakis , 2015 ) ) . Frequent subgraph mining is to look for the most common substructure among graphs ( Yan & Yan , 2002 ; Ketkar et al. , 2005 ; Zaki , 2005 ) . At node-level , researchers discover the vital substructure via the attention mechanism ( Velickovic et al. , 2017 ; Lee et al. , 2019 ; Knyazev et al. , 2019 ) . Ying et al . ( 2019 ) further identifies the important computational graph for node classification . Alsentzer et al . ( 2020 ) discovers subgraph representations with specific topology given subgraph-level annotation . Recently , it is popular to select a neighborhood subgraph of a central node to do message passing in node representation learning . DropEdge ( Rong et al. , 2020 ) relieves the over-smoothing phenomenon in deep GCNs by randomly dropping a portion of edges in graph data . Similar to DropEdge , DropNode ( Chen et al. , 2018 ; Hamilton et al. , 2017 ; Huang et al. , 2018 ) principle is also widely adopted in node representation learning . FastGCN ( Chen et al. , 2018 ) and ASGCN ( Huang et al. , 2018 ) accelerate GCN training via node sampling . GraphSAGE ( Hamilton et al. , 2017 ) leverages neighborhood sampling for inductive node representation learning . NeuralSparse ( Zheng et al. , 2020 ) select Top-K ( K is a hyper-parameter ) task-relevant 1-hop neighbors of a central node for robust node classification . Similarly , researchers discover the vital substructure at node level via the attention mechanism ( Velickovic et al. , 2017 ; Lee et al. , 2019 ; Knyazev et al. , 2019 ) . 3 NOTATIONS AND PRELIMINARIES . Let { ( G1 , Y1 ) , . . . , ( GN , YN ) } be a set of N graphs with their real value properties or categories , where Gn refers to the n-th graph and Yn refers to the corresponding properties or labels . We denote by Gn = ( V , E , A , X ) the n-th graph of size Mn with node set V = { Vi|i = 1 , . . . , Mn } , edge set E = { ( Vi , Vj ) |i > j ; Vi , Vj is connected } , adjacent matrix A ∈ { 0 , 1 } Mn×Mn , and feature matrix X ∈ RMn×d of V with d dimensions , respectively . Denote the neighborhood of Vi as N ( Vi ) = { Vj | ( Vi , Vj ) ∈ E } . We use Gsub as a specific subgraph and Gsub as the complementary structure of Gsub in G. Let f : G → R/ [ 0 , 1 , · · · , n ] be the mapping from graphs to the real value property or category , Y , G is the domain of the input graphs . I ( X , Y ) refers to the Shannon mutual information of two random variables . 3.1 GRAPH CONVOLUTIONAL NETWORK . Graph convolutional network ( GCN ) is widely adopted to graph classification . Given a graph G = ( V , E ) with node feature X and adjacent matrix A , GCN outputs the node embeddings X ′ from the following process : X ′ = GCN ( A , X ; W ) = ReLU ( D− 1 2 ÂD− 1 2XW ) , ( 1 ) where D refers to the diagonal matrix with nodes ’ degrees and W refers to the model parameters . One can simply sum up the node embeddings to get a fixed length graph embeddings ( Xu et al. , 2019 ) . Recently , researchers attempt to exploit hierarchical structure of graphs , which leads to various graph pooling methods ( Li et al. , 2019 ; Gao & Ji , 2019 ; Lee et al. , 2019 ; Diehl , 2019 ; Zhang et al. , 2018 ; Ranjan et al. , 2020 ; Ying et al. , 2018 ) . Li et al . ( 2019 ) enhances the graph pooling with self-attention mechanism to leverage the importance of different nodes contributing to the results . Finally , the graph embedding is obtained by multiplying the node embeddings with the normalized attention scores : E = Att ( X ′ ) = softmax ( Φ2tanh ( Φ1X ′T ) ) X ′ , ( 2 ) where Φ1 and Φ2 refers to the model parameters of self-attention . 3.2 OPTIMIZING INFORMATION BOTTLENECK OBJECTIVE . Given the input signal X and the label Y , the objective of IB is maximized to find the the internal code Z : maxZ I ( Z , Y ) − βI ( X , Z ) , where β refers to a hyper-parameter trading off informativeness and compression . Optimizing this objective will lead to a compact but informative Z. Alemi et al . ( 2017 ) optimize a tractable lower bound of the IB objective : LV IB = 1 N ∑N i=1 ∫ p ( z|xi ) log qφ ( yi|z ) dz − βKL ( p ( z|xi ) |r ( z ) ) , ( 3 ) where qφ ( yi|z ) is the variational approximation to pφ ( yi|z ) and r ( z ) is the prior distribution of Z . However , it is hard to estimate the mutual information in high dimensional space when the distribution forms are inaccessible , especially for irregular graph data .

Summary: I think this is a nice paper that successfully used information theoretic objective functions for graph representation learning. The authors leveraged the DONSKER approximation of mutual information for a global information bottleneck loss used on the input-space instead of learned latent-space. To help stabilise optimisation, the authors also use bi-level optimisation with more iterations on the inner $I(G, G_{sub})$ as well as automatic masks learning through their $L_{con}$ loss. The authors also showed many experiments on graph classification, denoising, and interpretation tasks.