Shift-Robust GNNs: Overcoming the Limitations of Localized Graph Training data

1 Introduction . The goal of graph-based semi-supervised learning ( SSL ) is to use relationships between data ( its graph inductive bias ) , along with a small set of labeled items , to predict the labels for the rest of a dataset . Unsurprisingly , varying exactly which nodes are labeled can have a profound effect on the generalization capability of a SSL classifier . Any bias in the sampling process to select nodes for training can create distributional differences between the training set and the rest of the graph . During inference any portion of the graph can be used , so any uneven labeling for training data can cause training and test data to have different distributions . An SSL classifier may then overfit to training data irregularities , thus hurting the performance at inference time . Recently , GNNs have emerged as a way to combine graph structure with deep neural networks . Surprisingly , most work on semi-supervised learning using GNNs for node classification [ 15 , 11 , 1 ] have ignored this critical problem , and even the most recently proposed GNN benchmarks [ 12 ] assume that an independent and identically distributed ( IID ) sample is possible for training labels . 1Code and processed data are available at https : //github.com/GentleZhu/Shift-Robust-GNNs . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) , Online . This problem of biased training labels can be quite pronounced when GNNs are applied for semisupervised learning in practice . It commonly happens when the size of the dataset is so large that only a subset of it can afford to be labeled – the exact situation where graph-based SSL is supposed to have a value proposition ! While the specific source of bias can vary , we have encountered it in many different settings . For example , sometimes fixed heuristics are used to select a subset of data ( which shares some characteristics ) for labeling . Other times , human analysts individually choose data items for labeling , using complex domain knowledge . However , even this can be rooted in shared characteristics of data . In yet another scenario , a label source may have some latency , causing a temporal mismatch between the distribution of data at time of labeling and at the time of inference . In all of these cases , the core problem is that the GNN overfits to spurious regularities as the subset of labeled data could not be created in an IID manner . One particular area where this can apply is in the spam and abuse domain , a common area of application for GNNs [ 18 , 10 , 33 ] . However , the labels in these problems usually come from explicit human annotations , which are both sparse ( as human labelling is expensive ) , and also frequently biased . Since spam and abuse problems typically have very imbalanced label distributions ( e.g . in many problems there are relatively few abusers – typically less than 1:100 ) , labeling nodes IID results in discovering very few abusive labels . In this case choosing the points to request labels for in an IID manner is simply not a feasible option if one wants to have a reasonable number of data items from the rare class . In this paper , we seek to quantify and address the problem of localized training data in Graph Neural Networks . We frame the problem as that of transfer learning – seeking to transfer the model ’ s performance from a small biased portion of the graph to the entire graph itself . Our proposed framework for addressing this problem , Shift-Robust GNN ( SR-GNN ) , strives to adapt a biased sample of labeled nodes to more closely conform to the distributional characteristics present in an IID sample of the graph . It can handle two kinds of bias that occur in both deeper GNNs and more recent linearized ( shallow ) versions of these models . First we consider the case of addressing distributional shift for standard GNN models such as GCNs [ 15 ] , MPNNs [ 7 ] , and many more [ 5 ] . These models create deep networks which iteratively convolve information over graph structure . SR-GNN addresses this variety of distributional shift via a regularization over the hidden layers of the network . Second , we consider a class of linearized models ( APPNP [ 16 ] , SimpleGCN [ 34 ] , etc ) which decouple GNNs into non-linear feature encoding and linear message passing . These models present an interesting challenge for debiasing , as the graph can introduce bias over the features after all learnable layers . In cases like this , SR-GNN can use an instance reweighting paradigm to ensure that the training examples are as representative as possible over the graph data . We illustrate the effectiveness of our proposed method on both paradigms with an experimental framework that introduces bias to the train/test split in a GNN , which lets us simulate the ‘ localized discovery ’ pattern observed in real applications on fully-labeled academic datasets . With these experiments we show that that our method SR-GNN can recover at least 40 % of the performance lost when training a GCN on the same biased input . Specifically , our contributions are the following : 1 . We provide the first focused discussion on the distributional shift problem in GNNs . 2 . We propose generalized framework , Shift-Robust GNN ( SR-GNN ) , which can address shift in both shallow and deep GNNs . 3 . We create an experimental framework which allows for creating biased train/test sets for graph learning datasets . 4 . We run extensive experiments and analyze the results , proving that our methods can mitigate distributional shift . 2 Related Work . 2.1 Distributional shift and Domain adaption work . Standard learning theory ( i.e . PAC , empirical risk minimization , etc . ) assumes that training and inference data is drawn from the same distribution , but there are many practical cases where this does not hold . The question of dealing with different distributions has been widely explored as a part of the transfer learning literature . In transfer learning , the domain adaptation problem deals with transferring knowledge from the source domain ( used for learning ) to the target domain ( the ultimate inference distribution ) . One of the first theoretical works on domain adaptation [ 3 ] developed a distance function between a model ’ s performance on the source and target domains to describe how similar they are . To obtain a final model , training then happens on a reweighed combination of source and target data , where weights are a function of the domain ’ s distance . Much additional theoretical ( e.g . [ 22 ] ) and practical work expanded this idea and explored models which are co-trained on both source and target data . These models seek to optimize utility while minimizing the distance between extracted features distributions on both domains ; this in turn led to the field of Domain Invariant Representation Learning ( DIR ) [ 6 ] . DIR is commonly achieved via co-training on labeled source and ( unlabeled ) target data . A modification to the loss either uses an adversarial head or adds additional regularizations . More recently , various regularizations using discrepancy measures have been shown to be more stable to hyperparameters and result in better performance than adversarial heads , faster loss convergence and easier training . Maximum mean discrepancy ( MMD ) [ 19 , 20 ] is a metric that measures difference between means of distributions in some rich Hilbert kernel space . Central moment discrepancy ( CMD ) [ 38 ] extends this idea and matches means and higher order moments in the original space ( without the projection into the kernel ) . CMD has been shown to produce superior results and is less susceptible to the weight with which CMD regularization is added to the loss [ 20 ] . It is important to point out that MMD and CMD regularizations are commonly used with non-linear networks on some hidden layer ( e.g . on extracted features ) . For linear models or non-differentiable models , prior work for domain adaptation often employed importance-reweighting instead . To find the appropriate weights , the same MMD distance was often used : in kernel mean matching ( KMM ) to find the appropriate weights one essentially minimizes MMD distance w.r.t the instance weights [ 17 ] . 2.2 GNNs . Given a graph G = { V , E , X } , the nodes V are associated with their features X ( X ∈ R|V |×F ) and the set of edges E ( i.e . adjacency matrix A , A ∈ R|V |×|V | ) which form connections between them . Graph neural networks [ 15 ] are neural networks that operate on both node features and graph structures . The core assumption of GNNs is that the structure of data ( A ) can provide a useful inductive bias for many modeling problems . GNNs output node representations Z which are used for unsupervised [ 30 ] or semi-supervised [ 15 , 11 ] learning . We denote the labels for SSL as { yi } . The general architecture Φ of a GNN consists of K neural network layers which encode the nodes and their neighborhood information using some learnable weights θ . More specifically the output of layer k of a GNN contains a row ( hki ) which can be used as a representation for each node i , i.e . zki = h k i . Successive layers mix the node representations using graph information , for example : Hk = σ ( ÃHk−1θk ) ( 1 ) where Ã is an appropriately normalized adjacency matrix2 , σ is the activation function , and H0 = X . For brevity ’ s sake , we refer to the final latent representations ZK simply as Z throughout the work . Although no existing work studies the distributional shift problem in GNNs , transfer learning of GNNs [ 13 , 39 ] has explored different node-level and graph-level pre-training tasks across different graphs . Alternatively , domain adaption methods have been used to optimize a domain classifier between source and target graphs [ 21 , 35 ] . In addition , distribution discrepancy minimization ( MMD ) has been adopted to train network embedding across domains [ 28 ] . Other regularizations for the latent state of GNNs have been proposed for domains like fairness [ 23 ] . We are the first to notice the importance of distributional shift on the same graph in a realistic setting and analyze its influence on different kinds of GNNs . 2The exact use of the adjacency matrix varies with specific GNN methods . For instance , the GCN [ 15 ] performs mean-pooling using matrix multiplication : Ã = D− 1 2 ( A + I ) D− 1 2 where I is the identity matrix , and D is the degree matrix of ( A+ I ) , Dii = ∑ j ( A+ I ) ij . 3 Distributional shift in GNNs . To learn an SSL classifier , a cross-entropy loss function l is commonly used , L = 1 M M∑ i=1 l ( yi , zi ) , ( 2 ) where zi is the node representation for the node i learned from a graph neural network , and M is the number of training examples . When the training and testing data come from the same domain ( i.e . Prtrain ( X , Y ) = Prtest ( X , Y ) ) , optimizing the cross entropy loss over the training data ensures that a classifier is well-calibrated for performing inference on the testing data .

Summary: Graph Neural Networks (GNNs) have been attracted much attention very recently because of surprising results in many applications. In general, GNNs assume that the training node labels follow the i.i.d. property. However, in practice, the labels could be very biased. To deal with the biased training nodes, the authors propose regularization methods that consider feature distance between training nodes and i.i.d sampled unlabeled nodes. CMD and KMM are key components of the regularization: CMD measures the distance between train nodes and general nodes, and KMM computes balanced weights of training loss functions. SR-GNN is the main framework that includes both CMD and KMM ideas. The experimental section shows that the biased data can hurt the performance, and SR-GNN can mitigate performance degradation.