Motif-Driven Contrastive Learning of Graph Representations

1 INTRODUCTION . Graph-structured data , such as molecules and social networks , is ubiquitous in many scientific research areas and real-world applications . To represent graph characteristics , graph motifs were proposed in Milo et al . ( 2002 ) as significant subgraph patterns occurring frequently in graphs and uncovering graph structural principles . For example , functional groups are important motifs that can determine molecule properties . Like the hydroxide ( –OH ) usually implies higher water solubility , and for proteins , Zif268 can mediate protein-protein interactions in sequence-specific DNA-binding proteins . ( Pabo et al. , 2001 ) . Graph motifs has been studied for years . Meaningful motifs can benefit many important applications like quantum chemistry and drug discovery ( Ramsundar et al. , 2019 ) . However , extracting motifs from large graph datasets remains a challenging question . Traditional motif discovery approaches ( Milo et al. , 2002 ; Kashtan et al. , 2004 ; Chen et al. , 2006 ; Wernicke , 2006 ) rely on discrete counting or statistical estimation , which are hard to generalize to large-scale graph datasets with continuous and high-dimension features , as often the case in real-world applications . Recently , Graph Neural Networks ( GNNs ) have shown great expressive power for learning graph representations without explicit feature engineering ( Kipf & Welling , 2016 ; Hamilton et al. , 2017 ; Veličković et al. , 2017 ; Xu et al. , 2018 ) . In addition , GNNs can be trained in a self-supervised manner without human annotations to capture important graph structural and semantic properties ( Veličković et al. , 2018 ; Hu et al. , 2020c ; Qiu et al. , 2020 ; Bai et al. , 2019 ; Navarin et al. , 2018 ; Wang et al. , 2020 ; Sun et al. , 2019 ; Hu et al. , 2020b ) . This motivates us to rethink about motifs as more general representations than exact structure matches and ask the following research questions : • Can we use GNNs to automatically extract graph motifs from large graph datasets ? • Can we leverage the learned graph motif to benefit self-supervised GNN learning ? In this paper , we propose MICRO-Graph : a framework for MotIf-driven Contrastive leaRning Of Graph representations . The key idea of this framework is to learn graph motifs as prototypical cluster centers of subgraph embeddings encoded by GNNs . In this way , the discrete counting problem is transfered to a fully-differentiable framework that can generalize to large-scale graph datasets with continuous and high-dimensional features . In addition , the learned motifs can help generate more informative subgraphs for graph-to-subgraph contrastive learning . The motif learning and contrastive learning are mutually reinforced to pre-train a more generalizable GNN encoder . For motif learning , given a graph dataset , a motif-guided subgraph segmenter generates subgraphs from each graph , and a GNN encoder turns these subgraphs into vector representations . We then learn graph motifs through clustering , where we keep the K prototypical cluster centers as representations of motifs . Similar and significant subgraphs are assigned to the same motif and become closer to their corresponding motif representation . We train our model in an Expectation-Maximization ( EM ) fashion to update both the motif assignment of each subgraph and the motif representations . For leveraging learned motifs , we propose a graph-to-subgraph contrastive learning framework for GNN pre-training . One of the key components for contrastive learning is to generate semantically meaningful views of each instance . For example , a continuous span within a sentence ( Joshi et al. , 2020 ) or a random crop of an image ( Chen et al. , 2020 ) . For graph data , previous approaches leverage node-level views , which is not sufficient to capture high-level graph structural information Sun et al . ( 2019 ) . As motifs can represent the key graph properties by its nature , we propose to leverage the learned motifs to generate more informative subgraph views . For example , alpha helix and beta sheet can come together as a simple ββα fold to form a zinc finger protein with unique properties . By learning such subgraph co-occurrence via contrastive learning , the pre-trained GNN can capture higher-level information of the graph that node-level contrastive can ’ t capture . The pre-trained GNN using MICRO-Graph on the ogbg-molhiv molecule dataset can successfully learn meaningful motifs , including Benzene rings , nitro , acetate , and etc . Meanwhile , fine-tune this GNN on seven chemical property prediction benchmarks yielding 2.0 % average improvement over non-pretrained GNNs and outperforming other self-supervised pre-training baselines . Also , extensive ablation studies show the significance of the learned motifs for the contrastive learning . 2 RELATED WORK . The goal of self-supervised learning is to train a model to capture significant characteristics of data without human annotations . This paper studies whether we can use such approach to automatically extract graph motifs , i.e . the significant subgraph patterns , and leverage the learned motifs to benefit self-supervised learning . In the following , we first review graph motifs especially challenges for motif mining , and then discuss approaches for pre-training GNNs in a self-supervised manner . Graph motifs are building blocks of complex graphs . They reveal the interconnections of graphs and represent graph characteristics . Mining motifs can benefit many tasks from exploratory analysis to transfer learning ( Henderson et al. , 2012 ) . For many years , various motif mining algorithms have been proposed . There are generally two categories , either exact counting as in Milo et al . ( 2002 ) ; Kashtan et al . ( 2004 ) ; Schreiber & Schwöbbermeyer ( 2005 ) ; Chen et al . ( 2006 ) , or sampling and statistical estimation as in Wernicke ( 2006 ) . However , both approaches can not scale to large graph datasets with high-dimension and continuous features , which is common in real-world applications . In this paper , we proposes to turn the discrete motif mining problem into a GNN-based differentiable cluster learning problem that can generalize to large-scale datasets . Another GNN-based work related to graph motifs is the GNNExplainer , which focuses on post-process model interpretation ( Ying et al. , 2019 ) . It can identify substructures that are important for graph property prediciton , e.g . motifs . The difference between GNNExplainer and MICRO-Graph is that the former identify motifs at a single graph level , and the later learns motifs across the whole dataset . Contrastive learning is one of the state-of-the-art self-supervised representation learning algorithms . It achieves great results for visual representation learning ( Chen et al. , 2020 ; He et al. , 2019 ) . Contrastive learning forces views generated from the same instance ( e.g . different crops of the same image ) to become closer , while views from different instances apart . One key component in contrastive learning is to generate informative and diverse views from each data instance . In computer vision , researchers use various techniques , including cropping , color distortion , and Gaussian blurs to generate views . However , when it comes to graphs , constructing informative view of graph is a challenging task . In our framework , we utilize the learned motifs , which are significant subgraph patterns , to guide view ( subgraph ) generation , and conduct graph-to-subgraph contrastive learning . Self-supervised learning for GNNs also draws many attention recently . For graphs , representations can be at different levels , e.g . node level and ( sub ) graph level . Veličković et al . ( 2018 ) ; Hu et al . ( 2020c ) ; Qiu et al . ( 2020 ) mainly focus on node-level representation learning in a single large graph , as opposed to the focus of this paper , which is representation learning of whole graphs . Hu et al . ( 2020b ) provides a systematic analysis of pre-training strategies on graphs for both node-level and graph-level . However , only the node-level learning is self-supervised , and annotated labels are utilized for supervised learning at the graph level . For graph level self-supervised representation learning , Sun et al . ( 2019 ) proposed a contrastive framework , InfoGraph , to maximize the mutual information between graph representations and node representations . In Rong et al . ( 2020 ) , the GROVER model and a motif based self-supervised learning task was proposed , where the discrete motifs are first extracted using a professional software , and then these motifs are used as prediction labels for pre-training the model . The difference between motifs in GROVER and in MICRO-Graph is that GROVER uses discrete structures , but MICRO-Graph uses continuous vector embeddings To alleviate these issues , we propose graph-to-subgraph view self-supervised contrastive learning , and the subgraph generation is guided by the learned motifs . 3 METHODOLOGY . The goal of this paper is to train a GNN encoder that can automatically extract graph motifs , i.e . significant subgraph patterns . Motif discovery on discrete graph structures is a combinatorial problem , and it is hard to generalize to large datasets with continuous features . We thus propose to formalize this problem as a differentiable clustering learning problem and solve it via self-supervised GNN learning . In this section , we formalize the problem and introduce the overall framework of MICROGraph in Section 3.1 , and then describe each module in details in the following sections . 3.1 THE OVERALL FRAMEWORK OF MICRO-Graph Given a dataset with M graphs G = { G1 , ... , GM } , the differentiable clustering learning problem is meant to learn two things . One is a GNN-based graph encoder E ( · ) that maps input ( sub ) graphs an embedding vector . The other is a K-slot embedding table { m1 , ... , mK } , where each slot is a motif vector m corresponding to a cluster center of embeddings of frequently occurred subgraphs . To tackle this problem , we introduce the MICRO-Graph framework which consists three modules : 1 ) a Motif-guided Segmenter to extract important subgraphs ; 2 ) a Motif-Learner to cluster sampled subgraphs and identify motifs ; 3 ) a constrastive learning module for graph-to-subgraph contrastive learning . The overall framework is shown in Figure1 . We describe details of each module in the following sections . The Motif Learner is introduced first in Section 3.2 , and then the Motif-guided Segmenter and the constrastive learning module in Section 3.3 and 3.4 respectively . 3.2 MOTIF LEARNER VIA EM CLUSTERING . The Motif Learner learns motifs by applying an Expectation-Maximization ( EM ) style clustering algorithm on sampled subgraphs . To start clustering , the Motif-guided Segmenter first extracts N subgraphs { gj } Nj=1 from the input whole graphs . For each subgraph gj , we generate its embedding ej = E ( gj ) and calculate the cosine similarity between ej and each of the K motif vectors { m1 , ... , mK } . We denote the similarity between ej and the kth motif vector as Sk , j = φ ( mk ) Tφ ( ej ) 1 . In vector notation , the K-dimensional vector of similarities between gj and all K motif vectors is denoted as sj , and the K-by-N-dimensional motif-to-subgraph similarity matrix is denoted as S , where the j-th column of S is sj and the entry ( k , j ) of S is Sk , j . E-Step . The goal of the E-step is to come up with motif-based cluster assignments for subgraph embeddings { ej } Nj=1 . The assignments can be represented by a K-by-N-dimensional matrix Q = [ q1 , ... , qN ] , where the j-th column qj contains the probabilities of assigning the j-th subgraph to each of the K motifs . Each qj can be a one-hot vector for hard clustering or a probability vector with all the entries sum up to one for soft-clustering . This vanilla version clustering problem boils down to maximizing the objective Tr ( QTS ) , which corresponds to an assignment Q that maximizes similarities between embeddings and its assigned motif . This objective works fine for a traditional EM clustering algorithm when embeddings are fixed . However , since representations will change when doing representation learning , this vanilla objective in the E-step can lead to a degenerate solution , i.e . all representations collapse to a single cluster center . To avoid this issue , we follow YM . et al . ( 2020 ) to introduce an entropy term and an equal-size constraint on Q for clusters to have similar sizes . Our final objective is : max Q∈Q Tr ( QTS ) + 1 λ H ( Q ) ( 1 ) where H ( Q ) = − ∑ i , j Qi , j logQi , j is the entropy function , and the constraint set Q requires the marginal projection of Q onto its columns and rows to be uniform . Q = { Q ∈ RK , N+ |Q1N = 1K K , Q1K = 1N N } ( 2 ) where 1N and 1K are all one vectors . This constraint optimization problem turns out to be an optimal transportation problem with a closed-form solution as ( 3 ) and can be solved efficiently using a fast Sinkhorn-Knopp algorithm . Q∗ = diag ( u ) · exp ( λS ) · diag ( v ) ( 3 ) Here u and v are normalization vectors . The derivations can be found in Cuturi ( 2013 ) . M-Step . The goal of the M-step is to maximize the log-likelihood of our data given the cluster assignment matrix Q estimated in the E-step . We update parameters in the GNN encoder and the motif embedding table through the M-step . This step is equivalent to a supervised K-class classification problem with labels Q and prediction scores S. Thus , we first apply a columnwise softmax normalization with temperature τg to S to convert all entries of S to probabilities , i.e . S̃k , j = softmaxk ( Sk , j/τg ) . Then we use the negative likelihood as the loss function . Lm = − 1 N N∑ j=1 K∑ k=1 Qk , j log S̃k , j ( 4 )

Summary: This paper proposes to learn the sub-graph patterns from a collection of training graphs. The key idea is to partition each graph into segments and enforce a global clustering of the subgraphs. The partitioning is also guided through contrastive learning, i.e., subgraphs should have a larger similarity with the graph it is drawn from, compared with other graphs. The learned GNN (that generates node embedding) will then be used to some downstream learning tasks with or without further fine-tuning.