FakeEdge: Alleviate Dataset Shift in Link Prediction
Abstract
Link prediction is a crucial problem in graph-structured data.
Due to the recent success of graph neural networks (GNNs), a variety of GNN-based models were proposed to tackle the link prediction task.
Specifically, GNNs leverage the message passing paradigm to obtain node representation, which relies on link connectivity.
However, in a link prediction task, links in the training set are always present while ones in the testing set are not yet formed, resulting in a discrepancy of the connectivity pattern and bias of the learned representation.
It leads to a problem of dataset shift which degrades the model performance.
In this paper, we first identify the dataset shift problem in the link prediction task and provide theoretical analyses on how existing link prediction methods are vulnerable to it.
We then propose FakeEdge, a model-agnostic technique, to address the problem by mitigating the graph topological gap between training and testing sets.
Extensive experiments demonstrate the applicability and superiority of FakeEdge on multiple datasets across various domains.
Introduction
Graph structured data is ubiquitous across a variety of domains, including social networks [1], protein-protein interactions [2], movie recommendations [3] and citation networks [4].
It provides a non-Euclidean structure to describe the relations among entities.
The link prediction task is to predict missing links or new forming links in an observed network [5].
Recently, with the success of graph neural networks (GNNs) for graph representation learning [6–9], several GNN-based methods have been developed [10–14] to solve link prediction tasks.
These methods encode the representation of target links with the topological structures and node/edge attributes in their local neighborhood.
After recognizing the pattern of observed links (training sets), they predict the likelihood of forming new links between node pairs (testing sets) where no link is yet observed.
Nevertheless, existing methods pose a discrepancy of the target link representation between training and testing sets.
As the target link is never observed in the testing set by the nature of the task, it will have a different local topological structure when compared to its counterpart from the training set.
Thus, the corrupted topological structure shifts the target link representation in the testing set, which we recognize it as a dataset shift problem [15, 16] in link prediction.
Note that there are some existing work [11] applying edge masking to moderate such a problem, similar to our treatment.
However, they tend to regard it as an empirical trick and fail to identify the fundamental cause as a problem of dataset shift.
We give a concrete example to illustrate how dataset shift can happen in the link prediction task, especially for GNN-based models with message passing paradigm [17] simulating the 1-dimensional
Weisfeiler-Lehman (1-WL) test [18].
In Figure 1, we have two local neighborhoods sampled as subgraphs from the training (top) and testing (bottom) set respectively.
The node pairs of interest,
[Figure] [Figure]
in
[Figure]
Link in
[Figure]
ObservedUnobserved
[Figure]
AggregateNeighbors
[Figure]
=
[Figure] [Figure]
Dataset shift can substantially degrade model performance since it violates the common assumption that the joint distribution of inputs and outputs stays the same in both the training and testing set.
The root cause of this phenomenon in link prediction is the unique characteristic of the target link: the link always plays a dual role in the problem setting and determines both the input and the output for a link prediction task.
The existence of the link apparently decides whether it is a positive or negative sample (output).
Simultaneously, the presence of the link can also influence how the representation is learned through the introduction of different topological structures around the link (input).
Thus, it entangles representation learning and labels in the link prediction problem.
To decouple the dual role of the link, we advocate a framework, namely subgraph link prediction , which disentangles the label of the link and its topological structure.
As most practical link prediction methods make a prediction by capturing the local neighborhood of the link [1, 11, 12, 19, 20], we unify them all into this framework, where the input is the extracted subgraph around the focal node pair and the output is the likelihood of forming a link incident with the focal node pair in the subgraph.
From the perspective of the framework, we find that the dataset shift issue is mainly caused by the presence/absence of the focal link in the subgraph from the training/testing set.
This motivates us to propose a simple but effective technique, FakeEdge , to deliberately add or remove the focal link in the subgraph so that the subgraph can stay consistent across training and testing.
FakeEdge is a model-agnostic technique, allowing it to be applied to any subgraph link prediction model.
It assures that the model would learn the same subgraph representation regardless of the existence of the focal link.
Lastly, empirical experiments prove that diminishing the dataset shift issue can significantly boost the link prediction performance on different baseline models.
We summarize our contributions as follows.
We first unify most of the link prediction methods into a common framework named as subgraph link prediction, which treats link prediction as a subgraph classification task.
In the view of the framework, we theoretically investigate the dataset shift issue in link prediction tasks, which motivates us to propose FakeEdge, a model-agnostic augmentation technique, to ease the distribution gap between the training and testing.
We further conduct extensive experiments on a variety of baseline models to reveal the performance improvement with FakeEdge to show its capability of alleviating the dataset shift issue on a broad range of benchmarks.
Related work
Dataset Shift.
Dataset shift is a fundamental issue in the world of machine learning.
Within the collection of dataset shift issues, there are several specific problems based on which part of the data experience the distributional shift, including covariate shift, concept shift, and prior probability shift.
[16] gives a rigorous definition about different dataset shift situations.
In the context of GNNs, [21] investigates the generalization ability of GNN models, and propose a self-supervised task to improve the size generalization.
[22] studies the problem that the node labels in training set are not uniformly sampled and suggests applying a regularizer to reduce the distributional gap between training and testing.
[23] proposes a risk minimization method by exploring multiple context of the observed graph to enable GNNs to generalize to out-of-distribution data.
[24] demonstrates that the existing link prediction models can fail to generalize to testing set with larger graphs and designs a structural pairwise embedding to achieve size stability.
[25–27] study the dataset shift problem for graph-level tasks, especially focusing on the graphs in the training and testing set with varying sizes.
Graph Data Augmentation.
Several data augmentation methods are introduced to modify the graph connectivity by adding or removing edges [28].
DropEdge [29] acts like a message passing reducer to tackle over-smoothing or overfitting problems [30].
Topping et al.
modify the graph’s topological structure by removing negatively curved edges to solve the bottleneck issue [32] of message passing [31].
GDC [33] applies graph diffusion methods on the observed graph to generate a diffused counterpart as the computation graph.
For the link prediction task, CFLP [34] generates counterfactual links to augment the original graph.
Edge Proposal Set [35] injects edges into the training graph, which are recognized by other link predictors in order to improve performance.
A proposed unified framework for link prediction
In this section, we formally introduce the link prediction task and formulate several existing GNNbased methods into a common general framework.
3.1 Preliminary
Let G =
( V, E, x V , x E ) be an undirected graph.
V is the set of nodes with size n , which can be indexed as { i } ni =1 .
E ⊆
V × V is the observed set of edges.
x Vi ∈ X V represents the feature of node i .
x Ei,j ∈ X E represents the feature of the edge ( i, j ) if ( i, j ) ∈ E .
The other unobserved set of edges is E c ⊆
V × V \ E , which are either missing or going to form in the future in the original graph G .
d ( i, j ) denotes the shortest path distance between node i and j .
The r -hop enclosing subgraph G ri,j for node i, j is the subgraph induced from G by node sets V ri,j = { v | v ∈ V, d ( v, i ) ≤ r or d ( v, j ) ≤ r } .
The edges set of G ri,j are E ri,j = { ( p, q ) | ( p, q )
∈ E and p, q ∈ V ri,j } .
An enclosing subgraph
G ri,j = ( V ri,j , E ri,j , x VV ri,j , x EE ri,j ) contains all the information in the neighborhood of node i, j .
The node set { i, j } is called the focal node pair , where we are interested in if there exists (observed) or should exist (unobserved) an edge between nodes i, j .
In the context of link prediction, we will use the term subgraph to denote enclosing subgraph in the following sections.
3.2 Subgraph link prediction
In this section, we discuss the definition of Subgraph Link Prediction and investigate how current link prediction methods can be unified in this framework.
We mainly focus on link prediction methods based on GNNs, which propagate the message to each node’s neighbors in order to learn the representation.
We start by giving the definition of the subgraph’s properties:
Definition 1.
Given a graph G =
( V, E, x V , x E ) and the unobserved edge set E c , a subgraph G ri,j have the following properties
: 1. a label y ∈ { 0 , 1 } of the subgraph indicates if there exists, or will form, an edge incident with focal node pair { i, j } .
That is, G ri,j label y
=
if and only if ( i, j ) ∈
E ∪ E c .
Otherwise, label y = 0 .
the existence e ∈ { 0 , 1 } of an edge in the subgraph indicates whether there is an edge observed at the focal node pair { i, j } .
If ( i, j ) ∈ E , e = 1 .
Otherwise e = 0 .
a phase c ∈ { train , test } denotes whether the subgraph belongs to training
or testing stage.
Especially for a positive subgraph ( y = 1 ), if ( i, j ) ∈ E , then c = train.
If ( i, j ) ∈ E c , then c = test.
Note that, the label y = 1 does not necessarily indicate the observation of the edge at the focal node pair { i, j } .
A subgraph in the testing set may have the label y
=
but the edge may not be present.
The existence e = 1 only when the edge is observed at the focal node pair.
Definition 2.
Given a subgraph G ri,j , Subgraph Link Prediction is a task to learn a feature h of the subgraph G ri,j and uses it to predict the label y ∈ { 0 , 1 } of the subgraph.
Generally, subgraph link prediction regards the link prediction task as a subgraph classification task.
The pipeline of subgraph link prediction starts with extracting the subgraph G ri,j around the focal node pair { i, j } , and then applies GNNs to encode the node representation Z .
The latent feature h of the subgraph is obtained by pooling methods on Z .
In the end, the subgraph feature h is fed into a classifier.
In summary, the whole pipeline entails:
Subgraph Extraction : Extract the subgraph G ri,j around the focal node pair { i, j } ; 2.
Node Representation Learning : Z =
GNN ( G ri,j ) , where Z ∈ R | V ri,j |× F hidden is the node embedding matrix learned by the GNN encoder; 3.
Pooling : h =
Pooling ( Z ; G ri,j ) , where h ∈ R F pooled is the latent feature of the subgraph G ri,j ; 4.
Classification : y =
Classifier ( h ) .
There are two main streams of GNN-based link prediction models.
Models like SEAL [11] and
WalkPool [12] can naturally fall into the subgraph link prediction framework, as they thoroughly follow the pipeline.
In SEAL, SortPooling [36] serves as a readout to aggregate the node’s features in the subgraph.
WalkPool designs a random-walk based pooling method to extract the subgraph feature h .
Both methods take advantage of the node’s representation from the entire subgraph.
In addition, there is another stream of link prediction models, such as GAE [10] and PLNLP [14], which learns the node representation and then devises a score function on the representation of the focal node pair to represent the likelihood of forming a link.
We find that these GNN-based methods with the message passing paradigm also belong to a subgraph link prediction task.
Considering a
GAE with l layers, each node v essentially learns its embedding from its l -hop neighbors { i | i ∈
V, d ( i, v ) ≤ l } .
The score function can be then regarded as a center pooling on the subgraph, which only aggregates the features of the focal node pair as h to represent the subgraph.
For a focal node pair { i, j } and GAE with l layers, an l -hop subgraph G li,j sufficiently contains all the information needed to learn the representation of nodes in the subgraph and score the focal node pair { i, j } .
Thus, the GNN-based models can also be seen as a citizen of subgraph link prediction.
In terms of the score function, there are plenty of options depending on the predictive power in practice.
In general, the common choices are: (1) Hadamard product: h
=
z i ◦ z j ; (2) MLP: h =
MLP ( z i ◦ z j ) where
MLP is the Multi-Layer Perceptron; (3) BiLinear: h
= z i W z
j where W is a learnable matrix; (4)
BiLinearMLP: h =
MLP ( z i ) ◦ MLP ( z j ) .
In addition to GNN-based methods, the concept of the subgraph link prediction can be extended to loworder heuristics link predictors, like Common Neighbor [1], Adamic–Adar index [20], Preferential
Attachment [37], Jaccard Index [38], and Resource Allocation [39].
The predictors with the order r can be computed by the subgraph G ri,j .
The scalar value can be seen as the latent feature h .
FakeEdge: Mitigates dataset shift in subgraph link prediction
In this section, we start by giving the definition of dataset shift in the general case, and then formally discuss how dataset shift occurs with regard to subgraph link prediction.
Then we propose FakeEdge as a graph augmentation technique to ease the distribution gap of the subgraph representation between the training and testing sets.
Lastly, we discuss how FakeEdge can enhance the expressive power of any GNN-based subgraph link prediction model.
4.1 Dataset shift
Definition 3.
Dataset Shift happens when the joint distribution between train and test is different.
That is, p ( h , y | c = train ) ̸
= p ( h , y | c = test ) .
A simple example of dataset shift is an object detection system.
If the system is only designed and trained under good weather conditions, it may fail to capture objects in bad weather.
In general, dataset shift is often caused by some unknown latent variable, like the weather condition in the example above.
The unknown variable is not observable during the training phase so the model cannot fully capture the conditions during testing.
Similarly, the edge existence e ∈ { 0 , 1 } in the subgraph poses as an "unknown" variable in the subgraph link prediction task.
Most of the current
GNN-based models neglect the effect of the edge existence on encoding the subgraph’s feature.
Definition 4.
A subgraph’s feature h is called Edge Invariant if p ( h , y | e ) =
p ( h , y ) .
To explain, the Edge Invariant subgraph embedding stays the same no matter if the edge is present at the focal node pair or not.
It disentangles the edge’s existence and the subgraph representation
[Figure]
learning.
For example, common neighbor predictor is Edge Invariant because the existence of an edge at the focal node pair will not affect the number of common neighbors that two nodes can have.
However, Preferential Attachment, another widely used heuristics link prediction predictor, is not
Edge Invariant because the node degree varies depending on the existence of the edge.
Theorem 1.
GNN cannot learn the subgraph feature h to be Edge Invariant.
Recall that the subgraphs in Figure 1 are encoded differently between the training and testing set because of the presence/absence of the focal link.
Thus, the vanilla GNN cannot learn the Edge
Invariant subgraph feature.
Learning Edge Invariant subgraph feature is crucial to mitigate the dataset shift problem.
Here, we give our main theorem about the issue in the link prediction task:
Theorem 2.
Given p ( h , y | e , c ) =
p ( h , y | e ) , there is no Dataset Shift in the link prediction if the subgraph embedding is Edge Invariant .
That is, p ( h , y | e ) =
p ( h , y )
=
⇒ p
( h , y | c ) =
p ( h , y ) .
The assumption p ( h , y | e , c ) = p ( h , y | e ) states that when the edge at the focal node pair is taken into consideration, the joint distribution keeps the same across the training and testing stages, which means that there is no other underlying unobserved latent variable shifting the distribution.
The theorem shows an Edge Invariant subgraph embedding will not cause a dataset shift phenomenon.
Theorem 2 gives us the motivation to design the subgraph embedding to be Edge Invariant.
When it comes to GNNs, the practical GNN is essentially a message passing neural network [17].
The existence of the edge incident at the focal node pair can determine the computational graph for message passing when learning the node representation.
4.2 Proposed methods
Having developed conditions of dataset shift phenomenon in link prediction, we next introduce a collection of subgraph augmentation techniques named as FakeEdge (Figure 2), which satisfies the conditions in Theorem 2.
The motivation is to mitigate the distribution shift of the subgraph embedding by eliminating the different patterns of target link existence between training and testing sets.
Note that all of the strategies follow the same discipline: align the topological structure around the focal node pair in the training and testing datasets, especially for the isomorphic subgraphs.
Therefore, we expect that it can gain comparable performance improvement across different strategies.
Compared to the vanilla GNN-based subgraph link prediction methods, FakeEdge augments the computation graph for node representation learning and subgraph pooling step to obtain an Edge
Invariant embedding for the entire subgraph.
Edge Plus A simple strategy is to always make the edge present at the focal node pair for all training and testing samples.
Namely, we add an edge into the edge set of subgraph by E r
+ i,j = E r
i,j ∪ { ( i, j ) } , and use this edge set to calculate the representation h plus of the subgraph G r + i,j .
Edge Minus Another straightforward modification is to remove the edge at the focal node pair if existing.
That is, we remove the edge from the edge set of subgraph by E r − i,j = E ri,j \{ ( i, j ) } , and obtain the representation h minus from G r − i,j .
For GNN-based models, adding or removing edges at the focal node pair can amplify or reduce message propagation along the subgraph.
It may also change the connectivity of the subgraph.
We are interested to see if it can be beneficial to take both situations into consideration by combining them.
Based on Edge Plus and Edge Minus , we further develop another two Edge Invariant methods:
Edge Mean To combine Edge Plus and Edge Minus , one can extract these two features and fuse them into one view.
One way is to take the average of the two latent features by h mean = h plus + h minus.
Edge Att Edge Mean weighs G r
+ i,j and G r − i,j equally on all subgraphs.
To vary the importance of two modified subgraphs, we can apply an adaptive weighted sum operation.
Similar to the practice in the text translation [40], we apply an attention mechanism to fuse the h plus and h minus by:
[Equation]
where w · = SoftMax ( q ⊺ · tanh ( W · h · + b ))
[Equation]
4.3 Expressive power of structural representation
[Figure]
[Equation]
[Figure]
In addition to solving the issue of dataset shift,
FakeEdge can tackle another problem that impedes the expressive power of link prediction methods on the structural representation [41].
In general, a powerful model is expected to discriminate most of the non-isomorphic focal node pairs.
For instance, in Figure 3 we have two isomorphic subgraphs A and B , which do not have any overlapping nodes.
Suppose that the focal node pairs we are interested in are { u, w } and { v, w } .
Obviously, those two focal node pairs have different structural roles in the graph, and we expect different structural representations for them.
With GNN-based methods like GAE, the node representation of the node u and v will be the same z u
=
z v , due to the fact that they have isomorphic neighborhoods.
GAE applies a score function on the focal node pair to pool the subgraph’s feature.
Hence, the structural representation of node sets { u, w } and { v, w } would be the same, leaving them inseparable in the embedding space.
This issue is caused by the limitation of GNNs, whose expressive power is bounded by 1-WL test [42].
Zhang et al.
address this problem by assigning distinct labels between the focal node pair and the rest of the nodes in the subgraph [19].
FakeEdge manages to resolve the issue by augmenting the neighborhoods of those two isomorphic nodes.
For instance, we can utilize the Edge Plus strategy to deliberately add an edge between nodes u and w (shown as the dashed line in Figure 3).
Note that the edge between v and w has already existed.
There is no need to add an edge between them.
Therefore, the node u and v will have different neighborhoods ( u has 4 neighbors and v has 3 neighbors), resulting in the different node representation between the node u and v after the first iteration of message propagation with GNN.
In the end, we can obtain different representations for two focal node pairs.
Other FakeEdge methods like Edge Minus can also tackle the issue in a similar way.
According to Theorem 2 in [19], such non-isomorphic focal node pairs { u, w } , { v, w } are not sporadic cases in a graph.
Given an n nodes graph whose node degree is O (log 1 − ϵ 2 r n ) for any constant ϵ > 0 , there exists ω ( n 2 ϵ ) pairs of such kind of { u, w } and { v, w } , which cannot be distinguished by GNN-based models like GAE.
However, FakeEdge can enhance the expressive power of link prediction methods by modifying the subgraph’s local connectivity.
[Table]
[Figure]
PLNLP
[Table]
[Figure] [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table] [Figure] [Table] [Figure]
[Table]
[Figure] [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
Experiments
In this section
[Table]
we conduct extensive experiments to evaluate how FakeEdge can mitigate the dataset shift issue on various baseline models in the link prediction task.
Then we empirically show the distribution gap of the subgraph representation between the training and testing and discuss how the dataset shift issue can worsen with deeper GNNs.
The code for the experiment can be found at https://github.com/Barcavin/FakeEdge .
5.1 Experimental setup
Baseline methods.
We show how FakeEdge techniques can improve the existing link prediction methods, including GAE-like models [10], PLNLP [14], SEAL [11], and WalkPool [12].
To examine the effectiveness of FakeEdge, we compare the model performance with subgraph representation learned on the original unmodified subgraph and the FakeEdge augmented ones.
For GAE-like models, we apply different GNN encoders, including GCN [9], SAGE [43] and GIN [42].
SEAL and
WalkPool have already been implemented in the fashion of the subgraph link prediction.
However, a subgraph extraction preprocessing is needed for GAE and PLNLP, since they are not initially implemented as the subgraph link prediction.
GCN, SAGE, and PLNLP use a score function to pool the subgraph.
GCN and SAGE use the Hadamard product as the score function, while MLP is applied for PLNLP (see Section 3.2 for discussions about the score function).
Moreover, GIN applies a subgraph-level pooling strategy, called "mean readout" [42], whose pooling is based on the entire subgraph.
Similarly, SEAL and WalkPool also utilize the pooling on the entire subgraph to aggregate the representation.
More details about the model implementation can be found in Appendix D.
Benchmark datasets.
For the experiment, we use 3 datasets with node attributes and 8 without attributes.
The graph datasets with node attributes are three citation networks: Cora [44], Citeseer [45], and Pubmed [46].
The graph datasets without node attributes are eight graphs in a variety of domains: USAir [47], NS [48], PB [49], Yeast [50], C.ele [51], Power [51], Router [52], and
E.coli [53].
More details about the benchmark datasets can be found in Appendix E.
Evaluation protocols.
Following the same experimental setting as of [11, 12], the links are split into 3 parts: 85% for training, 5% for validation, and 10% for testing.
The links in validation and testing are unobserved during the training phase.
We also implement a universal data pipeline for different methods to eliminate the data perturbation caused by train/test split.
We perform 10 random data splits to reduce the performance disturbance.
Area under the curve (AUC) [54] is used as the evaluation metrics and is reported by the epoch with the highest score on the validation set.
5.2 Results
FakeEdge on GAE-like models.
The results of models with ( Edge Plus , Edge Minus , Edge Mean , and Edge Att ) and without ( Original ) FakeEdge are shown in Table 1.
We observe that FakeEdge is a vital component for all different methods.
With FakeEdge, the link prediction model can obtain a significant performance improvement on all datasets.
GAE-like models and PLNLP achieve the most remarkable performance improvement when FakeEdge alleviates the dataset shift issue.
FakeEdge boosts them by 2% - 11% on different datasets.
GCN, SAGE, and PLNLP all have a score function as the pooling methods, which is solely based on the focal node pair.
In particular, the focal node pair is incident with the target link, which determines how the message passes around it.
Therefore, the most severe dataset shift issues happen at the embedding of the focal node pair during the node representation learning step.
FakeEdge is expected to bring a notable improvement to these situations.
Encoder matters.
In addition, the choice of encoder plays an important role when GAE is deployed on the Original subgraph.
We can see that SAGE shows the best performance without FakeEdge among these 3 encoders.
However, after applying FakeEdge, all GAE-like methods achieve comparable better results regardless of the choice of the encoder.
We come to a hypothesis that the plain SAGE itself leverages the idea of FakeEdge to partially mitigate the dataset shift issue.
Each node’s neighborhood in SAGE is a fixed-size set of nodes, which is uniformly sampled from the full neighborhood set.
Thus, when learning the node representation of the focal node pair in the positive training sets, it is possible that one node of the focal node pair is not selected as the neighbor of the other node during the neighborhood sampling stage.
In this case, the FakeEdge technique Edge
Minus is applied to modify such a subgraph.
FakeEdge on subgraph-based models.
In terms of SEAL and WalkPool, FakeEdge can still robustly enhance the model performance across different datasets.
Especially for datasets like Power and Router, FakeEdge increases the AUC by over 10% on both methods.
Both methods achieve better results across different datasets, except WalkPool model on datasets Cora and Citeseer.
One of the crucial components of WalkPool is the walk-based pooling method, which actually operates on both the Edge Plus and Edge Minus graphs.
Different from FakeEdge technique, WalkPool tackles the dataset shift problem mainly on the subgraph pooling stage.
Thus, WalkPool shows similar model performance between the Original and FakeEdge augmented graphs.
Moreover, SEAL and WalkPool have utilized one of the FakeEdge techniques as a trick in their initial implementations.
However, they have failed to explicitly point out the fix of dataset shift issue from such a trick in their papers.
Different FakeEdge techniques.
When comparing different FakeEdge techniques, Edge Att appears to be the most stable, with a slightly better overall performance and a smaller variance.
However, there is no significant difference between these techniques.
This observation is consistent with our expectation since all FakeEdge techniques follow the same discipline to fix the dataset shift issue.
5.3 Further discussions
In this section, we conduct experiments to more thoroughly study why FakeEdge can improve the performance of the link prediction methods.
We first give an empirical experiment to show how severe the distribution gap can be between training and testing.
Then, we discuss the dataset shift issue with deeper GNNs.
Last but not the least, we explore how FakeEdge can even improve the performance of heuristics predictors.
5.3.1 Distribution gap between the training and testing
FakeEdge aims to produce Edge Invariant subgraph embedding during the training and testing phases in the link prediction task, especially for those positive samples p ( h | y = 1) .
That is, the subgraph representation of the positive samples between the training and testing should be difficult, if at all, to be distinguishable from each other.
Formally, we ask whether p ( h | y = 1 , c = train ) =
p ( h | y = 1 , c = test ) , by conducting an empirical experiment on the subgraph embedding.
We retrieve the subgraph embedding of the positive samples from both the training and testing stages, and randomly shuffle the embedding.
Then we classify whether the sample is from training (c = train)
Cora
Citeseer
Pubmed
USAir
NS
PB
Yeast
C.ele
Power
Router
E.coli 0.
Original Plus
Minus Mean
Att
Dataset
[Figure]
Distribution gap exists ⬆
Layers
[Table]
or testing (c = test).
The shuffled positive samples are split 80%/20% as train and inference sets.
Note that the train set here, as well as the inference set, contains both the shuffled positive samples from the training and testing set in the link prediction task.
Then we feed the subgraph embedding into a 2-layer MLP classifier to investigate whether the classifier can differentiate the training samples (c = train) and the testing samples (c = test).
In general, the classifier will struggle to undertake the classification if the embedding of training and testing samples is drawn from the same underlying distribution, which indicates there is no significant dataset shift issue.
We use GAE with the GCN as the encoder to run the experiment.
AUC is used to measure the discriminating power of the classifier.
The results are shown in Figure 4.
Without FakeEdge, the classifier shows a significant ability to separate positive samples between training and testing.
When it comes to the subgraph embedding with FakeEdge, the classifier stumbles in distinguishing the samples.
The comparison clearly reveals how different the subgraph embedding can be between the training and testing, while FakeEdge can both provably and empirically diminish the distribution gap.
5.3.2 Dataset shift with deeper GNNs
Given two graphs with n nodes in each graph, 1-WL test may take up to n iterations to determine whether two graphs are isomorphic [55].
Thus, GNNs, which mimic 1-WL test, tend to discriminate more non-isomorphic graphs when the number of GNN layers increases.
SEAL [19] has empirically witnessed a stronger representation power and obtained more expressive link representation with deeper GNNs.
However, we notice that the dataset shift issue in the subgraph link prediction becomes more severe when GNNs try to capture long-range information with more layers.
We reproduce the experiments on GIN by using l = 1 , 2 , 3 message passing layers and compare the model performance by AUC scores with and without FakeEdge.
Here we only apply Edge Att as the FakeEdge technique.
The relative AUC score improvement of Edge Att is reported, namely ( AUC EdgeAtt − AUC Original ) /AUC
Original .
The results are shown in Table 2.
As we can observe, the relative performance improvement between Edge Att and Original becomes more significant with more layers, which indicates that the dataset shift issue can be potentially more critical when we seek deeper GNNs for greater predictive power.
To explain such a phenomenon, we hypothesize that GNNs with more layers will involve more nodes in the subgraph, such that their computation graph is dependent on the existence of the edge at the focal node pair.
For example, select a node v from the subgraph G ri,j , which is at least l hops away from the focal node pair { i, j } , namely l = min ( d ( i, v ) , d ( j, v )) .
If the GNN has only l layers, v will not include the edge ( i, j ) in its computation graph.
But with a GNN with l + 1 layers, the edge ( i, j ) will affect v ’s computation graph.
We leave the validation of the hypothesis to future work.
[Table]
5.3.3 Heuristic methods with FakeEdge
FakeEdge, as a model-agnostic technique, not only has the capability of alleviating the dataset shift issue for GNN-based models, but also can tackle the problem for heuristics methods.
The heuristics link predictors assign a score to each focal node pair, indicating the likelihood of forming a new edge.
Some of the conventional heuristic link predictors, like Common Neighbor [1], Adamic–Adar index [20], or Resource Allocation [39], are Edge Invariant because these predictors are independent of the existence of the target link.
However, other link predictors, including Preferential Attachment (PA) [37] and Jaccard Index (Jac) [38], are not Edge Invariant.
The existence/absence of the target link can change the values of the predictors, which in turn changes the ranking of focal node pairs.
The original PA for a focal node pair i, j is PA ( i, j )
= |N ( i ) ||N ( j ) | , where N ( i ) is the neighbors of node i . After applying Edge
Plus , PA plus ( i, j ) = |N ( i ) ∪ { j }||N ( j ) ∪ { i }| . Similarly, Jac plus ( i, j ) = |N ( i ) ∩ N ( j ) | /
|N ( i ) ∪ N ( j ) ∪ { i, j }| . We follow the same protocol in the previous experiment. As shown in Table 3, Edge Plus can significantly improve the performance of the PA predictor on several datasets.
With FakeEdge, PA performs over 10% better on Pubmed.
Surprisingly, even though PA is not able to predict the links on
Router dataset with AUC score lower than 50%, PA with Edge Plus achieves 74% AUC score and becomes a functional link predictor.
In terms of Jac, we observe that Jac with FakeEdge can only gain marginal improvement.
This is because that even though Jac is dependent on the existence of target link, the change of Jac index is relatively small when the existence of the target link flips.
Conclusion
Dataset shift is arguably one of the most challenging problems in the world of machine learning.
However, to the best of our knowledge, none of the previous studies sheds light on this notable phenomenon in link prediction.
In this paper, we studied the issue of dataset shift in link prediction tasks with GNN-based models.
We first unified several existing models into a framework of subgraph link prediction.
Then, we theoretically investigated the phenomenon of dataset shift in subgraph link prediction and proposed a model-agnostic technique FakeEdge to amend the issue.
Experiments with different models over a wide range of datasets verified the effectiveness of FakeEdge.
A Related work of link predictions
Early studies on link prediction problems mainly focus on heuristics methods, which require expertise on the underlying trait of network or hand-crafted features, including Common Neighbor [1],
Adamic–Adar index [20] and Preferential Attachment [37], etc.
WLNM [56] suggests a method to encode the induced subgraph of the target link as an adjacency matrix to represent the link.
With the huge success of GNN [9], GNN-based link prediction methods have become dominant across different areas.
Graph Auto Encoder(GAE) and Variational Graph Auto Encoder(VGAE)
[10] perform link prediction tasks by reconstructing the graph structure.
SEAL [11] and DE [13] propose methods to label the nodes according to the distance to the focal node pair.
To better exploit the structural motifs [57] in distinct graphs, a walk-based pooling method (WalkPool)
[12] is designed to extract the representation of the local neighborhood.
PLNLP [14] sheds light on pairwise learning to rank the node pairs of interest.
Based on two-dimensional Weisfeiler-Lehman tests, Hu et al.
propose a link prediction method that can directly obtain node pair representation [58].
B Proof of Theorem
We restate the Theorem 1: GNN cannot learn the subgraph feature h to be Edge Invariant.
Proof.
Recall that the computation of subgraph feature h involves steps such as:
Subgraph Extraction : Extract the subgraph G ri,j around the focal node pair { i, j } ; 2.
Node Representation Learning : Z =
GNN ( G ri,j ) , where Z ∈ R | V ri,j |× F hidden is the node embedding matrix learned by the GNN encoder; 3.
Pooling : h =
Pooling ( Z ; G ri,j ) , where h ∈ R F pooled is the latent feature of the subgraph G ri,j ;
Here, GNN is Message Passing Neural Network [17].
Given a subgraph G =
( V, E, x V , x E ) , GNN with T layers applies following rules to update the representation of node i ∈ V :
[Equation]
where N ( i ) is the neighborhood of node i in G , M t is the message passing function at layer t and U t is the node update function at layer t .
The hidden states at the first layer are set as h (0) i
= x
Vi .
The hidden states at the last layer are the outputs Z i
= h ( T ) i .
Given any subgraph G r i,j = ( V r i,j , E r
i,j , x V V r
i,j , x E E r i,j ) with the edge present at the focal node pair ( i,
j ) ∈
E ri,j , we construct another isomorphic subgraph G r ¯ i, ¯ j = ( V r ¯ i, ¯ j , E r ¯ i, ¯ j , x VV r ¯ i, ¯ j , x EE r ¯ i, ¯ j ) , but remove the edge (¯ i, ¯ j ) from the edge set E r ¯ i, ¯ j of the subgraph.
G r ¯ i, ¯ j can be seen as the counterpart of G ri,j in the testing set.
Thus, for the first iteration of node updates t = 1 :
[Equation]
Note that N (¯ i ) ∪ { j } = N ( i ) .
We have:
[Equation]
As U t is injective, p ( h (1) i , y = 1 | e = 1) ̸
= p ( h (1)¯ i , y = 1) =
p ( h (1) i , y = 1 | e = 0) .
Similarly, we can conclude that p ( h ( T ) i , y = 1 | e = 1) ̸
= p ( h ( T ) i , y = 1 | e = 0) .
As we use the last iteration of node updates h ( T ) i as the final node representation
Z , we have p ( Z , y | e = 1) ̸
= p ( Z , y | e = 0) , which leads to p ( h , y | e = 1) ̸
= p ( h , y | e = 0) and concludes the proof.
C Proof of Theorem
We restate the Theorem 2:
Given p ( h , y | e , c ) =
p ( h , y | e ) , there is no Dataset Shift in the link prediction if the subgraph embedding is Edge Invariant .
That is, p ( h , y | e ) =
p ( h , y )
=
⇒ p
( h , y | c ) =
p ( h , y ) .
Proof.
[Equation]
D Details about the baseline methods
To verify the effectiveness of FakeEdge, we tend to introduce minimal modification to the baseline models and make them compatible with FakeEdge techniques.
The baseline models in our experiments are mainly from the two streams of link prediction models.
One is the GAE-like model, including
GCN [9], SAGE [43], GIN [42] and PLNLP [14].
The other includes SEAL [11] and WalkPool [12].
GCN, SAGE and PLNLP learn the node representation and apply a score function on the focal node pair to represent the link.
As GAE-like models are not implemented in the fashion of subgraph link prediction, the subgraph extraction step is necessary for them as preprocessing.
We follow the code from the labeling trick [19], which implements the GAE models as the subgraph link prediction task.
In particular, GIN concatenates the node embedding from different layers to learn the node representation and applies a subgraph-level readout to aggregate as the subgraph representation.
As suggested by [13, 19], we always inject the distance information with the Double-Radius Node
Labeling [11] (DRNL) to enhance the model performance of GAE-like models.
In terms of the selection of hyperparameters, we use the same configuration as [19] on datasets Cora, Citeseer and
Pubmed.
As they do not have experiments on other 8 networks without attributes, we set the subgraph hop number as 2 and leave the rest of them as default.
For PLNLP, we also add a subgraph extraction step without modifying the core part of the pairwise learning strategy.
We find that the performance of PLNLP under subgraph setting is very unstable on different train/test splits.
In particular, the performance’s standard deviation of PLNLP is over 10% on each experiment.
Therefore, we also apply DRNL to stabilize the model.
SEAL and WalkPool have applied one of the FakeEdge techniques in their initial implementation.
SEAL uses a Edge Minus strategy to remove all the edges at focal node pair as a preprocessing step, while WalkPool applies Edge Plus to always inject edges into the subgraph for node representation learning.
Additionally, WalkPool has the walk-based pooling method operating on both the Edge
Plus and Edge Minus graphs.
This design is kept in our experiment.
Thus, our FakeEdge technique only takes effect on the node representation step for WalkPool.
From the results in Section 5.2, we can conclude that the dataset shift issue on the node representation solely would significantly impact the model performance.
We also use the same hyperparameter settings as originally reported in their paper.
E Benchmark dataset descriptions
The graph datasets with node attributes are three citation networks: Cora [44], Citeseer [45] and
Pubmed [46].
Nodes represent publications and edges represent citation links.
The graph datasets without node attributes are: (1) USAir [47]: a graph of US Air lines; (2) NS [48]: a collaboration network of network science researchers; (3) PB [49]: a graph of links between web pages on US political topic;
(4) Yeast [50]: a protein-protein interaction network in yeast; (5) C.ele
[51]: the neural network of Caenorhabditis elegans; (6) Power [51]: the network of the western US´s electric grid; (7) Router [52]: the Internet connection at the router-level; (8) E.coli [53]: the reaction network of metabolites in Escherichia coli.
The detailed statistics of the datasets can be found in Table 4.
[Table] [Table]
Models
Cora
Citeseer
Pubmed
USAir
NS
PB
Yeast
C.
[Table]
F Results measured by Hits@20 and statistical significance of results
We adopt another widely used metrics in the link prediction task [59], Hits@20, to evaluate the model performance with and without FakeEdge.
The results are shown in Table 5.
FakeEdge can boost all the models predictive power on different datasets.
Note that the AUC scores on several datasets are almost saturated in Table 1.
To further verify the statistical significance of the improvement, a two-sided t -test is conducted with the null hypothesis that the augmented Edge Att and the Original representation learning would reach at the identical average scores.
The p -values of different methods can be found in Table 6.
Recall that the p -value smaller than 0 .
05 is considered as statistically significant.
GAE-like methods obtain significant improvement on almost all of the datasets, except GCN on C.ele.
SEAL shows significant improvement with Edge
Att on 7 out of 11 datasets.
For WalkPool, more than half of the datasets are significantly better.
G FakeEdge with extremely sparse graphs
In real applications, the size of testing set often outnumbers the training set.
When it happens to a link prediction task, the graph will become more sparse because of the huge number of unseen links.
We are interested to see how FakeEdge can handle situations where the ratio of training set is low and there exists a lot of “true” links missing in the training graph.
We reset the train/test split as 20% for training, 30% for validation and 50% for testing and reevaluate the model performance.
The results can be found in Table 7.
As shown in the table, FakeEdge can still consistently improve the model performance under such an extreme setting.
It shows that the
[Table]
dataset shift for link prediction is a common issue and FakeEdge has the strength to alleviate it in various settings
[Table]
However, we still observe a significant performance drop when compared to the 85/5/10 evaluation setting.
This degradation may be caused by a more fundamental dataset shift problem of link prediction: the nodes in a graph are not sampled independently.
Existing link prediction models often assume that the likelihood of forming a link relies on its local neighborhood.
Nevertheless, an intentionally-sparsified graph can contain a lot of missing links from the testing set, leading to corrupted local neighborhoods of links which cannot reflect the real environments surrounding.
FakeEdge does not have the potential to alleviate such a dataset shift.
We leave this as a future work.
H Concatenation as another valid Edge Invariant subgraph embedding
Edge Concat To fuse the feature from Edge Plus and Edge Minus , another simple and intuitive way is to concatenate two embedding into one representation.
Namely, h concat =
[ h plus ; h minus ] , where [ · ; · ] is the concatenation operation.
h concat is also an Edge Invariant subgraph embedding.
In Table 8, we observe that Edge Concat has the similar performance improvement like other FakeEdge methods on all different backbone models.
I Dataset shift vs expressiveness: which contributes more with FakeEdge?
In Section 4.3, we discussed how FakeEdge can enhance the expressive power of GNN-based models on non-isomorphic focal node pairs.
Meanwhile, we have witnessed the boost of model performance brought by FakeEdge in the experiments.
One natural question to ask is whether resolving the dataset shift issue or lifting up expressiveness is the major contributor to make the model perform better.
To answer the question, we first revisit the condition of achieving greater expressiveness.
FakeEdge will lift up the expressive power when there exists two nodes being isomorphic in the graph, where we can construct a pair of non-isomorphic focal node pairs which GNNs cannot distinguish.
Therefore, how often such isomorphic nodes exist in a graph will determine how much improvement FakeEdge can make by bringing greater expressiveness.
Even though isomorphic nodes are common in specific types of graphs like regular graphs, it can be rare in the real world datasets [60].
Thus, we tend to conclude that the effect of solving dataset shift issue by FakeEdge contributes more to the performance improvement rather than greater expressive power.
But fully answering the question needs a further rigorous study.
[Table]
J
[Table]
FakeEdge can align the embedding of isomorphic subgraphs in training and testing sets.
However, it can pose a limitation that hinders one aspect of the GNN-based model’s expressive power.
Figuregives an example where subgraphs are from training and testing phases, respectively.
Now consider that those two subgraphs are both from training set (c = train).
Still, the top subgraph has edge observed at focal node pair (y = 1 ), while the other does not (y = 0 ).
With FakeEdge, two subgraphs will be modified to be isomorphic, yielding the same representation.
However, they are non-isomorphic before the modification.
To the best of our knowledge, no existing method can simultaneously achieve the most expressive power and get rid of dataset shift issue, because the edge at the focal node pair in the testing set can never be observed under a practical problem setting.