Slow Learning and Fast Inference: Efficient Graph Similarity Computation via Knowledge Distillation

1 Introduction . Measuring the similarity across graphs , i.e. , Graph Similarity Computation ( GSC ) , is one of the core problems of graph data mining , centered around by multiple downstream tasks such as graph retrieval [ 1 , 2 ] , plagiarism/anomaly detection [ 22 , 41 ] , graph clustering [ 39 ] , etc . As shown in Fig . 1 , the graph similarity can be defined as distances between graphs , such as Graph Edit Distance ( GED ) . The conventional solutions towards GSC are the exact computation of these graph distances , which , however , is an NP-hard problem . Therefore , such exact solutions are less favorable when handling large-scale graphs due to the expensive computation cost . Computational time , especially run time in inference stage , is particularly important in industrial scenarios . As a motivating example , in graph-structured molecules or chemical compounds query for in-silico drug screening , fast identifying similar compounds in a large database is a key process [ 25 ] . Leveraging the strong representational power of graph neural network ( GNN ) [ 21 , 13 , 43 , 42 ] , the GNN-based approximate GSC solutions have gained increasing popularity . To adapt GNNs to the GSC task , the target similarity score ( e.g. , GED ) is normalized into the range of ( 0 , 1 ] . In this way , the ∗Corresponding Author : qin.ca @ northeastern.edu 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . GSC can be regarded as a single-value regression problem that outputs a similarity score given two graphs as inputs . A standard design can be summarized as a twin of GNNs bridged by a co-attention with a Multi-layer Perceptron ( MLP ) stacked as the regression head . Such approaches can be trained in a fully supervised way using the Mean Square Error ( MSE ) loss computed over the ground truth similarity score . Many GNN-based GSC methods [ 1 , 2 , 22 ] followed such strategy , which , however , suffers from the fusion issue . The paper presents a novel solution to both effectively and efficiently address the task of approximate GSC . Compared to the commonly used graph convolutional network as the backbone [ 1 , 2 ] , this paper adopts a more robust network , i.e. , Graph Isomorphism Network ( GIN ) [ 43 ] . Cross-graph fusion is essential to the model . The multi-scale features within different GIN layers are fused with a new design . We have adopted an attention layer stacked over the concatenated cross-graph features for smooth feature fusion . To this end , similar features will be assigned with more weights to contribute to the desired task . Moreover , to make the model easier to deploy , we take an MLP for feature learning which is simple but effective to achieve cutting-edge performance . Intuitively , speed and accuracy can be considered as a trade-off . GSC naturally requires dense connections/interactions between the two input graphs , which will consequently cause increasing computations as the cost . This paper focuses on the efficiency of inference speed which can be addressed by either model compression or a faster data loading pipeline . Especially in industrial scenarios , the raw graph data are usually pre-processed as the embeddings off-line that can be easily applied to the real-time downstream tasks , e.g. , molecular graph retrieval . However , as shown in Fig . 2 , most of the co-attention-based GSC solutions employ feature fusion in the early stage , which only outputs the joint embedding of pairing graphs . Inspired by [ 26 ] , we propose a lightweight model that removes all the early feature fusion modules in the encoder for efficient GSC . In this way , as shown in Fig . 2 , the individual embedding of each graph can be collected by a Siamese GNN . Such pairing graph embeddings will be fused with an attention layer to predict the final similarity score . To overcome the accuracy drop of such a small network , we take a novel paradigm of Knowledge Distillation ( KD ) specifically designed for our task . As shown in Fig . 3 , we propose an early-feature fusion network regarded as the teacher model , and the student model is a siamese network without co-attention . It is found that the direct distillation of joint embeddings fails to work where the KD loss disturbs largely during training . To solve this , we generate the pseudo individual embeddings of the teacher model and use them for KD by minimizing their relational distances [ 29 ] . To ensure pseudo individual embeddings fully cover the information of raw graphs , we further apply an MSE loss on the reconstructed joint embeddings concatenated from pseudo individual ones . We have verified that there is only a marginal accuracy drop compared with the original joint embeddings , which justifies the claim above . To sum up , our contributions can be summarized in three folds : • We introduce a new early-feature fusion model to achieve the competitive accuracy by designing a strong co-attention network and taking the GIN as the backbone . • For efficient inference and off-line embedding collection , we propose a novel Knowledge Distillation method for GSC where the joint embeddings are decomposed to distill . • Extensive experiments on the popular GED benchmarks demonstrate the superiority of our model over the state-of-the-art GSC methods on both accuracy and efficiency . Compared with the co-attention models , there is a more than 10 times faster in inference speed compared with the best competitor on AIDS dataset . † 2 Related Works . 2.1 Graph Similarity Computation ( GSC ) . Graph similarity computation measures the similarity of two given graphs , where similarity metrics can be defined as Graph Edit-Distance [ 6 ] , Graph Isomorphism [ 8 ] , and Maximum Common Subgraph [ 7 ] . Exact computation of these metrics is generally an NP-complete problem [ 48 ] . To speed up the computation , kernel-based methods have been extensively proposed to approximate the exact solvers [ 44 , 3 , 28 , 45 ] . Recently , inspired by the strong representation power of deep neural network , a number of neural network based methods have been proposed and demonstrated a huge success [ 47 , 1 , 2 , 24 , 22 , 41 , 40 ] . Among them , regression-based similarity learning has a great promise due to the competitive performance in both efficiency and efficacy [ 1 , 22 , 2 ] . The intuition here is to learn an embedding vector using a graph neural network ( GNN ) , and then measure the similarity of graph embeddings . While such a graph-level embedding encoded by GNN alone is not sufficient to well distinguish the nuances of subgraph level structures . To integrate subgraph information for final similarity computation , several methods are proposed recently , such as node-level pairwise comparison [ 1 ] , cross-graph attention-based matching [ 22 ] , multi-scale neighbor aggregation [ 2 ] , etc . Despite the superior efficacy reported under various metrics ( such as Accuracy , Mean Squared Error ( mse ) , Spearman ’ s Rank Correlation Coefficient ) , the complex subgraph matching/fusion components ( termed ‘ early-fusion ’ in Fig . 2 ) in different layers dramatically slow down the similarity measure . Moreover , early-fusion prevents pre-computing the embeddings for all candidate graphs for further reducing inference time in the graph retrieval scenario . Motivated by this , we propose a slow learning and fast inference method by leveraging the knowledge distillation idea to transfer the fine-grained but slowly learned early-fusion teach model to the fast-inference student model . 2.2 Knowledge Distillation ( KD ) . Knowledge distillation is a general neural network training method , where a ( typically pretrained ) teacher network is introduced to guide the learning of a student network . Its idea was first pioneered by Bucilua et al . [ 5 ] to compress large machine learning models , where they proposed to transfer the knowledge of a model ensemble into a neural network by labeling unlabeled data as transfer set . This idea was later refined by Hinton et al . [ 16 ] , where they adopted softened probabilities of the teacher as a target for the student to learn and coined the term “ knowledge distillation ” . Ever since , many methods have been proposed revolving around the central question in KD : “ what is the definition of knowledge to be distilled ” . Popular definitions include feature distance [ 34 ] , feature map attention [ 46 ] , feature distribution [ 30 ] , activation boundary [ 15 ] , inter-sample distance structure [ 29 , 32 , 23 , 36 ] , and mutual information [ 35 ] . See [ 38 , 11 ] for a more comprehensive survey . [ 26 ] is proposed to distill separate models from a co-attention one . Despite the progress , they mainly focus on convolutional neural networks for vision tasks ( mainly image recognition ) or recurrent neural networks for sequential data tasks ( e.g. , for natural language understanding [ 19 ] ) . 3 Approach . This section will introduce 1 ) the architecture of the early-fusion network ( i.e. , teacher model ) ; 2 ) the KD process and its interpretation . Before that , we start from the formalized problem definition . 3.1 Problem Formulation . Formally , a graph G is defined upon the node set V and edge set E as G = ( V , E ) . In specific , the edge linking a pair of nodes including u ∈ V and v ∈ V can be denoted as the ( u , v ) ∈ E . In our setting , all the accessible graphs are undirected , i.e. , ( u , v ) ∈ E ↔ ( v , u ) ∈ E . The quantity of nodes is represented as N = |V| . A convenient way to represent the graphs is the adjacency matrix †The code is uploaded on https : //github.com/canqin001/Efficient_Graph_Similarity_ Computation A ∈ R|V|×|V| . We denote the presence of edges as A [ u , v ] = 1 if ( u , v ) ∈ E and A [ u , v ] = 0 otherwise . Mostly , graph attributes ( e.g. , node labels ) are available . Such node-level features can be denoted as a real-value matrix X ∈ R|V|×m with the m dimension and the order of feature matrix X is consistent with the adjacency matrix [ 12 ] . In GSC task , we have the access to pairing graphs Gi and Gj ∈ D , whereD = { G0 , G1 , ... } is the graph set . The similarity of such two graphs can be represented as Graph Edit Distance ( GED ) or Maximum Common Subgraph ( MCS ) . As shown in Fig . 1 , the GED is defined as the number of edit operations in the optimal trajectory to transform the source graph to the target . The MCS is the maximum subgraph common to both two graphs . To well fit GNN , the standard GED value is normalized as the nGED , i.e. , nGED ( Gi , Gj ) = GED ( Gi , Gj ) ( |G1|+|G2| ) /2 . In the following , nGED should be transformed to the value ranging ( 0 , 1 ] as the ground truth similarity score sij , i.e. , sij = exp ( −nGED ( Gi , Gj ) ) ∈ S , where S ∈ R|D|×|D| indicates the similarity matrix among all the graphs [ 2 ] .

Summary: The paper studies efficient graph similarity computation and proposes two connected techniques: a complex teacher model slowly learning through the graph and a simpler student model jointly learned and proposed to be used for faster inference. The proposed technique uses well-established Fusion network structures. The joint learning through knowledge distillation enables the simpler student model to learn more accurate latent representations through the teacher network. Both accuracy and run-time results have been provided using benchmark data sets.