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: The paper presents a new graph representation learning method for the whole graph under both unsupervised and semi-supervised setting. Different from existing ones using graph kernel, or graph2vec, the proposed InfoGraph is able to extract graph-level representation with fixed-length features that are generalized well. Basically, InfoGraph is parameterized by graph neural networks, but guided by mutual information loss. Experiments on both unsupervised and semi-supervised experiments on popular benchmarks demonstrate the effectiveness of InfoGraph and InfoGraph*