HANDLING DISTRIBUTION SHIFTS ON GRAPHS: AN INVARIANCE PERSPECTIVE
Year: 2022

Abstract:
There is increasing evidence suggesting neural networks’ sensitivity to distribution shifts, so that research on out-of-distribution (OOD) generalization comes into the spotlight. Nonetheless, current endeavors mostly focus on Euclidean data, and its formulation for graph-structured data is not clear and remains under-explored, given the two-fold fundamental challenges: 1) the inter-connection among nodes in one graph, which induces non-IID generation of data points even under the same environment, and 2) the structural information in the input graph, which is also informative for prediction. In this paper, we formulate the OOD problem for node-level prediction on graphs and develop a new domain-invariant learning approach, named Explore-to-Extrapolate Risk Minimization, that facilitates GNNs to leverage invariant graph features for prediction. The key difference to existing invariant models is that we design multiple context explorers (specified as graph editers in our case) that are adversarially trained to maximize the variance of risks from multiple virtual environments. Such a design enables the model to extrapolate from a single observed environment which is the common case for node-level prediction. We prove the validity of our method by theoretically showing its guarantee of a valid OOD solution and further demonstrate its power on various real-world datasets for handling distribution shifts from artificial spurious features, cross-domain transfers and dynamic graph evolution.

heading: 1 INTRODUCTION
text: As the demand for handling in-the-wild unseen instances draws increasing concerns, out-ofdistribution (OOD) generalization (Mansour et al., 2009; Blanchard et al., 2011; Muandet et al., 2013; Gong et al., 2016) occupies a central role in the ML community. Yet, recent evidence suggests that deep neural networks can be sensitive to distribution shifts, exhibiting unsatisfactory performance within new environments, e.g., Beery et al. (2018); Su et al. (2019); Recht et al. (2019); Mancini et al. (2020). A more concerning example is that a model for COVID-19 detection exploits undesired ‘shortcuts’ from data sources (e.g., hospitals) to boost training accuracy (DeGrave et al., 2020).
Recent studies of the OOD generalization problem like Rojas-Carulla et al. (2018); Bühlmann (2018); Gong et al. (2016); Arjovsky et al. (2019) treat the cause of distribution shifts between training and testing data as a potential unknown environmental variable e. Assuming that the goal is to predict target label y given associated input x, the environmental variable would impact the underlying data generating distribution p(x,y|e) = p(x|e)p(y|x, e). With E as the support of environments, f(·) as ∗This work was done during authors’ internship at AWS Shanghai AI Lab.
a prediction model and l(·, ·) as a loss function, the OOD problem could be formally represented as
min f max e∈E
E(x,y)∼p(x,y|e=e)[l(f(x), y)|e]. (1)
Such a problem is hard to solve since the observations in training data cannot cover all the environments in practice. Namely, the actual demand is to generalize a model trained with data from p(x,y|e = e1) to new data from p(x,y|e = e2). Recent research opens a new possibility via learning domain-invariant models (Arjovsky et al., 2019) under a cornerstone data-generating assumption: there exists a portion of information in x that is invariant for prediction on y across different environments. Based on this, the key idea is to learn a equipredictive representation model h that gives rise to equal conditional distribution p(y|h(x), e = e) for ∀e ∈ E . The implication is that such a representation h(x) will bring up equally (optimal) performance for a downstream classifier under arbitrary environments. The model p̂(y|x) with such a property is called as invariant model/predictor. Several up-to-date studies develop new objective designs and algorithms for learning invariant models, showing promising power for tackling OOD generalization (Chang et al., 2020; Ahuja et al., 2020; Krueger et al., 2021; Liu et al., 2021; Creager et al., 2021; Koyama & Yamaguchi, 2021).
While the OOD problem is well-established in certain settings (where the dataset can be obviously modeled as a set of i.i.d. generated pairs (x, y) from p(x,y|e)), its formulation for graph-structured data, especially for node-level tasks on graphs (where we note that each node in the graph corresponds to an instance), is not clear and remains as an open problem. Compared with classic data format (e.g. vision or texts), graph-structured data have two fundamental differences: 1) the non-independent and non-identically distributed nature exists in data generation even within the same environment, embodied with the inter-connection among data points in one graph; 2) the structural information also plays a role for prediction beside the node features. These differences bring up unique technical challenges for handling distribution shifts of node-level tasks on graphs.
Distribution shifts indeed widely exist in real-world graphs. For instance, in citation networks, the distributions for paper citations (the input) and subject areas/topics (the label) would go through significant change as time goes by (Hu et al., 2020b). In social networks, the distributions for users’ friendships (the input) and their activity (the label) would highly depend on when/where the networks are collected (Fakhraei et al., 2015). In financial networks (Pareja et al., 2020), the payment flows between transactions (the input) and the appearance of illicit transactions (the label) would have strong correlation with some external contextual factors (like time and market). In these cases, neural models built on graph-structured data, particularly, Graph Neural Networks (GNNs) which are the common choice, need to effectively deal with OOD data during test time. Moreover, as GNNs have become popular and easy-to-implement tools for modeling relational structures in broad AI areas (vision, texts, audio, etc.), enhancing its robustness to distribution shifts is a pain point for building general AI systems, especially applied to high-stake applications like autonomous driving (Dai & Gool, 2018), medical diagnosis (AlBadawy et al., 2018), criminal justice (Berk et al., 2018), etc.
In this paper, we endeavor to 1) formulate the OOD problem for node-level tasks on graphs, 2) develop a new learning approach based on an invariance principle, 3) provide theoretical results to dissect its rationale, and 4) design comprehensive experiments to show its practical efficacy. Concretely:
1. To accommodate the non-IID nature of nodes in a graph, we fragment a graph into a set of egographs for centered nodes and decompose the data-generating process into: 1) sampling a whole input graph and 2) sampling each node’s label conditioned the ego-graph. Based on this, we can inherit the spirit of Eq. 1 to formulate the OOD problem for node-level tasks over graphs (see Section 2.1).
2. To take into account structural information that is useful for prediction, we first re-formulate the invariant assumption used in prior arts with recursive computation on the induced BFS trees of ego-graphs, inspired by the Weisfeiler-Lehman test (Weisfeiler & Lehman, 1968). Then, to endow GNNs with enough ability for handling distribution shifts, we devise a new learning approach, entitled Explore-to-Extrapolate Risk Minimization, that aims GNNs at minimizing the mean and variance of risks from multiple environments that are simulated by adversarial context generators (instantiated as graph editers), as shown in Fig. 1(a) (see Section 2.2 and 3).
3. To shed more insights on the rationales of the proposed approach and its relationship with the formulated OOD problem, we prove that our objective can guarantee a valid solution for the formulated OOD problem given some mild conditions and furthermore, an upper bound on the OOD error can be effectively controlled when minimizing the training error (see Section 4).
4. To evaluate the approach, we design a comprehensive set of experiments on diverse real-world nodelevel prediction datasets that entail distribution shifts from artificial spurious features, cross-domain transfers and dynamic graph evolution. We also apply our approach to distinct GNN backbones (GCN, GAT, GraphSAGE, GCNII and GPRGNN), and the results show that it consistently outperforms standard empirical risk minimization with promising improvements on OOD data (see Section 5).

heading: 2 PROBLEM FORMULATION
text: In this section, we present our problem formulation for the OOD node-level prediction problem on graphs and then introduce a cornerstone invariance assumption for data generation. All the random variables are denoted as bold letters while the corresponding realizations are denoted as thin letters.

heading: 2.1 OUT-OF-DISTRIBUTION PROBLEM FOR GRAPH-STRUCTURED DATA
text: An input graph G = (A,X) contains two-fold information1: an adjacency matrix A = {avu|v, u ∈ V } and node features X = {xv|v ∈ 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 = {yv|v ∈ V }. We define G as a random variable of input graphs and 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 G and the target as Y, and then the data generation can be characterized as p(G,Y|e) = p(G|e)p(Y|G, e) where 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 v as a random variable of nodes. We define node v’s L-hop neighbors as Nv (where L is an arbitrary integer) and the nodes in Nv form an ego-graph Gv which consists of a (local) node feature matrix Xv = {xu|u ∈ Nv} and a (local) adjacency matrix Av = {auw|u,w ∈ Nv}. Use Gv as a random variable of ego-graphs2 whose realization is Gv = (Av, Xv). Besides, we define y as a random variable of node labels. In this way, we can fragment a whole graph as a set of instances {(Gv, yv)}v∈V where Gv denotes an input and yv is a target. Notice that the ego-graph can be seen as a Markov blanket for the centered node, so the conditional distribution p(Y|G, e) can be decomposed as a product of |V | independent and identical marginal distributions p(y|Gv, e). Therefore, the data generation of {(Gv, yv)}v∈V from a distribution p(G,Y|e) can be considered as a two-step procedure: 1) the entire input graph is generated via G ∼ p(G|e) which can then be fragmented into a set of ego-graphs {Gv}v∈V ; 2) each node’s label is generated via y ∼ p(y|Gv = Gv, e). Then the OOD node-level prediction problem can be formulated as: given training data
1Our formulation and method can be trivially extended to cover edge features which we omit here for brevity. 2We use a subscript v here to remind that it is an ego-graph from the view of a target node.
{Gv, yv}v∈V from p(G,Y|e = e), the model needs to handle testing data {Gv, yv}v∈V ′ from a new distribution p(G,Y|e = e′). Denote E as the support of environments, f as a predictor model with ŷ = f(Gv) and l(·, ·) as a loss function. More formally, the OOD problem can be written as:
min f max e∈E EG∼p(G|e=e)
[ 1
|V | ∑ v∈V Ey∼p(y|Gv=Gv,e=e)[l(f(Gv), y)]
] . (2)
We remark that the first-step sampling G ∼ p(G|e = e) can be ignored since in most cases one only has a single input graph in the context of node-level prediction tasks.

heading: 2.2 INVARIANT FEATURES FOR NODE-LEVEL PREDICTION ON GRAPHS
text: 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.
With our definition in Section 2.1, for node-level prediction on graphs, each input instance is an ego-graph Gv with target label yv. It seems not straightforward for how to define invariant features on graphs given two observations: 1) the ego-graph possesses a hierarchical structure for associated nodes (i.e., Gv induces a BFS tree rooted at v where the l-th layer contains the l-order neighbored nodes N (l)v ) and 2) the nodes in each layer are permutation-invariant and variable-length. Inspired by Weisfeiler-Lehman test, we can adapt the definition of invariance assumption in prior arts (RojasCarulla et al., 2018; Gong et al., 2016; Arjovsky et al., 2019; Koyama & Yamaguchi, 2021; Liu et al., 2021) to accommodate structural information in graph data: Assumption 1. (Invariance Property of Node-Level Prediction) Assume input feature dimension as d0. There exists a sequence of (non-linear) functions {h∗l }Ll=0 where h∗l : Rd0 → Rd and a permutation-invariant function Γ : Rdm → Rd, which gives a node-level readout rv = r(L)v that is calculated in a recursive way: r(l)u = Γ{r(l−1)w |w ∈ N (1)u ∪{u}} for l = 1, · · · , L and r(0)u = h∗l (xu) if u ∈ N (l)v . Denote r as a random variable of rv and it satisfies that 1) (Invariance condition): p(y|r, e) = p(y|r), and 2) (Sufficiency condition): y = c∗(r) + n, where c∗ is a non-linear function and n is an independent noise.
A more intuitive illustration for the above computation is presented in Fig. 1(b). The node-level readout rv aggregates the information from neighbored nodes recursively along the structures of BFS tree given by Gv. Essentially, the above definition assumes that in each layer the neighbored nodes contain a portion of causal features that contribute to stable prediction for y across different 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 Gv . Therefore, this formulation gives rise to enough flexibility and capacity for modeling on graph data.

heading: 3 METHODOLOGY
text: We next present our solution for the challenging OOD problem. Before going into the formal method, we first introduce a motivating example based on Assumption 1 to provide some high-level intuition.

heading: 3.1 MOTIVATING EXAMPLE
text: We consider a linear toy example and assume 1-layer graph convolution for illustration. Namely, the ego-graph Gv (and Nv) only contains the centered node and its 1-hop neighbors. We simplify the h∗ and c∗ in Assumption 1 as identity mappings and instantiate Γ as a mean pooling function. Then we assume 2-dim node features xv = [x1v, x 2 v] and
yv = 1 |Nv| ∑ u∈Nv x1u + n 1 v, x 2 v = 1 |Nv| ∑ u∈Nv yu + n 2 v + , (3)
where n1v and n 2 v are independent standard normal noise and 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.
Based on this, we consider a vanilla GCN as the predictor model ŷv = 1|Nv| ∑ u∈Nv θ1x 1 u + θ2x 2 u. Then the ideal solution for the predictor model is [θ1, θ2] = [1, 0]. This indicates that the GCN identifies the invariant feature, i.e., x1v insensitive to environment changes. However, here we show a negative result when using standard empirical risk minimization.
Proposition 1. Let the risk under environment e be R(e) = 1|V | ∑ v∈V Ey|Gv=Gv [‖ŷv − yv‖22]. The unique optimal solution for objective minθ Ee[R(e)] would be [θ1, θ2] = [ 1+σ 2 e
2+σ2e , 12+σ2e ] where σe > 0 denotes the standard deviation of across environments.
This indicates that directly minimizing the expectation of risks across environments would inevitably lead the model to rely on spurious correlation (x2v depends on environments). Also, such a reliance would be strengthened with smaller σe, i.e., when there is less uncertainty for the effect from environments. To mitigate the issue, fortunately, we can prove another result that implies a new objective as a sufficient condition for the ideal solution.
Proposition 2. The objective minθ Ve[R(e)] reaches the optimum if and only if [θ1, θ2] = [1, 0].
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.

heading: 3.2 STABLE LEARNING WITH EXPLORE-TO-EXTRAPOLATE RISK MINIMIZATION
text: We now return to the general case where we have {(Gv, yv)} for training and leverage a GNN model as the predictor: ŷv = fθ(Gv). The intuition in Section 3.1 implies a new learning objective:
min θ
Ve[L(Ge, Y e; θ)] + βEe[L(Ge, Y e; θ)], (4)
where L(Ge, Y e; θ) = 1|Ve| ∑ v∈Ve l(fθ(G e v), y e v) and β is a trading hyper-parameter. If we have training graphs from a sufficient number of environments Etr = {e} and the correspondence of each graph to a specific e, i.e., {Ge, Y e}e∈Etr which induces {{Gev, yev}v∈Ve : e ∈ Etr}, 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 gwk(G) (k = 1, · · · ,K) that aim to generate K-fold graph data {Gk}Kk=1 (which induces {{Gkv}v∈V : 1 ≤ k ≤ 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:
min θ
Var({L(gw∗k(G), Y ; θ) : 1 ≤ k ≤ K}) + β
K K∑ k=1 L(gw∗k(G), Y ; θ),
s. t. [w∗1 , · · · , w∗K ] = arg max w1,··· ,wK
Var({L(gwk(G), Y ; θ) : 1 ≤ k ≤ K}), (5)
where L(gwk(G), Y ; θ) = L(G k, Y ; θ) = 1|V | ∑ v∈V l(fθ(G k v), yv).
One remaining problem is how to specify gwk(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 Bk = {0, 1}N×N (k = 1, · · · ,K) and denote the supplement graph of A as A = 11> − I −A, where I is an identity matrix. Then the modified graph for view k is Ak = A+Bk ◦ (A−A) where ◦ denotes element-wise product. The optimization for Bk 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.

heading: 4 THEORETICAL DISCUSSIONS
text: We next present theoretical analysis to shed insights on the objective and its relationship with our formulated OOD problem in Section 2.1. To begin with, we introduce some building blocks. The GNN model f can be decomposed into an encoder h for representation and a classifier c for prediction, i.e., f = c ◦ h and we have zv = h(Gv), ŷv = c(zv). Besides, we assume I(x;y) stands for the mutual information between x and y and I(x;y|z) denotes the conditional mutual information given z. To keep notations simple, we define pe(·) = p(·|e = e) and Ie(·) = I(·|e = e). Another tricky point is that in computation of the KL divergence and mutual information, we require the samples from the joint distribution pe(G,Y), which also results in difficulty for handling data generation of interconnected nodes. Therefore, we again adopt our perspective in Section 2.1 and consider a two-step sampling procedure. Concretely, for any probability function f1, f2 associated with ego-graphs Gv and node labels y, we define computation for KL divergence as
DKL(f1(Gv,y)‖f2(Gv,y)) := EG∼p(G)
[ 1
|V | ∑ v∈V Eyv∼p(y|Gv=Gv) [ log f1(Gv = Gv,y = yv) f2(Gv = Gv,y = yv) ]] .
(6)

heading: 4.1 RELATIONSHIP BETWEEN INVARIANCE PRINCIPLE AND OOD PROBLEM
text: We will show that the objective Eq. 4 can guarantee a valid solution for OOD problem Eq. 2. To this end, we rely on another assumption for data-generating distribution. Assumption 2. (Environment Heterogeneity): For (Gv, r) that satisfies Assumption 1, there exists a random variable r such that Gv = m(r, r) according to the functional representation lemma. We assume that p(y|r, e = e) would arbitrarily change across environments e ∈ E .
Assumptions 1 and 2 essentially distill two portions of features in input data: one is domain-invariant for prediction and the other contributes to sensitive prediction that depends on environments. The GNN model f = c ◦ h induces two model distributions q(z|Gv) (by the encoder) and q(y|z) (by the classifier). Based on this, we can dissect the effects of Eq. 4 which indeed forces the representation z to satisfy the invariance and sufficiency conditions illustrated in Assumption 1. Theorem 1. If q(y|z) is treated as a variational distribution, then 1) minimizing the expectation term in Eq. 4 contributes to maxq(z|Gv) I(y; z), i.e., enforcing the sufficiency condition on z for prediction, and 2) minimizing the variance term in Eq. 4 would play a role for minq(z|Gv) I(y; e|z), i.e., enforcing the invariance condition p(y|z, e) = p(y|z).
Based on these results, we can bridge the gap between the invariance principle and OOD problem. Theorem 2. Under Assumption 1 and 2, if the GNN encoder q(z|Gv) satisfies that 1) I(y; e|z) = 0 (invariance condition) and 2) I(y; z) is maximized (sufficiency condition), then the model f∗ given by Ey[y|z] is the solution to OOD problem in Eq. 2.
The proof for Theorem 2 follows a similar line of Liu et al. (2021). The above theorems indicate that the objective Eq. 4 can guarantee a valid solution for the formulated node-level OOD problem on graph-structured data, which serves as a theoretical justification for our approach.

heading: 4.2 INFORMATION-THEORETIC ERROR FOR OOD GENERALIZATION
text: We proceed to analyze the OOD generalization error given by our learning approach. Recall that we assume training data from p(G,Y|e = e) and testing data from p(G,Y|e = e′). In fact, the training error and OOD generalization error can be respectively measured by the discrepancy terms: DKL(pe(y|Gv)‖q(y|Gv)) and DKL(pe′(y|Gv)‖q(y|Gv)) which can be calculated based on our definition in Eq. 6. This allows us to generalize the information-theoretic framework (Federici et al., 2021) for analysis on graph data. Based on Theorem 1, we can arrive at the following theorem which reveals the effect of Eq. 4 that contributes to tightening the bound for the OOD error. Theorem 3. Optimizing Eq. 4 with training data can minimize the upper bound for DKL(pe′(y|Gv)‖q(y|Gv) on condition that Ie′(Gv;y|z) = Ie(Gv;y|z).
The condition can be satisfied once 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.

heading: 5 EXPERIMENTS
text: In this section, we aim to verify the effectiveness and robustness of our approach in a wide variety of tasks reflecting real situations, using different GNN backbones. Table 1 summarizes the information of experimental datasets and evaluation protocols, and we provide more dataset information in Appendix E. We compare our approach EERM with standard empirical risk minimization (ERM). Implementation details are presented in Appendix F. In the following subsections, we will investigate three scenarios that require the model to handle distribution shifts stemming from different causes.

heading: 5.1 HANDLING DISTRIBUTION SHIFTS WITH ARTIFICIAL TRANSFORMATION
text: 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 test graphs (T1∼ T8) of 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 consistently outperforms ERM. In Cora/Photo, we manage to achieve 9.1%/2.6% improvement on average, 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.

heading: 5.2 GENERALIZING TO UNSEEN DOMAINS
text: 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 4.8% 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 12.2% improvement. Furthermore, EERM maintains the superiority with different training graphs, which also verifies the robustness of our approach w.r.t. training data.

heading: 5.3 EXTRAPOLATING OVER DYNAMIC DATA
text: 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
T1 T2 T3 T4 T5 T6 T7 T8 T9 (a) GraphSAGE
50
60
70
F1 S
co re
ERM EERM
T1 T2 T3 T4 T5 T6 T7 T8 T9 (b) GPRGNN
50
60
70
F1 S
co re
ERM EERM
and the goal is to detect illicit transactions. We use 5/5/33 snapshots for training/validation/testing. 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 a recently proposed model 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 better F1 scores than ERM in most cases with 2.5% ∼ 28.1% improvements, which again verifies its superiority. 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., 2020b) 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. (2020b), 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 7.8% relative improvement, which shows that EERM is capable of improving GNN model’s learning for extrapolating to future data.

heading: 6 COMPARISON WITH EXISTING WORKS
text: We compare with some related works, highlight our differences and discuss the potential impacts on broad areas. Due to the space limit, we defer more discussions to Appendix B.
Generalization on Graph Data. Recent endeavors (Scarselli et al., 2018; Garg et al., 2020; Verma & Zhang, 2019) derive generalization error bounds for GNNs on node-level tasks. Yet, they focus on in-distribution generalization and put little emphasis on distribution shifts, which are the main focus of our work. Furthermore, some up-to-date works explore GNN’s extrapolation ability for OOD data, e.g. unseen features/structures (Xu et al., 2021) and varying graph sizes (Yehudai et al., 2021; Bevilacqua et al., 2021). However, they concentrate on graph-level tasks (e.g., graph classification), where each input instance is a graph (usually with less than 100 nodes) and one dataset contains massive graphs for training and testing. By contrast, in node-level tasks, i.e., what this paper studies, each input is a node in one graph (usually with ∼1K to ∼1M nodes) and a dataset usually contains only a single graph. The graph-level problems have straightforward relationship to the general setting (in Eq. 1) since one can treat input graphs as x and graph labels as y and then the data from one environment becomes a set of i.i.d. generated pairs (x, y). Differently, node-level problems cannot be tackled in this way due to the inter-connection among data points that results in non-IID nature in data generation within one environment. To resolve this case, our work introduces a new perspective for problem formulation, backed up with a concrete approach for problem solving. Also, EERM can be adapted to the general setting especially for generalization from a single observed environment.
Benchmarking OOD Problems. A surge of recent works release and study OOD benchmarks, e.g., Worrall et al. (2017); Hendrycks et al. (2019); Gulrajani & Lopez-Paz (2020); Xiao et al. (2020); Santurkar et al. (2020); Ye et al. (2021). A very recent literature (Koh et al., 2021) summarizes the ways that prior works resort to for introducing distribution shifts to datasets, including 1) artificial
transformations, 2) synthetic-to-real transfers, 3) constrained splits and 4) cross-dataset transfers. As far as we know, the majority of them focus on classic data format (vision, texts, tabular data, etc.), and there are few studies designed for graph-structured data, e.g. the OGB-MolPCBA (Hu et al., 2020b) for graph-level classification/regression. As a by-product, our experiment designs (including datasets, splits and evaluation protocols) for three distinct scenarios (artificial transformations, cross-graph transfers and (time)-constrained splits) could help to enrich the OOD benchmarking zoo, particularly for node-level tasks on graphs, which remains unexplored in the literature.

heading: 7 CONCLUSION
text: This work targets out-of-distribution generalization for graph-structured data with the focus on node-level problems where the inter-connection of data points hinders trivial extension from existing formulation and methods. To this end, we take a fresh perspective to formulate the problem in a principled way and further develop a new approach for extrapolation from a single environment, backed up with theoretical guarantees. We also design comprehensive experiments to show the practical power of our approach on various real-world datasets with diverse distribution shifts.

heading: A OPTIMIZATION AND ALGORITHM
text: We illustrate the details of using policy gradient for optimizing the graph editers in Eq. 5. Concretely, for view k, we consider a parameterized matrix Pk = {πknm}. For the n-th node, the probability for editing the edge between it and the m-th node would be p(aknm) = exp(πknm)∑ m′ exp(π k nm′ ) . We then sample s actions {bknmt} s t=1 from a multinomial distributionM(p(akn1), · · · , p(aknn)), which give the nonzero entries in the n-th row of Bk. The reward function R(Gk) can be defined as the inverse loss. We can use REINFORCE algorithm to optimize the generator with the gradient∇wk log pwk(Ak)R(Gk) where wk = Pk and pwk(A k) = ΠnΠ s t=1p(b k nmt). We present the training algorithm in Alg. 1.
Algorithm 1: Stable Learning for OOD Generalization in Node-Level Prediction on Graphs. 1 INPUT: training graph data G = (A,X) and Y , initialized parameters of GNN θ, initialized
parameters of graph editers w = {wk}, learning rates αg , αf . 2 while not converged or maximum epochs not reached do 3 for t = 1, · · · , T do 4 Obtain modified graphs Gk = (Ak, X) from graph editer gwk , k = 1, · · · ,K; 5 Compute loss J1(w) = Var({L(Gk, Y ; θ) : 1 ≤ k ≤ K}) ; 6 Update wk ← wk + αg∇wk log pwk(Ak)J1(w), k = 1, · · · ,K ; 7 if t == T then 8 Compute loss J2(θ) = Var({L(Gk, Y ; θ) : 1 ≤ k ≤ K}) + βK ∑K k=1 L(G
k, Y ; θ) ; 9 Update θ ← θ − αf∇θJ2(θ) ;
10 OUTPUT: trained parameters of GNN θ∗.

heading: B FURTHER RELATED WORKS

heading: B.1 OUT-OF-DISTRIBUTION GENERALIZATION AND INVARIANT MODELS
text: Out-of-distribution generalization has drawn extensive attention in the machine learning community. To endow the learning systems with the ability for handling unseen data from new environments, it is natural to learn invariant features under the setting of the causal factorization of physical mechanisms (Schölkopf et al., 2012; Peters et al., 2016). A recent work (Arjovsky et al., 2019) proposes Invariant Risk Minimization (IRM) as a practical solution for OOD problem via invariance principle. Based on this, follow-up works make solid progress in this direction, e.g., with group distributional robust optimization (Sagawa et al., 2019), causal attribution (Chang et al., 2020), game theory (Ahuja et al., 2020), lottery ticket hypothesis (Zhang et al., 2021), etc. Several works attempt to resolve extended settings. For instance, Ahmed et al. (2021) proposes to match the output distribution spaces from different domains via some divergence, while a recent work (Mahajan et al., 2021) also leverages a matching-based algorithm that resorts to shared representations of cross-domain inputs from the same object. Also, Creager et al. (2021) and Liu et al. (2021) point out that in most real situations, one has no access to the correspondence of each data point in the dataset with a specific environment, based on which they propose to estimate the environments as a latent variable.
Krueger et al. (2021) devises Risk Extrapolation (REX) which aims at minimizing the weighted combination of the variance and the mean of risks from multiple environments. Xie et al. (2020) contributes to a similar objective from different theoretical perspective. Also, Koyama & Yamaguchi (2021) extends the spirit of MAML algorithm and arrives at a similar objective form. In our model, we also consider minimization of the combination of variance and mean terms (in Eq. 4) and on top of that we further propose to optimize through a bilevel framework in Eq. 5. Compared to existing works, the differences of EERM are two-folds. First, we do not assume input data from multiple environments and the correspondence between each data point and a specific environment. Instead, our formulation enables learning and extrapolation from a single observed environment. Second, on methodology side, we introduce multiple context generators that aim to generate data of virtual environments in an adversarial manner. Besides, our formulation in this paper focus on node-level prediction on graphs, where the essential difference, as mentioned before, lies in the inter-connection
of data points in one graph (that corresponds to an environment), which hinders trivial adaption from existing works in the general setting.
There are also some recent studies that discuss the pitfalls of IRM in some cases by analyzing its performance on concrete examples (Rosenfeld et al., 2021; Nagarajan et al., 2021; Kamath et al., 2021). A recent work (Federici et al., 2021) harnesses an information-theoretic perspective to unify existing invariant models and provide insightful reflections on current endeavors.

heading: B.2 GRAPH NEURAL NETWORKS AND GENERALIZATION
text: Another line of research related to us attempts to enhance the generalization ability of graph neural networks via modifying the graph structures. One category of recent works is to learn new graph structures based on the input graph and node features. To improve the generalization power, a common practice is to enforce a certain regularization for the learned graph structures. For example, Jin et al. (2020) proposes to constrain the sparsity and smoothness of graphs via matrix norms and further adopts proximal gradient methods for handling the non-differentiable issue. Chen et al. (2020b); Zhang et al. (2019) also aim to regularize the sparsity and smoothness but differently harness energy function to enforce the constraints. From a different perspective, Xu et al. (2019) attempts to attack the graph topology for improving the model’s robustness and proposes to leverage projected gradient descent to make it tractable for the optimization of discrete graph structures.
Alternatively, several works focus on pruning the graph networks (Chen et al., 2021) or adaptively sparsifying the structures (Zheng et al., 2020; Hasanzadeh et al., 2020). While in our model, the introduced context generators are specified as graph editers which also attempt to optimize new graph structures, our big picture motivation and specific method are quite different from the abovementioned works. On problem setting side, we focus on out-of-distribution generalization and target handling distribution shifts over graphs, which is more difficult than the setting of previous methods that concentrate on in-distribution generalization. On methodology side, our context generators aim to generate data of multiple virtual environments and are learned from maximizing the variance of risks of multiple (virtual) environments. In other words, these context generators work adversarially among each other and collaboratively with the GNN backbone to enable the model for out-of-distribution generalization from a single observed environment. This is a new learning method. Also, one can specify our context generators with other existing graph generation/editing/attacking frameworks as mentioned above, which we leave as future works.
There are a few recent studies that focus on out-of-graph problems. For example, Hu et al. (2020a) considers transferable active learning over multiple graphs and treats the labeling process for nodes in an input graph as a decision process for optimization. Furthermore, Baek et al. (2020) deals with link prediction in knowledge graphs and proposes a meta-learning approach for extrapolating to unseen nodes out of the input graph. While these studies concentrate on transferring a model to new data, their main bodies including the formulations, approaches and experiment designs are not aimed at OOD generalization where there exist distribution shifts between training and testing data.
There is a very recent work (Baranwal et al., 2021) that endeavors to understand the out-of-distribution generalization ability of GNNs for semi-supervised node classification. It introduces contextual stochastic block model that is a mix-up of standard stochastic block model and a Gaussian mixture model with each node class corresponding to a component, based on which the authors show some cases where linear separability can be achieved. By contrast, our work possesses the following distinct technical contributions. First, we formulate OOD problem for node-level tasks in a more general setting without any assumption on specific distribution forms or the way for generation of graph structures3 Second, our proposed formulation, model and algorithm are rooted on invariant models, providing a new perspective and methodology for node-level prediction on graphs. Third, compared with Baranwal et al. (2021) that focus on synthetic datasets and simulated distribution shifts with artificially adding inter-class edges, we design and conduct comprehensive experiments on diverse real-world datasets that can reflect in-the-wild nature in real situations (e.g., cross-graph transfers and dynamic evolution) and demonstrate the power of our approach in three scenarios.
3Our Assumption 1 and 2 mainly focus on the existence of causal features and spurious features in data generation across different environments.

heading: C PROOFS FOR SECTION 3.1

heading: C.1 PROOF FOR PROPOSITION 1
text: We define the aggregated node feature av = 1|Nv| ∑ u∈Nv xu. According to the definition and setup in Section 3.1, we derive the risk under a specific environment R(e):
1 |V | ∑ v∈V Ey|Gv=Gv [ ‖ŷv − yv‖22 ] = 1
|V | ∑ v∈V En1,n2 [ ‖θ1a1v + θ2a2v − yv‖22 ] = 1
|V | ∑ v∈V En1,n2 [ ‖(θ1 + θ2 − 1)a1v + θ2(n1v + n2v + )− n1v‖22 ] .
(7)
Denote the objective for empirical risk minimization as L1 = Ee[R(e)] and we have its first-order derivative w.r.t. θ1 as
∂L1 ∂θ1 = Ee
[ 1
|V | ∑ v∈V En1,n2 [ 2[(θ1 + θ2 − 1)a1v + θ2(n1v + n2v + )− n1v] · a1v
]]
= Ee
[ 1
|V | ∑ v∈V En1,n2 [ 2[(θ1 + θ2 − 1) · a1v · a1v
]] ,
(8)
where the second step is given by independence among a1v, n 1 v, n 2 v and . Let ∂L1 ∂θ1 = 0, and we will obtain θ1 + θ2 = 1.
Also, the first-order derivative w.r.t. θ2 is
∂L1 ∂θ2 = Ee
[ 1
|V | ∑ v∈V En1,n2 [ 2[(θ1 + θ2 − 1)a1v + θ2(n1v + n2v + )− n1v] · (a1v + n1v + n2v + )
]]
= Ee
[ 1
|V | ∑ v∈V En1,n2 [ 2[(θ1 + θ2 − 1) · a1v · a1v + θ2(n1v · n1v + n2v · n2v + · )− n1v · n1v
]]
= Ee
[ 1
|V | ∑ v∈V En1,n2 [ 2[(θ1 + θ2 − 1) · a1v · a1v + θ2(1 + 1 + σ2e)− 1
]] ,
(9)
where the last step is according to Ex[x2] = E2x[x] + Vx[x]. We further let ∂L1∂θ2 = 0, and will get the unique solution
θ1 = 1 + σ2e 2 + σ2e , θ2 = 1 2 + σ2e . (10)

heading: C.2 PROOF FOR PROPOSITION 2
text: Let L2 = Ve[R(e)] = Ee[R2(e)]− E2e[R(e)] and l(e) = (θ1 + θ2 − 1)a1v + θ2(n1v + n2v + )− n1v . We derive the first-order derivation of L2 w.r.t. θ1 and θ2. Firstly,
∂L2 ∂θ1 = Ee
[ 1
|V | ∑ v∈V En1,n2 [ 4l3(e)a1v ]] + E2e [ 1 |V | ∑ v∈V En1,n2 [ 2l(e)a1v ]] , (11)
∂L2 ∂θ2 = Ee
[ 1
|V | ∑ v∈V En1,n2 [ 4l3(e) · (a1v + n1v + n2v + )
]]
+ E2e
[ 1
|V | ∑ v∈V En1,n2 [ 2l(e) · (a1v + n1v + n2v + )
]] .
(12)
By letting ∂L2∂θ1 = 0, we obtain the equation θ1 + θ2 = 1. Plugging it into ∂L2 ∂θ2 we have
∂L2 ∂θ2 = Ee
[ 1
|V | ∑ v∈V En1,n2 [ 4(θ2(n 1 v + n 2 v + )− n1v)3 · (a1v + n1v + n2v + )
]]
+ E2e
[ 1
|V | ∑ v∈V En1,n2 [ 2(θ2(n 1 v + n 2 v + )− n1v) · (a1v + n1v + n2v + )
]] ,
(13)
which is a function of e unless [θ1, θ2] = [1, 0] that gives rise to ∂L2∂θ2 = 0 for arbitrary distributions of environments. We thus conclude the proof.

heading: D PROOFS FOR SECTION 4

heading: D.1 PROOF FOR THEOREM 1
text: We first present a useful lemma that interprets the invariance and sufficiency conditions with the terminology of information theory. Lemma 1. The two conditions in Assumption 1 can be equivalently expressed as 1) (Invariance): I(y; e|r) = 0 and 2) (Sufficiency): I(y; r) is maximized.
Proof. For the invariance, we can easily arrive at the equivalence given the fact I(y; e|r) = DKL(p(y|e, r)‖p(y|r)). (14)
For the sufficiency, we first prove that for (Gv, r,y) satisfying that y = c∗(r) +n would also satisfy that r = arg maxr I(y; r). We prove it by contradiction. Suppose that r 6= arg maxr I(y; r) and r′ = arg maxr I(y; r) where r′ 6= r. Then there exists a random variable r such that r′ = m(r, r) where m is a mapping function. We thus have I(y; r′) = I(y; r, r) = I(c∗(r); r, r) = I(c∗(r); r) = I(y; r) which leads to contradiction.
We next prove that for (Gv, r,y) satisfying that r = arg maxr I(y; r) would also satisfy that y = c∗(r) + n. Suppose that y 6= c∗(r) + n and y = c∗(r′) + n where r′ 6= r. We then have the relationship I(c∗(r′); r) ≤ I(c∗(r′); r′) which yields that r′ = arg maxr I(y; r) and leads to contradiction.
Given the dependency relationship z ← Gv → y, we have the fact that maxq(z|Gv) I(y, z) is equivalent to minq(z|Gv) I(y,Gv|z). Also, we have (treating q(y|z) as a variational distribution)
I(y,Gv|z) = DKL(p(y|Gv, e)‖p(y|z, e)) = DKL(p(y|Gv, e)‖q(y|z))−DKL(p(y|z, e)‖q(y|z)) ≤ DKL(p(y|Gv, e)‖q(y|z)).
(15)
Based on this, we have the inequality I(y,Gv|z) ≤ min
q(y|z) DKL(p(y|Gv, e)‖q(y|z)). (16)
Also, we have (according to our definition in Eq. 6) DKL(p(y|Gv, e)‖q(y|z))
=EeEG∼pe(G)
[ 1
V ∑ v∈V Eyv∼pe(y|Gv=Gv)Ezv∼q(z|Gv=Gv) [ log pe(y = yv|Gv = Gv) q(y = yv|z = zv) ]]
≤EeEG∼pe(G)
[ 1
V ∑ v∈V Eyv∼pe(y|Gv=Gv) [ log pe(y = yv|Gv = Gv) Ezv∼q(z|Gv=Gv)[q(y = yv|z = zv)] ]] ,
(17)
where the second step is according to Jensen Inequality and the equality holds if q(z|Gv) is a delta distribution (induced by the GNN encoder h). Then the problem minq(z|Gv),q(y|z)DKL(p(y|Gv, e)‖q(y|z)) can be equivalently converted into
min f
Ee
[ 1
|Ve| ∑ v∈Ve l(f(Gev), y e v)
] = min
f Ee[L(Ge, Y e; f)]. (18)
We thus have proven that minimizing the expectation term in Eq. 4 is to minimize the upper bound of I(y,Gv|z) and contributes to maxq(z|Gv) I(y, z). Second, we have
I(y; e|z) =DKL(p(y|z, e)‖p(y|z)) =DKL(p(y|z, e)‖Ee[p(y|z, e)]) =DKL(q(y|z)‖Ee[q(y|z)])−DKL(q(y|z)‖p(y|z, e))−DKL(Ee[p(y|z, e)]‖Ee[q(y|z)]) ≤DKL(q(y|z)‖Ee[q(y|z)]). (19)
Besides, we have (according to the definition in Eq. 6)
DKL(q(y|z)‖Ee[q(y|z)])
=EeEG∼pe(G)
[ 1
|V | ∑ v∈V Eyv∼pe(y|Gv=Gv)Ezv∼q(z|Gv=Gv) [ log q(y = yv|z = zv) Ee[q(y = yv|z = zv)] ]] . (20)
Using Jensen Inequality, we will obtain that DKL(q(y|z)‖Ee[q(y|z)]) is upper bounded by
Ee[|L(Ge, Y e; f)− Ee[L(Ge, Y e; f)]|] = Ve[L(Ge, Y e; f)]. (21)
Hence we have proven that minimizing the variance term in Eq. 4 plays a role for solving minq(z|Gv) I(y; e|z).

heading: D.2 PROOFS FOR THEOREM 2
text: With Lemma 1, we know that 1) the representation z (given by GNN encoder z = h(Gv)) satisfies the invariant condition, i.e., p(y|z) = p(y|z, e) if and only if I(y; e|z) = 0 and 2) the representation z satisfies the sufficiency condition, i.e., y = c∗(z) + n if and only if z = arg maxz I(y; z).
We denote the GNN encoder that satisfies the invariance and sufficiency conditions as h∗ and the corresponding predictor model f∗(Gv) = Ey[y|h∗(Gv)] with f∗ = c∗ ◦ h∗. Since we assume the GNN encoder q(z|Gv) satisfies the conditions in Assumption 1, then according to Assumption 2, we know that there exists random variable z such that Gv = m(z, z) and p(y|z, e) would change arbitrarily across environments. Based on this, for any environment e that gives the distribution pe(y, z, z), we can construct environment e′ with the distribution pe′(y, z, z) that satisfies
pe′(y, z, z) = pe(y, z)pe′(z). (22)
Then we reproduce the proof of Theorem 2.1 in Liu et al. (2021) to prove the result by showing that for arbitrary function f = c ◦ h and environment e, there exists an environment e′ such that
EG∼pe′ (G)
[ 1
|V | ∑ v∈V Eyv∼pe′ (y|Gv=Gv)[l(f(Gv), yv)]
]
≥ EG∼pe(G)
[ 1
|V | ∑ v∈V Ey∼pe(y|Gv=Gv)[l(f ∗(Gv), yv)]
] .
(23)
Concretely we have
EG∼pe′ (G)
[ 1
|V | ∑ v∈V E(yv,zv,zv)∼pe′ (y,z,z|Gv=Gv)[l(c(zv, zv), yv)]
]
=EG′∼pe′ (G)
[ 1
|V ′| ∑ v′∈V ′ Ezv′∼pe′ (z|Gv=Gv′ )EG∼pe(G)
[ 1
|V | ∑ v∈V E(yv,zv)∼pe(y,z|Gv=Gv)[l(c(zv, zv′), yv)]
]]
≥EG′∼pe′ (G)
[ 1
|V ′| ∑ v′∈V ′ Ezv′∼pe′ (z|Gv=Gv′ )EG∼pe(G)
[ 1
|V | ∑ v∈V E(yv,zv)∼pe(y,z|Gv=Gv)[l(c ∗(zv, zv′), yv)]
]]
=EG′∼pe′ (G)
[ 1
|V ′| ∑ v′∈V ′ Ezv′∼pe′ (z|Gv=Gv′ )EG∼pe(G)
[ 1
|V | ∑ v∈V E(yv,zv)∼pe(y,z|Gv=Gv)[l(c ∗(zv), yv)]
]]
=EG∼pe(G)
[ 1
|V | ∑ v∈V E(yv,zv)∼pe(y,z|Gv=Gv)[l(c ∗(zv), yv)]
]
=EG∼pe(G)
[ 1
|V | ∑ v∈V E(yv,zv,zv)∼pe(y,z,z|Gv=Gv)[l(c ∗(zv), yv)]
] ,
(24)
where the first equality is given by Eq. 22 and the second/third steps are due to the sufficiency condition of h∗.

heading: D.3 PROOF FOR THEOREM 3
text: Recall that according to our definition in Eq. 6, the KL divergence DKL(pe(y|Gv)‖q(y|Gv)) would be
DKL(pe(y|Gv)‖q(y|Gv)) := EG∼pe(G)
[ 1
|V | ∑ v∈V Eyv∼pe(y|Gv=Gv) [ log pe(y = yv|Gv = Gv) q(y = yv|Gv = Gv) ]] .
(25) This newly defined KL divergence allows us to apply the information-theoretic framework (Federici et al., 2021) for our analysis on graph data. First, we can decompose the training error (resp. OOD error) into a representation error and a latent predictive error.
Lemma 2. For any GNN encoder q(z|Gv) and classifier q(y|z), we have
DKL(pe(y|Gv)‖q(y|Gv)) ≤ Ie(Gv;y|z) +DKL(pe(y|z)‖q(y|z)), (26)
DKL(pe′(y|Gv)‖q(y|Gv)) ≤ Ie′(Gv;y|z) +DKL(pe′(y|z)‖q(y|z)). (27)
Proof. Firstly, we have
DKL(pe(y|Gv)‖q(y|Gv))
=EG∼pe(G)
[ 1
|V | ∑ v∈V Eyv∼pe(y|Gv=Gv) [ log pe(y = yv|Gv = Gv) q(y = yv|Gv = Gv) ]]
=EG∼pe(G)
[ 1
|V | ∑ v∈V Eyv∼pe(y|Gv=Gv) [ log pe(y = yv|Gv = Gv) Ezv∼q(z|Gv=Gv)q(y = yv|z = zv) ]]
≤EG∼pe(G)
[ 1
|V | ∑ v∈V Eyv∼pe(y|Gv=Gv)Ezv∼q(z|Gv=Gv) [ log pe(y = yv|Gv = Gv) q(y = yv|z = zv) ]] =DKL(pe(y|Gv)‖q(y|z)),
(28)
where the third step is again due to Jensen Inequality and the equality holds once q(z|Gv) is a delta distribution.
Besides, we have
DKL(pe(y|Gv)‖q(y|z))
=EG∼pe(G)
[ 1
|V | ∑ v∈V Eyv∼pe(y|Gv=Gv)Ezv∼q(z|Gv=Gv) [ log pe(y = yv|Gv = Gv) q(y = yv|z = zv) ]]
=EG∼pe(G)
[ 1
|V | ∑ v∈V Eyv∼pe(y|Gv=Gv)Ezv∼q(z|Gv=Gv) [ log pe(y = yv|Gv = Gv)pe(y = yv|z = zv) pe(y = yv|z = zv)q(y = yv|z = zv) ]] =I(Gv;y|z) +DKL(pe(y|z)‖q(y|z)).
(29)
The result for DKL(pe′(y|Gv)‖q(y|Gv)) can be obtained in a similar way.
Lemma 3. For any q(z|Gv) and q(y|z), the following inequality holds for any z satisfying p(z = z|e = e) > 0, ∀e ∈ E .
DJSD(pe′(y|z)‖q(y|z)) ≤
(√ 1
2α I(y; e|z) +
√ 1
2 DKL(pe(y|z)‖q(y|z)
)2 . (30)
Proof. The proof can be adapted by from the Proposition 3 in Federici et al. (2021) by replacing e in our case with t.
The results of Lemma 2 and 3 indicate that if we aim to reduce the OOD error measured by DKL(pe′(y|Gv)‖q(y|Gv), one need to control three terms: 1) Ie(Gv;y|z), 2) DKL(pe(y|z)‖q(y|z) and 3) I(y; e|z). The next lemma unifies minimization for the first two terms.
Lemma 4. For any q(z|Gv) and q(y|z), we have
min q(z|Gv),q(y|z) DKL(pe(y|Gv)‖q(y|z))⇔ min q(z|Gv) Ie(Gv;y|z) + min q(y|z) DKL(pe(y|z)‖q(y|z)).
(31)
Proof. Recall that q(y|z) is a variational distribution. We have
Ie(Gv;y|z) = DKL(pe(y|Gv)‖pe(y|z))) = DKL(pe(y|Gv)‖q(y|z))−DKL(pe(y|z)‖q(y|z)) ≤ DKL(pe(y|Gv)‖q(y|z)).
(32)
Therefore, we can see that Ie(Gv;y|z) is upper bounded byDKL(pe(y|Gv‖q(y|z)) and the equality holds if and only if DKL(pe(y|z)‖q(y|z)) = 0. We thus conclude the proof.
Recall that according to Lemma 1 we have the fact that our objective in Eq. 4 essentially has the similar effect as
min q(z|Gv),q(y|z)
DKL(pe(y|Gv)‖q(y|z)) + I(y; e|z). (33)
Based on the Lemma 2, 3 and 4, we know that optimization for the objective Eq. 4 can reduce the upper bound of OOD error given by DKL(pe′(y|Gv)‖q(y|Gv) on condition that Ie′(Gv;y|z) = Ie(Gv;y|z). We conclude our proof for Theorem 3.

heading: E DATASETS AND EVALUATION PROTOCOLS
text: In this section, we introduce the detailed information for experimental datasets and also provide the details for our evaluation protocols including data preprocessing, dataset splits and the ways for calculating evaluation metrics. In the following subsections, we present the information for the three scenarios, respectively.
E.1 ARTIFICIAL DISTRIBUTION SHIFTS ON CORA AND AMAZON-PHOTO
Cora and Amazon-Photo are two commonly used node classification benchmarks and widely adopted for evaluating the performance of GNN designs. These datasets are of medium size with thousands of nodes. See Table 1 for more statistic information. Cora is a citation network where nodes represent papers and edges represent their citation relationship. Amazon-Photo is a copurchasing network where nodes represent goods and edges indicate that two goods are frequently bought together. In the original dataset, the available node features have strong correlation with node labels. To evaluate model’s ability for out-of-distribution generalization, we need to introduce distribution shifts into the training and testing data.
For each dataset, we use the provided node features to construct node labels and spurious environmentsensitive features. Specifically, assume the provided node features as X1. Then we adopt a randomly initialized GNN (with input of X1 and adjacency matrix) to generate node labels Y (via taking an argmax in the output layer to obtain one-hot vectors), and another randomly initialized GNN (with input of the concatenation of Y and an environment id) to generate spurious node features X2. After that, we concatenate two portions of features X = [X1, X2] as input node features for training and evaluation. In this way, we construct ten graphs with different environment id’s for each dataset. We use one graph for training, one for validation and report the classification accuracy on the remaining graphs. One may realize that this data generation is a generalized version of our motivating example in Section 3.1 and we replace the linear aggregation as a randomly initialized graph neural network to introduce non-linearity.
In fact, with our data generation, the original node features X1 can be seen as domain-invariant features that are sufficiently predictive for node labels and insensitive to different environments, while the generated features X2 are domain-variant features that are conditioned on environments. Therefore, in principle, the ideal case for the model is to identify and leverage the invariant features for prediction. In practice, there exist multiple factors that may affect model’s learning, including the local optimum and noise in data. Therefore, one may not expect the model to exactly achieve the ideal case since there also exists useful predictive information in X2 that may help the model to increase the training accuracy. Yet, through our experiments in Fig. 2(b) and 3(b), we show that the reliance of EERM on spurious features is much less than ERM, which we believe could serve as
concrete evidence that our approach is capable for guiding the GNN model to alleviate reliance on domain-variant features.

heading: E.2 CROSS-DOMAIN TRANSFERS ON MULTI-GRAPH DATA
text: A typical scenario for distribution shifts on graphs is the problem of cross-domain transfers. There are quite a few real-world situations where one has access to multiple observed graphs each of which is from a specific domain. For example, in social networks, the domains can be instantiated as where or when the networks are collected. In protein networks, there may exist observed graph data (protein-protein interactions) from distinct species which can be seen as distinct domains. In short, since most of graph data records the relational structures among a specific group of entities and the interactions/relationships among entities from different groups often have distinct characteristics, the data-generating distributions would vary across groups, which bring up domain shifts.
Yet, to enable transfer learning across graphs, the graphs in one dataset need to share the same input feature space and output space. We adopt two public datasets Twitch-Explicit and Facebook-100 that satisfy this requirement.
Twitch-Explicit contains seven networks where nodes represent Twitch users and edges represent their mutual friendships. Each network is collected from a particular region, including DE, ENGB, ES, FR, PTBR, RU and TW. These seven networks have similar sizes and different densities and maximum node degrees, as shown in Table 4. Also, in Fig. 6, we compare the ROC-AUC results on different leave-out data. We consider GCN, GAT and GCNII as the GNN backbones and train the model with standard empirical risk minimization (ERM). We further consider two ways for data splits. In the first case, which we call "OOD", we train the model on the nodes of one graph DE and report the highest ROC-AUC on the nodes of another graph ENGB. In the second case, which we call "IID", we train the model on 90% nodes of DE and evaluate the performance on the leave-out 10% data. The results in Fig. 6 show that the model performance exhibits a clear drop from "IID" to "OOD", which indicates that there indeed exist distribution shifts among different input graphs. This also serves as a justification for our evaluation protocol in Section 5.2 where we adopt the graph-level splitting to construct training/validation/testing sets.
Another dataset is Facebook-100 which consists of 100 Facebook friendship network snapshots from the year 2005, and each network contains nodes as Facebook users from a specific American university. We adopt fourteen networks in our experiments: John Hopkins, Caltech, Amherst,
Bingham, Duke, Princeton, WashU, Brandeis, Carnegie, Cornell, Yale, Penn, Brown and Texas. Recall that in Section 5.2 we use Penn, Brown and Texas for testing, Cornell and Yale for validation, and use three different combinations from the remaining graphs for training. These graphs have significantly diverse sizes, densities and degree distributions. In Fig. 7 we present a comparison which indicates that the distributions of graph structures among these graphs are different. Concretely, the testing graphs Penn and Texas are much larger (with 41554 and 31560 nodes, respectively) than training/validation graphs (most with thousands of nodes). Also, the training graphs Caltech and Amherst are much denser than other graphs in the dataset, while some graphs like Penn have nodes with very large degrees. These statistics suggest that our evaluation protocol requires the model to handle different graph structures from training/validation to testing data.

heading: E.3 TEMPORAL EVOLUTION ON DYNAMIC GRAPH DATA
text: Another common scenario is for temporal graphs that dynamically evolve as time goes by. The types of evolution can be generally divided into two categories. In the first case, there are multiple graph snapshots and each snapshot is taken at one time. As time goes by, there exists a sequence of graph snapshots which may contain different node sets and data distributions. Typical examples include financial networks that record the payment flows among transactions within different time intervals. In the second case, there is one graph that evolves with node/edge adding or deleting. Typical examples include some large-scale real-world graphs like social networks and citation networks where the distribution for node features, edges and labels would have strong correlation with time (in different scales). We adopt two public real-world datasets Elliptic and OGB-Arxiv for node classification experiments.
Elliptic contains a sequence of 49 graph snapshots. Each graph snapshot is a network of Bitcoin transactions where each node represents one transaction and each edge indicates a payment flow. Approximately 20% of the transactions are marked with licit or illicit ones and the goal is to identify illicit transaction in the future observed network. Since in the original dataset, the first six snapshots have extremely imbalanced classes (where the illicit transactions are less than 10 among thousands of nodes), we remove them and use the 7th-11th/12th-17th/17th-49th snapshots for training/validation/testing. Also, due to the fact that each graph snapshot has very low positive label rate, we group the 33 testing graph snapshots into 9 test sets according to the chronological order. In Fig. 8 we present the label rate and positive label rate for training/validation/testing sets. As we can see, the positive label rates are quite different in different data sets. Indeed, the model needs to handle distinct label distributions from training to testing data.
OGB-Arxiv is composed of 169,343 Arxiv CS papers from 40 subject areas and their citation relationship. The goal is to predict a paper’s subject area. In (Hu et al., 2020b), the papers published before 2017, on 2018 and since 2019 are used for training/validation/testing. Also, the authors adopt the transductive learning setting, i.e., the nodes in validation and test sets also exist in the graph for training. In our case, we instead adopt inductive learning setting where the nodes in validation and test sets are unseen during training, which is more akin to the real-world situation. Besides, for better evaluation on generalization, especially extrapolating to new data, we consider dataset splits with a larger year gap: we use papers published before 2011 for training, from 2011 to 2014 for validation, and after 2014 for test. Such a dataset splitting way would introduce distribution shift between
training and testing data, since several latent influential factors (e.g., the popularity of research topics) for data generation would change over time. In Fig. 9, we visualize the T-SNE embeddings of the nodes and mark the training/validation/testing nodes with different colors. From Fig. 9(a) to Fig. 9(c), we can see that testing nodes non-overlapped with the training/validation ones exhibit an increase, which suggests that the distribution shifts enlarge as time difference goes large. This phenomenon echoes the results we achieve in Table 3 where we observe that as the time difference between testing and training data goes larger, model performance suffers a clear drop, with ERM suffering more than EERM.
F IMPLEMENTATION DETAILS
In this section, we present the details for our implementation in Section 5 including the model architectures, hyper-parameter settings and training details in order for reproducibility. Most of our experiments are run on GeForce RTX 2080Ti with 11GB except some experiments requiring large GPU memory for which we adopt RTX 8000 with 48GB. The configurations of our environments and packages are listed below:
• Ubuntu 16.04
• CUDA 10.2
• PYTHON 3.7
• Numpy 1.20.3
• PyTorch 1.9.0
• PyTorch Geometric 1.7.2

heading: F.1 MODEL ARCHITECTURES
text: In our experiments in Section 5, we adopt different GNN architectures as the backbone. Here we introduce the details for them.
GCN. We use the GCNConv available in Pytorch Geometric for implementation. The detailed architecture description is as below:
• A sequence of L-layer GCNConv.
• Add self-loop and use batch normalization for graph convolution in each layer.
• Use ReLU as the activation.
GAT. We use the GATConv available in Pytorch Geometric for implementation. The detailed architecture description is as below:
• A sequence of L-layer GATConv with head number H .
• Add self-loop and use batch normalization for graph convolution in each layer.
• Use ELU as the activation.
GraphSAGE. We use the SAGEConv available in Pytorch Geometric for implementation. The detailed architecture description is as below:
• A sequence of L-layer SAGEConv.
• Add self-loop and use batch normalization for graph convolution in each layer.
• Use ReLU as the activation.
GCNII. We use the implementation4 provided by the original paper (Chen et al., 2020a). The associated hyper-parameters in GCNII model are set as: αGCNII = 0.1 and λGCNII = 1.0.
GPRGNN. We use the implementation5 provided by Chien et al. (2021). We adopt the PPR initialization and GPRprop as the propagation unit. The associated hyper-parameters in GPRGNN model are set as: αGPRGNN = 0.1.

heading: F.2 HYPER PARAMETER SETTINGS
text: The hyper-parameters for model architectures are set as default values in different cases. Other hyperparameters are searched with grid search on validation dataset. The searching space are as follows: learning rate for GNN backbone αf ∈ {0.0001, 0.0002, 0.001, 0.005, 0.01}, learning rate for graph editers αg ∈ {0.0001, 0.001, 0.005, 0.01}, weight for combination β ∈ {0.2, 0.5, 1.0, 2.0, 3.0}, number of edge editing for each node s ∈ {1, 5, 10}, number of iterations for inner update before one-step outer update T ∈ {1, 5}.

heading: F.2.1 SETTINGS FOR SECTION 5.1
text: We consider 2-layer GCN with hidden size 32. We use weight decay with coefficient set as 1e-3. Besides, we set αg = 0.005, αf = 0.01, β = 2.0, s = 5, T = 1.
4https://github.com/chennnM/GCNII 5https://github.com/jianhao2016/GPRGNN

heading: F.2.2 SETTINGS FOR SECTION 5.2
text: For GCN, we set the layer number L as 2. For GAT, we set L = 2 and H = 4. For GCNII, we set the layer number as 10. We use hidden size 32 and weight decay with coefficient set as 1e-3.
For Twitch-Explicit, other hyper-parameters are set as follows:
• GCN: αg = 0.001, αf = 0.01, β = 3.0, s = 5, T = 1. • GAT: αg = 0.005, αf = 0.01, β = 1.0, s = 5, T = 1. • GCNII: αg = 0.01, αf = 0.001, β = 1.0, s = 5, T = 1.
For Facebook-100, other hyper-parameters are set as: αg = 0.005, αf = 0.01, β = 1.0, s = 5, T = 1.

heading: F.2.3 SETTINGS FOR SECTION 5.3
text: For GraphSAGE and GPRGNN, we set the layer number as 5 and hidden size as 32.
For Elliptic, other hyper-parameters are set as follows:
• GraphSAGE: αg = 0.0001, αf = 0.0002, β = 1.0, s = 5, T = 1. • GPRGNN: αg = 0.005, αf = 0.01, β = 1.0, s = 5, T = 1.
For OGB-Arxiv, other hyper-parameters are set as follows:
• GraphSAGE: αg = 0.01, αf = 0.005, β = 0.5, s = 1, T = 5. • GPRGNN: αg = 0.001, αf = 0.01, β = 1.0, s = 1, T = 5.

heading: F.3 TRAINING DETAILS
text: For each method, we train the model with a fixed number of epochs and report the test result achieved at the epoch when the model provides the best performance on validation set.

heading: G MORE EXPERIMENT RESULTS
text: We provide additional experiment results in this section. In Fig.10 and 11 we present the distribution of test accuracy on Cora when using SGC and GAT, respectively, as the GNNs for data generation. In Fig. 12 and 13 we further compare with the training accuracy using all the features and removing the spurious ones for inference. These results are consistent with those presented in Section 5.1, which again verifies the effectiveness of our approach. Besides, the corresponding extra results on Photo are shown in Fig. 14, 15, 16 and 17, which also back up our discussions in Section 5.1.

Review summaries:
- This paper tackles the out-of-distribution problem for node-level prediction on graphs from the invariance perspective. It presents a novel approach in which the whole graph is divided into n ego-graphs where n is the number of nodes. All the ego-graphs can be treated as a set of IID. Based on this division, the model can be trained to defend the adversarial attack from multiple environments.  - This work extends the discussion of the OOD problem to node-level tasks on graphs. A new learning approach is proposed and theoretically proven to be correct. In empirical experiments on multiple datasets, the approach (EERM) outperforms its counterpart baseline ERM. 

The paper adapts a recently proposed invariant risk minimization approach for tackling distribution shift on node-level predictions on graphs. The main idea is to decompose the graph into a set of ego-graphs rooted at each node, thus incorporating structural information. Since the learning objective requires data from different environments they authors introduce auxiliary context generators trained to maximize the variance loss. 

This paper studies the problem of distribution shifts as out-of-distribution generalization. Specifically, it formulates the OOD problem as invariant risk minimization under different environments. The relation between these two has been extensively discussed in the paper. Multiple environment is done by graph editing using policy gradient. Experiments on three different kind of distribution shifts are presented and the effectiveness of EERM is validated.

The importance of out-of-domain (OOD) generalization has emerged, and there is much research for out-of-domain generalization [1,2,3,4]. However, the related works on the graph-structured datasets are not explored well. OOD generalization of the graph is not trivial because the graph has the interconnection among nodes and the existence of structural information. To solve the problems, this paper proposes a new method, multiple contexts explore that are adversarially trained to maximize the variance of risk. The proposed model is validated on many diverse datasets, Cora, Amazon-Photo, Twitch-explicit, and so on.  [1] Arjovsky, Martin, et al. "Invariant risk minimization." arXiv preprint arXiv:1907.02893 (2019).  [2] Sagawa, Shiori, et al. "Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization." arXiv preprint arXiv:1911.08731 (2019).  [3] Krueger, David, et al. "Out-of-distribution generalization via risk extrapolation (rex)." International Conference on Machine Learning. PMLR, 2021.  [4] Creager, Elliot, Jörn-Henrik Jacobsen, and Richard Zemel. "Environment inference for invariant learning." International Conference on Machine Learning. PMLR, 2021.