Few-shot graph link prediction with domain adaptation

1 INTRODUCTION . Real-world link prediction problems often deal with graphs with nodes from various domains . For example , we can separate the items from a product purchase network into multiple domains such as books , electronics , luxury items . A financial transaction network can have data coming from various countries . Different domains may have different distributions , but more importantly , domains are often not balanced in most real-world datasets . For example , we may expect to see hundreds of thousands of nodes for books but only a few hundred for luxury items . There might also be billions of financial transactions in developed countries but much fewer in developing countries . We observe a similar phenomenon even in academic graph datasets . For example , in the ogbn-arxiv ( Hu et al . ( 2020 ) ) dataset , which includes pre-prints from 40 computer science subjects , papers from the most popular subject take up to 16.13 % of all the nodes . In contrast , the least represented subject only has 0.02 % of the data . Training graph models with imbalanced data without precaution significantly downgrade model performances , especially for domains with fewer samples . The training process is overwhelmed by the popular domains that have abundant data . These popular domains act as noise , hampering the model to fit for small domains effectively . Despite the prevalence of this problem , current research on this topic is not adequate ( Zhao et al . ( 2021 ) ; Shi et al . ( 2020 ) ) . To improve model performance for small domains , we propose to train domain agnostic embedding that we can use for downstream tasks . The intuition is that if the latent representation learned by the graph model does not contain any domain-specific information , and therefore is domain-agnostic , it becomes more robust to domains with fewer data and even novel domains . We propose to use adversarial learning and force the graph model to learn domain agnostic embeddings . Our idea of building domain-agnostic features across domains aligns with the few-shot learning ( FSL ) problem in computer vision ( CV ) . FSL ( Fei-Fei et al . ( 2006 ) ; Fink ( 2005 ) ) refers to the type of ML algorithm that works with limited labeled data . A common FSL technique is to learn domainagnostic features first ( with/without external data ) and then let the model converge fast on to domains with limited shots ( Long & Wang ( 2015 ) ) . It is relatively easier for CV tasks because stacked convolutional neural networks used in CV usually learn features in a general-to-specific sequence . In contrast , as stacking graph neural networks ( GNN ) layers only explore larger neighborhoods , the same method couldn ’ t be applied directly to graph data . In addition , the CV literature has explored the idea of using adversarial learning to build better domain agnostic features . Domain-Adversarial Neural Network ( DANN ) ( Ganin et al . ( 2016 ) ) first introduces the idea of using a discriminator to align and separate semantic probability distributions . Further research ( Motiian et al . ( 2017 ) ) shows that it helps reduce the number of samples needed for each domain/class . Essentially , our work follows this path but further improves it with the enclosing subgraph idea to achieve better performances for graph link predictions . We find that the method of extracting enclosing subgraphs , proposed in the SEAL model Zhang & Chen ( 2018 ) suits well for our setting because it not only boosts the number of trainable samples but also gives them different topological features . It is similar to what “ cropping ” does for images . CV Research has shown that data augmentation method , such as cropping ( Qi et al . ( 2018 ) ; Zhang et al . ( 2018b ) ) , introduces invariances for models to capture and is effective for FSL problems . We summarize our main contributions as follows : • We propose the idea of learning domain-agnostic graph embedding to improve the quality of link prediction by using adversarial learning and enclosing subgraphs . • We define the concepts of “ shots ” for graph data , making it significantly easier to train , describe , and analyze graphs with multiple imbalanced domains . • We use t-SNE plots to demonstrate the effectiveness of our method . It also allows practitioners to decide the expected gain from our proposed ideas . 2 RELATED WORKS . 2.1 LINK PREDICTION . Link prediction focuses on predicting future or missing links between a pair of nodes . The current state-of-the-arts link prediction method is SEAL ( Zhang & Chen ( 2018 ) ) . The core idea of SEAL is to extract enclosing subgraphs for a set of sampled positive and negative links . To be more specific , a positive link is a pair of nodes connected with an edge , and a negative link is a pair of nodes that are not connected . An enclosing subgraph is the union set of h-hop neighboring nodes of the pair of nodes . This data-preprocessing step lets SEAL transform the link prediction problem to a ( sub ) graph classification problem . It uses graph neural networks ( GNNs ) to generate the graph embedding , which feeds into a classifier and generates new predictions . Other than SEAL , Graph Auto-Encoder ( GAE , Kipf & Welling ( 2016 ) ) is an alternative approach . GAE works by generating node embedding for all nodes and then model the link likelihood based on the concatenated node embedding vectors for any node pairs . However , recent researches have shown this method heavily relies on the assumption that similar nodes have similar connections . In addition to these GNN-based methods , traditional heuristic methods , such as common neighbor ( CN ) and Adamic-Adar ( AA ) , provide a completely different alternative . Given a pair of nodes , these non-parametric methods describe the topological structure of the surrounding neighborhood with a score . The advantage is that these methods could be used directly without training , but it ’ s challenging to utilize node/edge features with these methods . 2.2 FEW-SHOT LEARNING AND DOMAIN ADAPTATION . Few-shot learning is a special type of machine learning problem , which aims at obtaining good learning performances when supervised evidence information is limited ( Wang et al . ( 2020 ) ) . Common FSL methods include data augmentation ( cropping , rotating , generative models , etc ) and metalearning ( Sun et al . ( 2019 ) ) . Domain Adaptation is a type of Transfer Learning ( TL ) technique . It focuses on transferring learned knowledge from one domain to another while maintaining good performances . The core idea of domain adaptation is to learn a domain invariant representation , which is usually done by maximizing mutual information ( Hjelm et al . ( 2019 ) ) or training in an adversarial network ( Goodfellow et al . ( 2014 ) ) . Some recent works focus on the problem of few-shot learning with domain adaptation . The work from Motiian et al . ( 2017 ) also uses adversarial networks as the solution , but it introduced a 4-class confusion matrix classifier as the discriminator . The work from Zhao et al . ( 2020 ) explicitly added source/target per-class separation before doing embedding feature learning with domain adaptation . 2.3 ADVERSARIAL LEARNING . To learn domain agnostic models , we need to ensure that the learned graph embedding does not contain domain-specific information . To achieve this , a commonly used technique is adversarial learning , in which two models are trained against each other and play a minimax game based on game theory . For example , in the popular Generative Adversarial Network ( GAN ) ( Goodfellow et al . ( 2014 ) ) , adversarial learning helps generators build more realistic images by having a discriminator that identifies whether an image is synthetic or real . There have been prior works done on using adversarial training on graph-related tasks . The work from Wang et al . ( 2018 ) uses GAN to improve various graph-related tasks , such as link prediction and node property prediction , by generating “ fake ” graphs . The work from Lei et al . ( 2019 ) also uses GAN to solve the challenging temporal link prediction problem for weighted dynamic networks . 3 METHODOLOGY . 3.1 PROBLEM DEFINITION . We define the few-shot link prediction problem with imbalanced domains as follows : in a given graph space S , we have access to graphs from m ( m ∈ N , m ≥ 2 ) number of domains . Different domains may have variations in their marginal distributions . Among them domains , m−1 ( k ∈ N+ ) domains have k number of small graphs , which we define as graph “ shots ” . The last domain only has 1 shot . Our objective is to learn a link prediction model that leverages the domain differences to improve the performance in the domain with 1 shot of data . 3.2 NOTATIONS . In terms of graph-related notations , each graph is denoted by G = ( V , E , d ) , where V is the set of vertices , E is the set of edges and d is the type of domain , which could be a one-hot encoded vector with one extra digit representing any unseen domain . In the adjacency matrix A of G , Au , v = 1 if edge ( u , v ) ∈ E , and otherwise Au , v = 0 . For any node u , we use N ( u ) to denote the set of neighboring nodes of u . In addition , we use superscript as a shorthand to denote the domain class . For example , for a graph from domain i , we would denote it as Gi . 3.3 DOMAINS . For our setting , we have a total of m domains and each domain has one or more graphs or subgraphs Dk = { Gk1 , ... } mk=1 , where Dk denotes the kth domain and graph Gk1 is a realization of random variable Gk . The effective of our method compared to baseline is based on the assumption that the distribution of the random variable G is different across different domains p ( Gk ) 6= p ( Gk′ ) if k′ 6= k . 3.4 LINK PREDICTION . Here is the general form of any link prediction models : given ( u , v , G ) , where u and v are a pair of nodes in graph G , the likelihood of having a link between u and v , yu , v , is modeled by function f : ( u , v , G ) → yu , v . This function could be trained with a union set of positive and negative samples . Positive samples are sampled from existing edges while negative samples are built with node pairs with u and v′ , v′ /∈ N ( u ) . Based on the ideas from the SEAL framework , the first step is to extract an h-hop enclosing subgraph around edge ( u , v ) . We denote this subgraph as Su , v and the process of generating the enclosing subgraph as s ( u , v , G ) = Su , v . Function f could be rewritten as f : Su , v → yu , v . In general , f is composed with two separated functions g and c and could be written as f = g ◦ c . Function g generates the graph embedding vector g : Su , v → hSu , v . Function c takes in the graph embedding and generates the prediction ˆyu , v . Here we have the supervised loss function . Lcls ( g , c ) = E [ ` ( c ( g ( Su , v ) ) , yu , v ) ] = − 1 N ∑ yu , v · log ( c ( g ( Su , v ) ) ) + ( 1− yu , v ) · log ( 1− c ( g ( Su , v ) ) ) , where E is the statistical expectation of any loss function ` . In this case , ` is the standard binary cross-entropy loss . The input of the discriminator D is a pair of graph embedding hi and hj generated by g. As stated above , here we use superscript to denote the domain information . The objective is to determine whether these two subgraphs are from the same domain . Since we know from which graph G these subgraphs are generated , and we know the domain information of each G , we can easily get the ground truth di , j . The loss function here could be binary cross-entropy as well . An alternative design is to let the discriminator predict the exact domain that the graph embedding is from . We select the current one because it is easier to apply it to novel domains . Ldis ( D ) = E [ ` ( D ( g ( Siu , v ) , g ( S j u′ , v′ ) ) , d i , j ) ] Given the two loss functions above , we could get the total loss for the input of a pair of subgraphs from domain i and j . Since our objective is to make sure the graph embedding does not contain domain-specific information , we would like to minimize the supervised training loss while maximizing the discriminator ’ s loss . Ltotal ( g , c ) = E [ ∑m i=1 ` ( c ( g ( S i u , v ) ) , y i u , v ) − α · ∑m i=1 ∑m j=1 ` ( D ( g ( S i u , v ) , g ( S j u′ , v′ ) ) , d i , j ) ] , where α is hyperparameter to adjust the weight of the discrimination term compared to the supervised loss .

Summary: The paper proposes a method to resolve the issue of domain imbalanced dataset in graph link prediction. The proposed method uses adversarial training to generate graph embeddings that are domain agnostic in order to facilitate transfer learning cross domains. The paper uses T-SNE plot of graph embedding to gain insights of the best scenarios for applying the proposed methods. The paper compares the proposed method with heuristics-based and GNN-based domain adaptation methods using experiments.