Iterative Graph Self-Distillation

1 INTRODUCTION . Graphs are ubiquitous representations encoding relational structures across various domains . Learning low-dimensional vector representations of graphs is critical in various domains ranging from social science ( Newman & Girvan , 2004 ) to bioinformatics ( Duvenaud et al. , 2015 ; Zhou et al. , 2020 ) . Many graph neural networks ( GNNs ) ( Gilmer et al. , 2017 ; Kipf & Welling , 2016 ; Xu et al. , 2018 ) have been proposed to learn node and graph representations by aggregating information from every node ’ s neighbors via non-linear transformation and aggregation functions . However , the key limitation of existing GNN architectures is that they often require a huge amount of labeled data to be competitive but annotating graphs like drug-target interaction networks is challenging since it needs domainspecific expertise . Therefore , unsupervised learning on graphs has been long studied , such as graph kernels ( Shervashidze et al. , 2011 ) and matrix-factorization approaches ( Belkin & Niyogi , 2002 ) . Inspired by the recent success of unsupervised representation learning in various domains like images ( Chen et al. , 2020b ; He et al. , 2020 ) and texts ( Radford et al. , 2018 ) , most related works in the graph domain either follow the pipeline of unsupervised pretraining ( followed by fine-tuning ) or InfoMax principle ( Hjelm et al. , 2018 ) . The former often needs meticulous designs of pretext tasks ( Hu et al. , 2019 ; You et al. , 2020 ) while the latter is dominant in unsupervised graph representation learning , which trains encoders to maximize the mutual information ( MI ) between the representations of the global graph and local patches ( such as subgraphs ) ( Veličković et al. , 2018 ; Sun et al. , 2019 ; Hassani & Khasahmadi , 2020 ) . However , MI-based approaches usually need to sample subgraphs as local views to contrast with global graphs . And they usually require an additional discriminator for scoring local-global pairs and negative samples , which is computationally prohibitive ( Tschannen et al. , 2019 ) . Besides , the performance is also very sensitive to the choice of encoders and MI estimators ( Tschannen et al. , 2019 ) . Moreover , MI-based approaches can not be handily extended to the semi-supervised setting since local subgraphs lack labels that can be utilized for training . Therefore , we are seeking an approach that learns the entire graph representation by contrasting the whole graph directly without the need of MI estimation , discriminator and subgraph sampling . Motivated by recent progress on contrastive learning , we propose the Iterative Graph Self-Distillation ( IGSD ) , a teacher-student framework to learn graph representations by contrasting graph instances directly . The high-level idea of IGSD is based on graph contrastive learning where we pull sim- ilar graphs together and push dissimilar graph away . However , the performance of conventional contrastive learning largely depends on how negative samples are selected . To learn discriminative representations and avoid collapsing to trivial solutions , a large set of negative samples ( He et al. , 2020 ; Chen et al. , 2020b ) or a special mining strategy ( Schroff et al. , 2015 ; He et al. , 2020 ) are necessary . In order to alleviate the dependency on negative samples mining and still be able to learn discriminative graph representations , we propose to use self-distillation as a strong regularization to guide the graph representation learning . In the IGSD framework , graph instances are augmented as several views to be encoded and projected into a latent space where we define a similarity metric for consistency-based training . The parameters of the teacher network are iteratively updated as an exponential moving average of the student network parameters , allowing the knowledge transfer between them . As merely small amount of labeled data is often available in many real-world applications , we further extend IGSD to the semi-supervised setting such that it can effectively utilize graph-level labels while considering arbitrary amounts of positive pairs belonging to the same class . Moreover , in order to leverage the information from pseudo-labels with high confidence , we develop a self-training algorithm based on the supervised contrastive loss for fine-tuning . We experiment with real-world datasets in various scales and compare the performance of IGSD with state-of-the-art graph representation learning methods . Experimental results show that IGSD achieves competitive performance in both unsupervised and semi-supervised settings with different encoders and data augmentation choices . With the help of self-training , our performance can exceed state-of-the-art baselines by a large margin . To summarize , we make the following contributions in this paper : • We propose a self-distillation framework called IGSD for unsupervised graph-level representation learning where the teacher-student distillation is performed for contrasting graph pairs under different augmented views . • We further extend IGSD to the semi-supervised scenario , where the labeled data are utilized effectively with the supervised contrastive loss and self-training . • We empirically show that IGSD surpasses state-of-the-art methods in semi-supervised graph classification and molecular property prediction tasks and achieves performance competitive with state-of-the-art approaches in unsupervised graph classification tasks . 2 RELATED WORK . Contrastive Learning Modern unsupervised learning in the form of contrastive learning can be categorized into two types : context-instance contrast and context-context contrast ( Liu et al. , 2020 ) . The context-instance contrast , or so-called global-local contrast focuses on modeling the belonging relationship between the local feature of a sample and its global context representation . Most unsupervised learning models on graphs like DGI ( Veličković et al. , 2018 ) , InfoGraph ( Sun et al. , 2019 ) , CMC-Graph ( Hassani & Khasahmadi , 2020 ) fall into this category , following the InfoMax principle to maximize the the mutual information ( MI ) between the input and its representation . However , estimating MI is notoriously hard in MI-based contrastive learning and in practice tractable lower bound on this quantity is maximized instead . And maximizing tighter bounds on MI can result in worse representations without stronger inductive biases in sampling strategies , encoder architecture and parametrization of MI estimators ( Tschannen et al. , 2019 ) . Besides , the intricacies of negative sampling in MI-based approaches impose key research challenges like improper amount of negative samples or biased negative sampling ( Tschannen et al. , 2019 ; Chuang et al. , 2020 ) . Another line of contrastive learning approaches called context-context contrast directly study the relationships between the global representations of different samples as what metric learning does . For instance , a recently proposed model BYOL ( Grill et al. , 2020 ) bootstraps the representations of the whole images directly . Focusing on global representations between samples and corresponding augmented views also allows instance-level supervision to be incorporated naturally like introducing supervised contrastive loss ( Khosla et al. , 2020 ) into the framework for learning powerful representations . Graph Contrastive Coding ( GCC ) ( Qiu et al. , 2020 ) is a pioneer to leverage instance discrimination as the pretext task for structural information pre-training . However , our work is fundamentally different from theirs . GCC focuses on structural similarity to find common and transferable structural patterns across different graph datasets and the contrastive scheme is done through subgraph instance discrimination . On the contrary , our model aims at learning graph-level representation by directly contrasting graph instances such that data augmentation strategies and graph labels can be utilized naturally and effectively . Knowledge Distillation Knowledge distillation ( Hinton et al. , 2015 ) is a method for transferring knowledge from one architecture to another , allowing model compression and inductive biases transfer . Self-distillation ( Furlanello et al. , 2018 ) is a special case when two architectures are identical , which can iteratively modify regularization and reduce over-fitting if perform suitable rounds ( Mobahi et al. , 2020 ) . However , they often focus on closing the gap between the predictive results of student and teacher rather than defining similarity loss in latent space for contrastive learning . Semi-supervised Learning Modern semi-supervised learning can be categorized into two kinds : multi-task learning and consistency training between two separate networks . Most widely used semisupervised learning methods take the form of multi-task learning : argminθ Ll ( Dl , θ ) +wLu ( Du , θ ) on labeled data Dl and unlabeled data Du . By regularizing the learning process with unlabeled data , the decision boundary becomes more plausible . Another mainstream of semi-supervised learning lies in introducing student network and teacher network and enforcing consistency between them ( Tarvainen & Valpola , 2017 ; Miyato et al. , 2019 ; Lee , 2013 ) . It has been shown that semisupervised learning performance can be greatly improved via unsupervised pre-training of a ( big ) model , supervised fine-tuning on a few labeled examples , and distillation with unlabeled examples for refining and transferring the task-specific knowledge ( Chen et al. , 2020c ) . However , whether task-agnostic self-distillation would benefit semi-supervised learning is still underexplored . 3 PRELIMINARIES . 3.1 FORMULATION . Unsupervised Graph Representation Learning Given a set of unlabeled graphs G = { Gi } Ni=1 , we aim at learning the low-dimensional representation of every graph Gi ∈ G favorable for downstream tasks like graph classification . Semi-supervised Graph Representation Learning Consider a whole dataset G = GL ∪ GU composed by labeled data GL = { ( Gi , yi ) } li=1 and unlabeled data GU = { Gi } l+u i=l+1 ( usually u l ) , our goal is to learn a model that can make predictions on graph labels for unseen graphs . And with K augmentations , we get G′L = { ( G′k , y′k ) } Klk=1 and G′U = { G′k } K ( l+u ) k=l+1 as our training data . 3.2 GRAPH REPRESENTATION LEARNING . We represent a graph instance as G ( V , E ) with the node set V and the edge set E . The dominant ways of graph representation learning are graph neural networks with neural message passing mechanisms ( Hamilton et al. , 2017 ) : for every node v ∈ V , node representation hkv is iteratively computed from the features of their neighbor nodes N ( v ) using a differentiable aggregation function . Specifically , at the iteration k we get the node embedding as : hkv = σ ( Wk · CONCAT ( hk−1v , AGGREGATEk ( { hk−1u , ∀u ∈ N ( v ) } ) ) ) ( 1 ) Then the graph-level representations can be attained by aggregating all node representations using a readout function like summation or set2set pooling ( Vinyals et al. , 2015 ) . 3.3 GRAPH DATA AUGMENTATION . It has been shown that the learning performance of GNNs can be improved via graph diffusion , which serves as a homophily-based denoising filter on both features and edges in real graphs ( Klicpera et al. , 2019 ) . The transformed graphs can also serve as effective augmented views in contrastive learning ( Hassani & Khasahmadi , 2020 ) . Inspired by that , we transform a graph G with transition matrix T via graph diffusion and sparsification S = ∑∞ k=0 θkT k into a new graph with adjacency matrix S as an augmented view in our framework . While there are many design choices in coefficients θk like heat kernel , we employ Personalized PageRank ( PPR ) with θPPRk = α ( 1 − α ) k due to its superior empirical performance ( Hassani & Khasahmadi , 2020 ) . As another augmentation choice , we randomly remove edges of graphs to attain corrupted graphs as augmented views to validate the robustness of models to different augmentation choices .

Summary: Learning graph-level representations with only labels has been explored by many works. However, it's not easy to annotate every graph. This paper applies the ideas from semi-supervised classification task to improve the representation quality learned by graph neural network. Specifically the proposed solution combines several kinds of existing techniques including diffusion graph augmentation, mean teacher consistency, debiased contrastive loss and pseudo class consistency. Finally they are combined together to act as a regularization term by utilizing the unlabelled data. From this point of view, the novelty of this work is incremental, but it's still an interesting work for improving graph-level representations.