Handling Distribution Shifts on Graphs: An Invariance Perspective

An input graph $G=(A,X)$ contains two-fold information²²2Our formulation and method can be trivially extended to cover edge features which we omit here for brevity.: an adjacency matrix $A=\{a_{vu}|v,u\in V\}$ and node features $X=\{x_{v}|v\in V\}$ where $V$ denotes node set. Apart from these, each node in the graph has a label, which can be represented as a vector $Y=\{y_{v}|v\in V\}$. We define $\mathbf{G}$ as a random variable of input graphs and $\mathbf{Y}$ as a random variable of node label vectors. Such a definition takes a global view and treat the input graph as a whole. Based on this, one can adapt the definition of general OOD problem Eq. 1 via instantiating the input as $\mathbf{G}$ and the target as $\mathbf{Y}$, and then the data generation can be characterized as $p(\mathbf{G},\mathbf{Y}|\mathbf{e})=p(\mathbf{G}|\mathbf{e})p(\mathbf{Y}|\mathbf{G},\mathbf{e})$ where $\mathbf{e}$ is a random variable of environments that is a latent variable and impacts data distribution.

However, the above definition makes little sense in node-level problems where in most cases there is a single input graph that contains a massive number of nodes. To make the problem-solving reasonable, we instead take a local view and investigate each node’s ego-graph that has influence on the centered node. Assume $\mathbf{v}$ as a random variable of nodes. We define node $v$’s $L$-hop neighbors as $N_{v}$ (where $L$ is an arbitrary integer) and the nodes in $N_{v}$ form an ego-graph $G_{v}$ which consists of a (local) node feature matrix $X_{v}=\{x_{u}|u\in N_{v}\}$ and a (local) adjacency matrix $A_{v}=\{a_{uw}|u,w\in N_{v}\}$. Use $\mathbf{G}_{\mathbf{v}}$ as a random variable of ego-graphs³³3We use a subscript $\mathbf{v}$ here to remind that it is an ego-graph from the view of a target node. whose realization is $G_{v}=(A_{v},X_{v})$. Besides, we define $\mathbf{y}$ as a random variable of node labels. In this way, we can fragment a whole graph as a set of instances $\{(G_{v},y_{v})\}_{v\in V}$ where $G_{v}$ denotes an input and $y_{v}$ is a target. Notice that the ego-graph can be seen as a Markov blanket for the centered node, so the conditional distribution $p(\mathbf{Y}|\mathbf{G},\mathbf{e})$ can be decomposed as a product of $|V|$ independent and identical marginal distributions $p(\mathbf{y}|\mathbf{G}_{\mathbf{v}},\mathbf{e})$.

Therefore, the data generation of $\{(G_{v},y_{v})\}_{v\in V}$ from a distribution $p(\mathbf{G},\mathbf{Y}|\mathbf{e})$ can be considered as a two-step procedure: 1) the entire input graph is generated via $G\sim p(\mathbf{G}|\mathbf{e})$ which can then be fragmented into a set of ego-graphs $\{G_{v}\}_{v\in V}$; 2) each node’s label is generated via $y\sim p(\mathbf{y}|\mathbf{G}_{\mathbf{v}}=G_{v},\mathbf{e})$. Then the OOD node-level prediction problem can be formulated as: given training data $\{G_{v},y_{v}\}_{v\in V}$ from $p(\mathbf{G},\mathbf{Y}|\mathbf{e}=e)$, the model needs to handle testing data $\{G_{v},y_{v}\}_{v\in V^{\prime}}$ from a new distribution $p(\mathbf{G},\mathbf{Y}|\mathbf{e}=e^{\prime})$. Denote $\mathcal{E}$ as the support of environments, $f$ as a predictor model with $\hat{y}=f(G_{v})$ and $l(\cdot,\cdot)$ as a loss function. More formally, the OOD problem can be written as:

We remark that the first-step sampling $G\sim p(\mathbf{G}|\mathbf{e}=e)$ can be ignored since in most cases one only has a single input graph in the context of node-level prediction tasks.

To solve the OOD problem Eq. 2 is impossible without any prior domain knowledge or structural assumptions since one only has access to data from limited environments in the training set. Recent studies (Rojas-Carulla et al., 2018; Arjovsky et al., 2019) propose to learn invariant predictor models which resorts to an assumption for data-generating process: the input instance contains a portion of features (i.e., invariant features) that 1) contributes to sufficient predictive information for the target and 2) gives rise to equally (optimal) performance of the downstream classifier across environments.

A more intuitive illustration for the above computation is presented in Fig. 1(b). The node-level readout $r_{v}$ aggregates the information from neighbored nodes recursively along the structures of BFS tree given by $G_{v}$. Essentially, the above definition assumes that in each layer the neighbored nodes contain a portion of causal features that contribute to stable prediction for $\mathbf{y}$ across different $\mathbf{e}$. Such a definition possesses two merits: 1) the (non-linear) transformation $h_{l}^{*}$ can be different across layers, and 2) for arbitrary node $u$ in the original graph $G$, its causal effect on distinct centered nodes $v$ could be different dependent on its relative position in the ego-graph $G_{v}$. Therefore, this formulation gives rise to enough flexibility and capacity for modeling on graph data.

We consider a linear toy example and assume 1-layer graph convolution for illustration. Namely, the ego-graph $G_{v}$ (and $N_{v}$) only contains the centered node and its 1-hop neighbors. We simplify the $h^{*}$ and $c^{*}$ in Assumption 1 as identity mappings and instantiate $\Gamma$ as a mean pooling function. Then we assume 2-dim node features $x_{v}=[x_{v}^{1},x_{v}^{2}]$ and

where $n_{v}^{1}$ and $n_{v}^{2}$ are independent standard normal noise and $\epsilon$ is a random variable with zero mean and non-zero variance dependent on environment $e$. In Fig. 1(c) we show the dependency among these random variables in a graphical representation and instantiate them in an example of citation networks, where a paper’s published avenue is an invariant feature for predicting the paper’s sub-area while its citation index (a spurious feature) is affected by both the label and the environment.

The new objective tackles the variance across environments and guarantees the desirable solution. The enlightenment is that if the model yields equal performance on different $e$’s, it would manage to leverage the invariant features, which motivates us to devise a new objective for solving Eq. 2.

We now return to the general case where we have $\{(G_{v},y_{v})\}$ for training and leverage a GNN model as the predictor: $\hat{y}_{v}=f_{\theta}(G_{v})$. The intuition in Section 3.1 implies a new learning objective:

where $L(G^{e},Y^{e};\theta)=\frac{1}{|V_{e}|}\sum_{v\in V_{e}}l(f_{\theta}(G_{v}^{e}),y_{v}^{e})$ and $\beta$ is a trading hyper-parameter. If we have training graphs from a sufficient number of environments $\mathcal{E}_{tr}=\{e\}$ and the correspondence of each graph to a specific $e$, i.e., $\{G^{e},Y^{e}\}_{e\in\mathcal{E}_{tr}}$ which induces $\{\{G^{e}_{v},y^{e}_{v}\}_{v\in V_{e}}:e\in\mathcal{E}_{tr}\}$, we can use the empirical estimation with risks from different environments to handle Eq. 4 in practice, as is done by the Risk Extrapolation (Rex) approach (Krueger et al., 2021). Unfortunately, as mentioned before, for node-level tasks on graphs, the training data is often a single graph (without any correspondence of nodes to environments), and hence, one only has training data from a single environment. Exceptions are some multi-graph scenarios where one can assume each graph is from an environment, but there are still a very limited number of training graphs (e.g., less than five). The objective Eq. 4 would require data from diverse environments to enable the model for desirable extrapolation. To detour such a dilemma, we introduce $K$ auxiliary context generators $g_{w_{k}}(G)$ ($k=1,\cdots,K$) that aim to generate $K$-fold graph data $\{G^{k}\}_{k=1}^{K}$ (which induces $\{\{G^{k}_{v}\}_{v\in V}:1\leq k\leq K\}$) based on the input one $G$ and mimics training data from different environments. The generators are trained to maximize the variance loss so as to explore the environments and facilitate stable learning of the GNN:

where $L(g_{w_{k}}(G),Y;\theta)=L(G^{k},Y;\theta)=\frac{1}{|V|}\sum_{v\in V}l(f_{\theta}(G_{v}^{k}),y_{v})$.

One remaining problem is how to specify $g_{w_{k}}(G)$. Following recent advances in adversarial robustness on graphs (Xu et al., 2019; Jin et al., 2020), we consider editing graph structures by adding/deleting edges. Assume a Boolean matrix $B^{k}=\{0,1\}^{N\times N}$ ($k=1,\cdots,K$) and denote the supplement graph of $A$ as $\overline{A}=\mathbf{1}\mathbf{1}^{\top}-I-A$, where $I$ is an identity matrix. Then the modified graph for view $k$ is $A^{k}=A+B^{k}\circ(\overline{A}-A)$ where $\circ$ denotes element-wise product. The optimization for $B^{k}$ is difficult due to its non-differentiability and one also needs to constrain the modification within a threshold. To handle this, we use policy gradient method REINFORCE, treating graph generation as a decision process and edge editing as actions (see details in Appendix A). We call our approach in Eq. 5 Explore-to-Extrapolate Risk Minimization (Eerm) and present our training algorithm in Alg. 1.

The above results imply that the objective Eq. 4 can guarantee a valid solution for the formulated OOD problem on graph-structured data, which serves as a theoretical justification for our approach.

The condition can be satisfied once $\mathbf{z}$ is a sufficient representation across environments. Therefore, we have proven that the new objective could help to reduce the generalization error on out-of-distribution data and indeed enhance GNN model’s power for in-the-wild extrapolation.

We first consider artificial distribution shifts based on two public node classification benchmarks Cora and Amazon-Photo. For each dataset, we adopt two randomly initialized GNNs to 1) generate node labels based on the original node features and 2) generate spurious features based on the node labels and environment id, respectively (See Appendix E.1 for details). We generate 10-fold graph data with distinct environment id’s and use 1/1/8 of them for training/validation/testing.

We use a 2-layer vanilla GCN (Kipf & Welling, 2017) as the GNN model. We report results on 8 testing graphs (T1$\sim$ T8) of the two datasets in Fig. 2(a) and 3(a), respectively, where we also adopt 2-layer GCNs for data generation. The results show that our approach Eerm consistently outperforms Erm on Cora and Photo, which suggests the effectiveness of our approach for handling distribution shifts. We also observe that in Photo, the performance variances within one graph and across different test graphs are both much lower compared with those in Cora. We conjecture the reasons are two-fold. First, there is evidence that in Cora the (original) features from adjacent nodes are indeed informative for prediction while in Photo this information contributes to negligible gain over merely using centered node’s features. Based on this, once the node features are mixed up with invariant and spurious ones, it would be harder for distinguishing them in the former case that relies more on graph convolution.

In Fig. 2(b) and Fig. 3(b), we compare the averaged training accuracy (achieved by the epoch with the highest validation accuracy) given by two approaches when using all the input features and removing the spurious ones for inference (we still use all the features for training in the latter case). As we can see, the performance of Erm drops much more significantly than Eerm when we remove the spurious input features, which indicates that the GCN trained with standard Erm indeed exploits spurious features to increase training accuracy while our approach can help to alleviate such an issue and guide the model to focus on invariant features. Furthermore, in Fig. 2(c) and Fig. 3(c), we compare the test accuracy averaged on eight graphs when using different GNNs e.g. GCN, SGC (Wu et al., 2019) and GAT (Velickovic et al., 2018), for data generation (See Appendix G for more results). The results verify that our approach achieves consistently superior performance in different cases.

We proceed to consider another scenario where there are multiple observed graphs in one dataset and a model trained with one graph or a limited number of graphs is expected to generalize to new unseen graphs. The graphs of a dataset share the input feature space and output space and may have different sizes and data distributions since they are collected from different domains. We adopt two public social network datasets Twitch-Explicit and Facebook-100 collected by Lim et al. (2021).

Training with a Single Graph. In Twitch, we adopt a single graph DE for training, ENGB for validation and the remaining five networks (ES, FR, PTBR, RU, TW) for testing. We follow Lim et al. (2021) and use test ROC-AUC for evaluation. We specify the GNN model as GCN, GAT and a recently proposed model GCNII (Chen et al., 2020a) that can address the over-smoothing of GCN and enable stacking of deep layers. The layer numbers are set as 2 for GCN and GAT and 10 for GCNII. Fig. 4 compares the results on five test graphs, where Eerm significantly outperforms Erm in most cases with up to $7.0\%$ improvement on ROC-AUC. The results verify the effectiveness of Eerm for generalizing to new graphs from unseen domains.

Training with Multiple Graphs. In FB-100, we adopt three graphs for training, two graphs for validation and the remaining three for testing. We also follow Lim et al. (2021) and use test accuracy for evaluation. We use GCN as the backbone and compare using different configurations of training graphs, as shown in Table 2. We can see that Eerm outperforms Erm on average on all the test graphs with up to $7.2\%$ improvement. Furthermore, Eerm maintains the superiority with different training graphs, which also verifies the robustness of our approach w.r.t. training data.

We consider the third scenario where the input data are temporal dynamic graphs and the model is trained with dataset collected at one time and needs to handle newly arrived data in the future. Here are also two sub-cases that correspond to distinct real-world scenarios, as discussed below.

Handling Dynamic Graph Snapshots. We adopt a dynamic financial network dataset Elliptic (Pareja et al., 2020) that contains dozens of graph snapshots where each node is a Bitcoin transaction and the goal is to detect illicit transactions. We use 5/5/33 snapshots for train/valid/test. Following Pareja et al. (2020) we use F1 score for evaluation. We consider two GNN architectures as the backbone: GraphSAGE (Hamilton et al., 2017) and GPRGNN (Chien et al., 2021) that can adaptively combine information from node features and graph topology. The results are shown in Fig. 5 where we group the test graph snapshots into 9 folds in chronological order. Our approach yields much higher F1 scores than Erm in most cases with averagely $9.6\%/10.0\%$ improvements using GraphSAGE/GPR-GNN as backbones. Furthermore, there is an interesting phenomenon that both methods suffer a performance drop after T7. The reason is that this is the time when the dark market shutdown occurred (Pareja et al., 2020). Such an emerging event causes considerable variation to data distributions that leads to performance degrade for both methods, with Erm suffering more. In fact, the emerging event acts as an external factor which is unpredictable given the limited training data. The results also suggest that how neural models generalize to OOD data depends on the learning approach but its performance limit is dominated by observed data. Nonetheless, our approach contributes to better F1 scores than Erm even in such an extreme case.

Handling New Nodes in Temporally Augmented Graph. Citation networks often go through temporal augmentation with new papers published. We adopt OGB-Arxiv (Hu et al., 2020) for experiments and enlarge the time difference between training and testing data to introduce distribution shifts: we select papers published before 2011 for training, in-between 2011 and 2014 for validation, and within 2014-2016/2016-2018/2018-2020 for testing. Also different from the original (transductive) setting in Hu et al. (2020), we use the inductive learning setting, i.e., test nodes are strictly unseen during training, which is more akin to practical situations. Table 3 presents the test accuracy and shows that Eerm outperforms Erm in five cases out of six. Notably, when using GPRGNN as the backbone, Eerm manages to achieve up to $8.1\%$ relative improvement, which shows that Eerm is capable of improving GNN model’s learning for extrapolating to future data.