InfoGraph: Unsupervised and Semi-supervised Graph-Level Representation Learning via Mutual Information Maximization

1 INTRODUCTION . Graphs have proven to be an effective way to represent very diverse types of data including social networks Newman & Girvan ( 2004 ) , biological reaction networksPavlopoulos et al . ( 2011 ) , proteinprotein interactions Krogan et al . ( 2006 ) , the quantum mechanical properties of individual molecules Xie & Grossman ( 2018 ) ; Jin et al . ( 2018 ) , and many more . Graphs provide explicit information about the coupling between individual units in a larger part along with a well defined framework for assigning properties to the nodes and the edges connecting them . There has been a significant amount of previous work done studying many aspects of graphs including link prediction Gao et al . ( 2011 ) ; Wang et al . ( 2011 ) and node prediction Blei et al . ( 2003 ) . Due to its flexibility , graph-like data structures can capture rich information which is critical in many applications . At the lowest level , much work has been done on learning node representations– low-dimensional vector embeddings of individual nodes Perozzi et al . ( 2014 ) ; Tang et al . ( 2015 ) ; Grover & Leskovec ( 2016 ) . Another field that has attracted a large amount of attention recently is learning representations of entire graphs . Such a problem is critical in a variety of applications such as predicting the properties of molecular graphs in both drug discovery and material science Chen et al . ( 2019b ; a ) . There has been some recent progress based on neural message passing algorithms Gilmer et al . ( 2017 ) ; Xie & Grossman ( 2018 ) , which learn the representations of entire graphs in a supervised way . These methods have been shown achieving state-of-the-art results on a variety of different prediction tasks Kipf et al . ( 2018 ) ; Xie & Grossman ( 2018 ) ; Gilmer et al . ( 2017 ) ; Chen et al . ( 2019a ) . However , one of the most difficult obstacles for supervised learning on graphs is that it is often very costly or even impossible to collect annotated labels . For example , in the chemical domain labels are typically produced with a costly Density Functional Theory ( DFT ) calculation . One option is to use semi-supervised methods which combine a small handful of labels with a larger , unlabeled , dataset . In real-world applications , partially labeled datasets are common , making tools that are able to efficiently utilize the present labels particularly useful . Coming up with methods that are able to learn unsupervised representations of an entire graph , as opposed to nodes , is an important step in working with unlabeled or partially labeled graphs Narayanan et al . ( 2017 ) ; Hu et al . ( 2019 ) ; Nguyen et al . ( 2017 ) . For example , there exists work that explores pre-training techniques for graphs to improve generalization Hu et al . ( 2019 ) . Another common approach to unsupervised representation learning on graphs is through graph kernels Pržulj ( 2007 ) ; Kashima et al . ( 2003 ) ; Orsini et al . ( 2015 ) . However , many of these methods do not provide explicit graph embeddings which many machine learning algorithms operate on . Furthermore , the handcrafted features of graph kernels lead to high dimensional , sparse or non-smooth representations and thus result in poor generalization performance , especially on large datasets Narayanan et al . ( 2017 ) . Unsupervised learning of latent representations is also an important problem in other domains , such as image generation Kingma & Welling ( 2013 ) ; Kim & Mnih ( 2018 ) and natural language processing Mikolov et al . ( 2013a ) . A recent work introduced Deep Infomax , a method that maximizes the mutual information content between the input data and the learned representation Hjelm et al . ( 2018 ) . This method outperforms other methods on many unsupervised learning tasks . Motivated by Deep InfoMax Hjelm et al . ( 2018 ) , we aim to use mutual information maximization for unsupervised representation learning on the entire graph . Specifically , our objective is to maximize the mutual information between the representations of entire graphs and the representations of substructures of different granularity . We name our model InfoGraph . We also propose a semi-supervised learning model which we name InfoGraph * . We employ a student-teacher framework similar to Mean-Teacher method Tarvainen & Valpola ( 2017 ) . We maximize the mutual information between intermediate representations of the two models so that the student model learns from the teacher model . The student model is trained on the labeled data using a supervised objective function while the teacher model is trained on unlabeled data with InfoGraph . Using InfoGraph * , we achieve performance competitive with state-of-the-art methods on molecular property prediction . We summarize our contributions as follows : • We propose InfoGraph , an unsupervised graph representation learning method based on Deep InfoMax ( DIM ) Hjelm et al . ( 2018 ) . • We show that InfoGraph can be extended to semi-supervised prediction tasks on graphs . • We empirically show that InfoGraph surpasses state-of-the-art performance on graph classi- fication tasks with unsupervised learning and obtains performance comparable with state-ofart methods on molecular property prediction tasks using semi-supervised learning . 2 RELATED WORK . Representation learning for graphs has mainly dealt with supervised learning tasks . Recently , however , researchers have proposed algorithms that learn graph-level representations in an unsupervised manner Narayanan et al . ( 2017 ) ; Adhikari et al . ( 2018 ) . Concurrently to this work , information maximizing graph neural networks ( IGNN ) was introduced which uses mutual information maximization between edge states and transform parameters to achieve state-of-the-art predictions on a variety of supervised molecule property prediction tasks Chen et al . ( 2019b ) . In this work , our focus is on unsupervised and semi-supervised scenarios . Graph Kernels . Constructing graph kernels is a common unsupervised task in learning graph representations . These kernels are typically evaluated on node classification tasks . In graph kernels , a graph G is decomposed into ( possibly different ) { Gs } sub-structures . The graph kernel K ( G1 , G2 ) is defined based on the frequency of each sub-structure appearing in G1 and G2 respectively . Namely , K ( G1 , G2 ) = 〈fGs1 , fGs2 〉 , where fGs is the vector containing frequencies of { Gs } sub-structures , and 〈 , 〉 is an inner product in an appropriately normalized vector space . Much work has been devoted to deciding which sub-structures are more suitable than others ( refer to appendix A.1 ) . Instead of defining hand crafted similarity measures between substructures , InfoGraph adopts a more principled metric – mutual information . Contrastive methods . An important approach for unsupervised representation learning is to train an encoder to be contrastive between representations that capture statistical dependencies of interest and those that do not . For example , a contrastive approach may employ a scoring function , training the encoder to increase the score on “ real ” input ( a.k.a , positive examples ) and decrease the score on “ fake ” input ( a.k.a. , negative samples ) . For more detailed discussion , refer to appendix A.2 . Deep Graph InfoMax ( DGI ) Veličković et al . ( 2018 ) also belongs to this category , which aims to train a node encoder that maximizes mutual information between node representations and the pooled global graph representation . Although we built upon a similar methodology , our aim is different than theirs as our goal is to obtain embeddings at the whole graph level for unsupervised and semisupervised learning whereas DGI only evaluates node level embeddings . In order to differentiate our method with Deep Graph Infomax ( Veličković et al . ( 2018 ) ) , we term our model InfoGraph . Semi-supervised Learning . A comprehensive overview of semi-supervised learning ( SSL ) methods is out of the scope of this paper . We refer readers to Appendix B for a short overview or Zhu et al . ( 2003 ) ; Chapelle et al . ( 2006 ) ; Oliver et al . ( 2018 ) for more comprehensive discussions . Here , we discuss a state-of-the-art method applicable for regression tasks – Mean Teacher Tarvainen & Valpola ( 2017 ) .Mean Teacher adds a loss term which encourages the distance between the original network ’ s output and the teacher ’ s output to be small . The teacher ’ s predictions are made using an exponential moving average of parameters from previous training steps . Inspired by the “ studentteacher ” framework in Mean Teacher model , our semi-supervised model ( InfoGraph * ) deploys two separate encoders but instead of explicitly encouraging the output of the student model to be similar to the teacher model ’ s output , we enable the student model to learn from the teacher model by maximizing mutual information between intermediate representations learned by two models . 3 METHODOLOGY . Most recent work on graphs focus on supervised learning tasks or learning node representations . However , many graph analytic tasks such as graph classification , regression , and clustering require representing entire graphs as fixed-length feature vectors . Though graph-level representations can be obtained through the node-level representations implicitly , explicitly extracting the graph can be more straightforward and optimal for graph-oriented tasks . Another scenario that is important , yet attracts comparatively less attention in the graph related literature is semi-supervised learning . One of the biggest challenges in prediction tasks in biology Yan et al . ( 2017 ) ; Yang et al . ( 2014 ) or molecular machine learning Duvenaud et al . ( 2015 ) ; Gilmer et al . ( 2017 ) ; Jia & Liang ( 2017 ) is the extreme scarcity of labeled data . Therefore , semi-supervised learning , in which a large number of unlabeled samples are incorporated with a small number of labeled samples to enhance accuracy of models , will play a key role in these areas . In this section , we first formulate an unsupervised whole graph representation learning problem and a semi-supervised prediction task on graphs . Then , we present our method to learn graph-level representations . Afterwards we present our proposed model for the semi-supervised learning scenario . 3.1 PROBLEM DEFINITION . Unsupervised Graph Representation Learning . Given a set of graphs G = { G1 , G2 , ... } and a positive integer δ ( the expected embedding size ) , our goal is to learn a δ-dimensional distributed representation of every graph Gi ∈ G. We denote the number of nodes in Gi as |Gi| . We denote the matrix of representations of all graphs as Φ ∈ R|G|×δ . Semi-supervied Graph Prediction Tasks . Given a set of labeled graphs GL = { G1 , · · · , G|GL| } with corresponding output { o1 , · · · , o|GL| } , and a set of unlabeled samples GU = { G|GL|+1 , · · · , G|GL|+|GU | } , our goal is to learn a model that can make predictions for unseen graphs . Note that in most cases |GU | |GL| . 3.2 INFOGRAPH . We focus on graph neural networks ( GNNs ) —a flexible class of embedding architectures which generate node representations by repeated aggregation over local node neighborhoods . The representations of nodes are learned by aggregating the features of their neighborhood nodes , so we refer to these as patch representations . GNNs utilize a READOUT function to summarize all the obtained patch representations into a fixed length graph-level representation . Formally , the k-th layer of a GNN is h ( k ) v = COMBINE ( k ) ( h ( k−1 ) v , AGGREGATE ( k ) ( { ( h ( k−1 ) v , h ( k−1 ) u , euv ) : u ∈ N ( v ) } ) ) , ( 1 ) where h ( k ) v is the feature vector of node v at the k-th iteration/layer ( or patch representation centered at node i ) , euv is the feature vector of the edge between u and v , and N ( v ) are neighborhoods to node v. h ( 0 ) v is often initialized as node features . READOUT can be a simple permutation invariant function such as averaging or a more sophisticated graph-level pooling function Ying et al . ( 2018 ) ; Zhang et al . ( 2018 ) . We seek to obtain graph representations by maximizing the mutual information between graph-level and patch-level representations . By doing so , the graph representations can learn to encode aspects of the data that are shared across all substructures . Assume that we are given a set of training samples G : = { Gj ∈ G } Nj=1 with empirical probability distribution P on the input space . Let φ denote the set of parameters of a K-layer graph neural network . After the first k layers of the graph neural network , the input graph will be encoded into a set of patch representations { h ( k ) i } Ni=1 . Next , we summarize feature vectors at all depths of the graph neural network into a single feature vector that captures patch information at different scales centered at every node . Inspired by Xu et al . ( 2018b ) , we use concatenation . That is , hiφ = CONCAT ( { h ( k ) i } K k=1 ) ( 2 ) Hφ ( G ) = READOUT ( { hiφ } Ni=1 ) ( 3 ) where hiφ is the summarized patch representation centered at node i and Hφ ( G ) is the global representation after applying READOUT . Note that here we slightly abuse the notation of h. We define our mutual information ( MI ) estimator on global/local pairs , maximizing the estimated MI over the given dataset G : = { Gj ∈ G } Nj=1 : φ̂ , ψ̂ = arg max φ , ψ ∑ G∈G 1 |G| ∑ u∈G Iφ , ψ ( ~h u φ ; Hφ ( G ) ) . ( 4 ) Iφ , ψ is the mutual information estimator modeled by discriminator Tψ and parameterized by a neural network with parameters ψ . We use the Jensen-Shannon MI estimator ( following the formulation of Nowozin et al . ( 2016 ) ) , Iφ , ψ ( h i φ ( G ) ; Hφ ( G ) ) : = EP [ −sp ( −Tφ , ψ ( ~hiφ ( x ) , Hφ ( x ) ) ) ] − EP×P̃ [ sp ( Tφ , ψ ( ~h i φ ( x ′ ) , Hφ ( x ) ) ) ] ( 5 ) where x is an input sample , x′ ( negative sample ) is an input sampled from P̃ = P , a distribution identical to the empirical probability distribution of the input space , and sp ( z ) = log ( 1 + ez ) is the softplus function . In practice , we generate negative samples using all possible combinations of global and local patch representations across all graph instances in a batch . Since Hφ ( G ) is encouraged to have high MI with patches that contain information at all scales , this favours encoding aspects of the data that are shared across patches and aspects that are shared across scales . The algorithm is illustrated in Fig . 1 . It should be noted that our model is similar to Deep Graph Infomax ( DGI ) Veličković et al . ( 2018 ) , a model proposed for learning unsupervised node embeddings . However , there are important design differences due to the different problems that we are focusing on . First , in DGI they use random sampling to obtain negative samples due to the fact that they are mainly focusing on learning node embeddings on a graph . However , contrastive methods require a large number of negative samples to be competitive Hjelm et al . ( 2018 ) , thus the use of batch-wise generation of negative samples is crucial as we are trying to learn graph embeddings given many graph instances . Second , the choice of graph convolution encoders is also crucial . We use GIN Xu et al . ( 2018a ) while DGI uses GCN Kipf & Welling ( 2016 ) as GIN provides a better inductive bias for graph level applications . Graph neural network designs should be considered carefully so that graph representations can be discriminative towards other graph instances . For example , we use sum over mean for READOUT and that can provide important information regarding the size of the graph .

Summary: In this paper, the authors propose a graph-level representation, which extends the existing node-level representation learning mechanism. Besides, both unsupervised and semi-supervised learning are leveraged for InfoGraph and InfoGraph*, receptively. The authors naturally apply Deep Graph Infomax, a contrastive representation learning method, for the whole graph level instead of the previous node embedding learning. The experiments on graph classification and molecular property prediction indicate the effectiveness, even compared to the supervised methods.