Counterfactual Graph Learning for Link Prediction

1 INTRODUCTION . Link prediction seeks to predict the likelihood of edge existence between node pairs based on the observed graph . Given the omnipresence of graph-structured data , link prediction has copious applications such as movie recommendation ( Bennett et al. , 2007 ) , chemical interaction prediction ( Stanfield et al. , 2017 ) , and knowledge graph completion ( Kazemi & Poole , 2018 ) . Graph machine learning methods have been widely applied to solve this problem . Their standard scheme is to first learn the representation vectors of nodes and then learn the association between the representations of a pair of nodes and the existence of the link between them . For example , graph neural networks ( GNNs ) use neighborhood aggregation to create the representation vectors : the representation vector of a node is computed by recursively aggregating and transforming representation vectors of its neighboring nodes ( Kipf & Welling , 2016a ; Hamilton et al. , 2017 ; Wu et al. , 2020 ) . Then the vectors are fed into a binary classification model to learn the association . GNN methods have shown predominance in the task of link prediction ( Kipf & Welling , 2016b ; Zhang et al. , 2020a ) . Unfortunately , the causal relationship between graph structure and link existence was largely ignored in previous work . Existing methods that learn from association only were are not able to capture essential factors to accurately predict missing links in the test data . Take social network as an example . Suppose Alice and Adam live in the same neighborhood and they are close friends . The association between neighborhood belonging and friend closeness could be too strong to discover the essential factors of the friendship such as common interests or family relationship which could be the cause of being living in the same neighborhood . So , our idea is asking a counterfactual question : “ would Alice and Adam still be close friends if they were not living in the same neighborhood ? ” If a graph learning model could learn the causal relationship by answering the counterfactual questions , it would improve the accuracy of link prediction with the novel knowledge it captured . Generally , the questions can be described as “ would the link exist or not if the graph structure became different ? ” As known to many , counterfactual question is a key component of causal inference and have been well defined in the literature . A counterfactual question is usually framed with three factors : context ( as a data point ) , manipulation ( e.g. , treatment , intervention , action , strategy ) , and outcome ( van der Laan & Petersen , 2007 ; Johansson et al. , 2016 ) . ( To simplify the language , we use “ treatment ” to refer to the manipulation in this paper , as readers might be familiar more with the word “ treatment. ” ) Given certain data context , it asks what the outcome would have been if the treatment had not been the observed value . In the scenario of link prediction , we consider the information of a pair of nodes as context , graph structural properties as treatment , and link existence as outcome . Recall the social network example . The context is the representations of Alice and Adam that are learned from their personal attributes and relationships with others on the network . The treatment is living in the same neighborhood , which can be identified by community detection . And the outcome is their friendship . In this work , we present a counterfactual graph learning method for link prediction ( CFLP ) that trains graph learning models to answer the counterfactual questions . Figure 1 illustrates this twostep method . Suppose the treatment variable is defined as one type of global graph structure , e.g. , the neighborhood assignment discovered by spectral clustering or community detection algorithms . We are wondering how likely the neighborhood distribution makes a difference on the link ( non ) existence for each pair of nodes . So , given a pair of nodes ( like Alice and Adam ) and the treatment value on this pair ( in the same neighborhood ) , we find a pair of nodes ( like Helen and Bob ) that satisfies two conditions : ( 1 ) it has a different treatment ( in different neighborhoods ) and ( 2 ) it is the most similar pair with the given pair of nodes . We call these matched pair of nodes as “ counterfactual links. ” Note that the outcome of the counterfactual link can be either 1 or 0 , depending on whether there exists an edge between the matched pair of nodes ( Helen and Bob ) . The counterfactual link provides unobserved outcome to the given pair of nodes ( Alice and Adam ) under a counterfactual condition ( in different neighborhoods ) . After counterfactual links are created for all ( positive and negative ) training examples , CFLP trains a link predictor ( which can be GNN-based ) to learn the representation vectors of nodes to predict both the observed factual links and the counterfactual links . In this Alice-Adam example , the link predictor is trained to estimate the individual treatment effect ( ITE ) of neighborhood assignment as 1 − 1 = 0 , where ITE is a metric for the effect of treatment on the outcome and zero indicates the given treatment has no effect on the outcome . So , the learner will try to discover the essential factors on the friendship between Alice and Adam . CFLP leverages causal models to find these factors for graph learning models to accurately predict missing links . Contributions . Our main contributions can be summarized as follows . ( 1 ) This is the first work that proposes to improve link prediction by causal inference , specifically , learning to answer counterfactual questions about link existence . ( 2 ) This work introduces CFLP that trains GNN-based link predictors to predict both factual and counterfactual links . It learns the causal relationship between global graph structure and link existence . ( 3 ) CFLP outperforms competitive baseline methods on several benchmark datasets . We analyze the impact of counterfactual links as well as the choice of treatment variable . This work sheds insights for improving graph machine learning with causal analysis , which has not been extensively studied yet , when the other direction ( machine learning for causal inference ) has been studied for a long time . 2 PROBLEM DEFINITION . Notations Let G = ( V , E ) be an undirected graph of N nodes , where V = { v1 , v2 , . . . , vN } is the set of nodes and E ⊆ V × V is the set of observed links . We denote the adjacency matrix as A ∈ { 0 , 1 } N×N , whereAi , j = 1 indicates nodes vi and vj are connected and vice versa . We denote the node feature matrix as X ∈ RN×F , where F is the number of node features and xi ( bolded ) indicates the feature vector of node vi ( the i-th row of X ) . In this work , we follow the commonly accepted problem definition of link prediction on graph data ( Zhang & Chen , 2018 ; Zhang et al. , 2020a ; Cai et al. , 2021 ) : Given an observed graph G ( with validation and testing links masked off ) , predict the link existence between every pair of nodes . More specifically , for the GNN-based link prediction methods , they learn low-dimensional node representations Z ∈ RN×H , where H is the dimensional size of latent space such that H F , and then use Z for the prediction of link existence . 3 PROPOSED METHOD . 3.1 IMPROVING GRAPH LEARNING WITH CAUSAL MODEL Z ? T Y Z ? T Y Treatment Effect Estimation Graph Representation Learning Figure 2 : Causal modeling ( not the target of our work but related to the idea we propose ) : Given Z and observed outcomes , find treatment effect of T on Y . Z ? T Y Z ? T Y Treatment Effect Estimation Graph Representation Learning Figure 3 : Graph learning with causal model ( the proposed idea ) : leverage the estimated ATE ( Y |T ) to improve the learning of Z. Leveraging Causal Model ( s ) Counterfactual causal inference aims to find out the causal relationship between treatment and outcomes by asking the counterfactual questions such as ” would the outcome be different if the treatment was different ? ” ( Morgan & Winship , 2015 ) . Figure 2 is a typical example , in which we denote the context ( confounder ) as Z , treatment as T , and the outcome as Y . Given the context , treatments , and their corresponding outcomes , counterfactual inference methods aim to find the effect of treatment on the outcome , which is usually measured by individual treatment effect ( ITE ) and its expectation averaged treatment effect ( ATE ) ( van der Laan & Petersen , 2007 ; Weiss et al. , 2015 ) . For a binary treatment variable T = { 0 , 1 } , denoting g ( z , T ) as the outcome of z given the treatment T , we have ITE ( z ) = g ( z , 1 ) − g ( z , 0 ) , and ATE = Ez∼Z ITE ( z ) . Ideally , we need all possible outcomes of the contexts under all kinds of treatments to study the causal relationships ( Morgan & Winship , 2015 ) . However , in reality , the fact that we can only observe one potential outcome under one particular treatment prevents the ITE from being known ( Johansson et al. , 2016 ) . Traditional causal inference methods use statisti- cal learning approaches such as Neyman–Rubin casual model ( BCM ) and propensity score matching ( PSM ) to predict the value of ATE ( Rubin , 1974 ; 2005 ) . In this work , we look at link prediction with graph learning , which is essentially learning the best node representations Z for the prediction of link existence . Therefore , as shown in Figure 3 , where the outcome Y is the link existence , the objective is different from classic causal inference . In graph learning , we can estimate the effect of treatment on the outcome ( ATE ( Y |T ) ) , and we want to improve the learning of Z with the estimation . More specifically , in graph learning for link prediction , for each pair of nodes ( vi , vj ) , its ITE can be estimated with ITE ( vi , vj ) = g ( ( zi , zj ) , 1 ) − g ( ( zi , zj ) , 0 ) ( 1 ) and we use this information to improve the learning of Z , i.e. , P ( Z|Y ) . We denote the observed adjacency matrix as the factual outcomes A and the unobserved adjacency matrix when the treatment is different as the counterfactual outcomes ACF . We denote T ∈ { 0 , 1 } N×N as the binary factual treatment matrix , where Ti , j indicates the treatment of the node pair ( vi , vj ) . We denote TCF as the counterfactual treatment matrix where TCFi , j = 1 − Ti , j . We are interested in ( 1 ) estimating the counterfactual outcomes ACF via observed data , ( 2 ) learning with the counterfactual adjacency matrix ACF to enhance link prediction , and ( 3 ) learning the causal relationship between graph structural information ( treatment ) and link existence ( outcome ) . Treatment Variable Previous work on graph machine learning ( Velickovic et al. , 2019 ; Park et al. , 2020 ) showed that the graph ’ s global structural information could improve the quality of representation vectors of nodes learned by GNNs . This is because the message passing-based GNNs aggregate local information in the algorithm of representation vector generation and the global structural information is complementary with the aggregated information . Therefore , for a pair of nodes , one option of defining the treatment variable is its global structural role in the graph . Without the loss of generality , we use Louvain ( Blondel et al. , 2008 ) , an unsupervised approach that has been widely used for community detection , as an example . Louvain discovers community structure of a graph and assigns each node to one community . Then we can define the binary treatment variable as whether these two nodes in the pair belong to the same community . Let c : V → N be any graph mining/clustering method that outputs the index of community/cluster/neighborhood that each node belongs to . The treatment matrix T is defined as Ti , j = 1 if c ( vi ) = c ( vj ) , and Ti , j = 0 otherwise . For the choice of c , we suggest methods that group nodes based on global graph structural information , including but not limited to Louvain ( Blondel et al. , 2008 ) , K-core ( Bader & Hogue , 2003 ) , and spectral clustering ( Ng et al. , 2001 ) .

Summary: The authors propose a model for link prediction that takes into account the counterfactuals. The main idea is to learn an embedding of the nodes and then use that embedding to train a classifier model for link prediction. The key part is that the embedding is learned in such a way as to minimize the discrepancy between the factual and counterfactual distributions. Experiments of various commonly-used datasets show that the proposed method results in improved link prediction performance compared to state-of-the-art methods.