Invariant Node Representation Learning under Distribution Shifts with Multiple Latent Environments

Consider a graph \(G = (V, E)\), the node feature matrix \(X = \lbrace x_v | v \in V \rbrace \in \mathbb {R}^{|V| \times F}\) (where \(F\) denotes the node feature dimension), and labels \(Y = \lbrace y_v | v \in V\rbrace\). The adjacency matrix is denoted as \(A=\lbrace a_{v,v^\prime } | v,v^\prime \in V\rbrace \in \mathbb {R}^{|V| \times |V|}\), where \(a_{v,v^\prime }=1\) means there exists an edge connecting node \(v\) and \(v^\prime\), and \(a_{v,v^\prime }=0\) otherwise. We assume the nodes \(V\) are collected from multiple environments, i.e., \(V = \lbrace V^e\rbrace _{e \in {\rm supp}(\mathcal {E}_{tr})}\), where \(V^e\) denotes the nodes from environment \(e\), \(\mathrm{supp}(\mathcal {E}_{tr})\) is the support of the environmental variable. We use \(\mathbf {v}\) and \(\mathbf {y}\) to denote the random variables of node and label, respectively. We summarize the key notations of this article and the corresponding descriptions in Table 1.

Recently, invariant learning has received surging attention to enable generalizing to distribution shifts, i.e., out-of-distribution (OOD)generalization. It aims to exploit the invariant relationships between the input data and labels across distribution shifts while filtering out the variant spurious correlations.¹ Following the invariant learning literature [2, 4, 11, 40, 42, 64], we formulate the problem of learning invariant node representations capable of generalizing to distribution shifts, i.e., out-of-distribution (OOD) generalized node representation learning, as:

Note that \(\mathrm{supp}(\mathcal {E}_{tr}) \subset \mathrm{supp}(\mathcal {E})\). Distribution shifts indicate that \(P^e(\mathbf {v},\mathbf {y}) \ne P^{e^{\prime }}(\mathbf {v},\mathbf {y}), e \in \mathrm{supp}(\mathcal {E}_{tr}), e^{\prime } \in \mathrm{supp}(\mathcal {E}) \setminus \mathrm{supp}(\mathcal {E}_{tr})\), i.e., the joint distribution of node and label is different in training and testing data. The testing nodes are not available in the training stage, meaning that we can not obtain a prior distribution of testing nodes for training.² However, Problem 1 is difficult to be directly solved, since (1) the nodes which are connected by graph structure are non-independent inducing obstacle for predictions, and (2) the environment labels for the nodes are unobserved [4, 40], which are usually unavailable or prohibitively expensive to collect for most real scenarios.

In this article, we focus on learning invariant node representation by adopting message-passing GNNs. Since only the immediate neighbors of nodes are aggregated in each message-passing layer, the representation of nodes only depends on their \(L\)-hop neighbors, where \(L\) is the number of message-passing layers. Therefore, we learn representations of nodes by only focusing on their \(L\)-order ego-graph, which is the common assumption for most message-passing GNNs [34, 38, 78]. Denote the node \(v\)’s \(L\)-hop neighbors as \(\mathcal {N}_v = \lbrace u|d(v,u) \le L\rbrace\), where \(d(v,u)\) is the shortest path distance between node \(v\) and \(u\). The nodes in \(\mathcal {N}_v\) and their connections form the ego-graph \(G_v\) of node \(v\), which is represented as a local node feature matrix \(X_v = \lbrace x_u | u \in \mathcal {N}_v\rbrace\) and local adjacency matrix \(A_v = \lbrace a_{u,u^\prime } | u,u^\prime \in \mathcal {N}_v\rbrace\). We use \(\mathbf {G_v}\) and \(G_v\) to denote the random variable and instance of ego-graphs and use \(\mathbf {G}\) and \(\mathbf {Y}\) to denote the random variable of input graph and node label vector, respectively. Then, we can reformulate the problem by using ego-subgraphs, i.e., an ego-graph dataset defined as \(\mathcal {G} = \lbrace \mathcal {G}^e\rbrace _{e \in {\rm supp}(\mathcal {E}_{tr})}\), where \(\mathcal {G}^e = \lbrace (G^e_v, y^e_v) | v \in V^e\rbrace\) denotes the ego-graphs in environment \(e\). Notice that ego-graphs are not independent samples, but they can be seen as a Markov blanket [34, 78], so the conditional distribution can be decomposed (conditional independence), i.e., \(P(\mathbf {Y}|\mathbf {G}) = \prod _{\mathbf {v}} P(\mathbf {y}|\mathbf {G_v})\).

To enable OOD generalization, recent studies on invariant learning [2, 4, 11, 40, 42, 64] propose to train a predictor using only a portion of features of each input instance that capture the invariant and sufficiently predictive relations with labels. Since we have transformed the node representation learning task into only using ego-graphs \(G_v\), we assume that each ego-graph instance has an invariant subgraph, i.e., ego-subgraph \(G_v^{I} \subset G_v\), that possesses invariant and sufficiently predictive information to the node’s label \(y_v\) in different environments under distribution shifts. We refer to the rest of each ego-graph, i.e., the complement of \(G_v^I\), as the variant ego-subgraph and denote it as \(G_v^S\). \(G_v^S\) represents the surrounding part of the node \(v\) whose relationship with the label is variant across different environments, e.g., spurious correlations for predicting node \(v\). The graph model will have a better OOD generalization ability if it can identify the invariant ego-subgraph \(G_v^I\) for each node accurately and learn node representation based on \(G_v^I\) for predictions.

Formally, we denote a generator for each node’s ego-graph to obtain the invariant ego-subgraph as \(G_v^I = \Psi (G_v)\). Following the invariant learning literature [2, 11, 40, 42, 45, 50], we make the assumption:

The invariance assumption means that the node representations learned on invariant ego-subgraphs have an invariant relation to the node labels across different environments. The sufficiency assumption means that the node representations learned on invariant ego-subgraphs are sufficiently predictive to the node labels.

In this article, we instantiate \(\Psi (\cdot)\) using two learnable masks on node features and graph structures (i.e., edges). First, the edge mask is responsible for splitting the local adjacency matrix \(A_v\) of the ego-graph \(G_v\) into the local adjacency matrix \(A_v^I\) of the invariant ego-subgraph \(G_v^I\) and the local adjacency matrix \(A_v^S\) of the variant ego-subgraph \(G_v^S\). A straightforward strategy is to train a binary mask matrix \(M^{A_v} = \lbrace 0, 1\rbrace ^{|\mathcal {N}_v| \times |\mathcal {N}_v|}\) on the local adjacency matrix \(A_v\). However, directly optimizing such a mask matrix is a discrete optimization problem and intractable in practice, especially for large-scale graphs [88]. Besides, learning a mask for each ego-subgraph cannot share knowledge among different nodes. Therefore, we adopt a learnable GNN (denoted as \(\mathrm{GNN^{M}}\)) to parameterize the mask matrix. Specifically, we relax edge masks from binary variables to continuous variables in \([0,1]\). The soft mask for each edge \((u,u^\prime), u,u^\prime \in \mathcal {N}_v\) in ego-graph \(G_v\) is:

Besides the edge mask, we also adopt a soft \(F\)-dimensional feature mask \(M^X \in [0,1]^{F}\) shared by all the nodes for selecting the invariant node features in the ego-graph \(G_v\). The invariant ego-subgraph \(G_v^I = (A_v^I, X_v^I)\) and variant ego-subgraph \(G_v^S = (A_v^S, X_v^S)\) of \(G_v\) are calculated as:where \(\odot\) is the element-wise matrix multiplication. Using the above method, we can generate all the invariant ego-subgraphs \(\lbrace G_v^I|v \in V\rbrace\) and variant ego-subgraphs \(\lbrace G_v^S|v \in V\rbrace\).

After splitting the nodes’ ego-graphs into invariant and variant subgraphs, we can infer the environment label \(\mathcal {E}_{infer}\) using variant subgraphs \(\lbrace G_v^S|v \in V\rbrace\). The intuition is that, since the invariant ego-subgraphs capture the invariant relationships between predictive node features and graph structures with the node labels, the variant ego-subgraphs in turn capture variant spurious correlations under different distributions. Consider two nodes \(v, v^\prime\) from the same environment (e.g., two proteins from the same species or two papers published in the same period). Their variant ego-subgraphs \(G_v^S\) and \(G_{v^\prime }^s\) are likely show similar environment patterns. Based on the graph homophily assumption [57] that similar nodes are more likely to connect to each other, the nodes from the same environment will tend to be more densely connected in their variant ego-subgraphs than nodes from different environments (an illustrating example is shown in Figure 1). Therefore, we can infer the environments by conducting graph clustering based on the variant node features and edges.

Specifically, let \(X^S\) and \(A^S\) denote the node features and edges in \(\lbrace G_v^S|v \in V\rbrace\). Assuming there are \(K\) latent environments in graph, we design a contrastive modularity-based clustering method to infer the environments by learning a cluster assignment matrix \(C = \lbrace C_v|v \in V\rbrace\), where \(C_v\) is \(K\)-dimensional one-hot vector indicating the environment of node \(v\). We propose to minimize the following contrastive objective for clustering the nodes denoted by \((X^S, A^S)\):whereIn Equation (5), \(\mathbf {d}\) and \(m\) indicate the degree vector and the number of edges calculated by \(A^S\), respectively. \({\rm diag}(\cdot)\) means only keeping the diagonal elements of the input matrix. \(B \in \mathbb {R}^{K \times K}\) is the modularity matrix [60], whose entry \(B_{k,k\prime }\) is the probability of an edge existing between cluster \(k\) and \(k^\prime\). Optimizing Equation (4) can maximize the connection probability between nodes from the same clusters (i.e., positive pairs) and minimize the connecting probability between nodes from the different clusters (i.e., negative pairs) via a contrastive scheme [13], encouraging to form clear clusters. Since optimizing the binary cluster assignment matrix is proven to be NP-hard [8], we follow Reference [73] to relax \(C \in [0,1]^{|V| \times K}\) as a soft cluster assignment and adopt a GNN to calculate the assignment matrix, i.e., \(C = {\rm Softmax}({\rm GNN}^{\rm C}(X^S, A^S)).\) Finally, the optimal cluster assignment \(C^*\) can be used to indicate the inferred environments \(\mathcal {E}_{infer}\) of nodes.

After obtaining the inferred invariant ego-subgraphs \(\lbrace G_v^I | v \in V\rbrace\) and environment labels \(\mathcal {E}_{infer}\), we propose the invariance regularization module, which can make the graph model to generate node representations capable of OOD generalization under distribution shifts. Specifically, we aim to learn the optimal generator \(\Psi ^*\) in Assumption 1 by proposing and optimizing the maximal invariant ego-subgraph generator criterion. Following the invariant learning literature [11, 40, 50, 51], we give the following definition:

Then, we show that the optimal generator \(\Psi ^*\) satisfies the following theorem:

Theorem 1 provides us an objective function to optimize the invariant ego-subgraph generator. However, directly solving according to Theorem 1 for a non-linear \(\Psi\) is difficult [40]. Following the invariant learning literature [40], we minimize the following invariance regularizer:where \(f(\cdot) = w \circ g \circ \Psi\), \(\mathcal {E}_{infer}\) is the inferred environment label, and \(\theta\) denotes all the learnable parameters. Recall that \(g(\cdot)\) is the representation learning function of the invariant ego-subgraphs and \(w(\cdot)\) is the classifier. We instantiate \(g\) as another GNN as: \(\mathbf {Z}_{I} = \mathrm{GNN^{I}}(G_v^I)\), where \(\mathbf {Z}_{I}\) are the node representations capturing invariant patterns from the ego-subgraphs. \(w(\cdot)\) is instantiated as a multilayer perceptron with the ReLU [1] activation function, followed by the softmax function. By optimizing Equation (8), we can get our desired generator \(\Psi\) and the ego-subgraph representation learning function \(g(\cdot)\), which collectively serve as our representation learning method \(h(\cdot)\), i.e., \(h = g \circ \Psi\).

We further theoretically analyze our INL model by showing that the maximal invariant ego-subgraph generator can achieve OOD optimality.

Intuitively, Theorem 2 shows that we can transform the OOD generalization problem into finding the optimal invariant ego-subgraphs while maintaining the optimality. The proofs of the above theorems are inspired by the invariant learning literature [45, 50, 51, 78], and a motivating example for better understanding is provided in Section 3.6. It indicates that our method can get rid of spurious correlations and learn OOD generalized node representations based on the identified invariant ego-subgraphs.

We present the pseudocode of INL in Algorithm 1 to show the training procedure. Specifically, we first obtain the invariant and variant ego-subgraphs for all nodes with the learnable masks on node features and edges. Then, we infer the environments for all nodes with the variant node features and edges from variant ego-subgraphs. And we learn the invariant node representations with invariance regularization based on the inferred invariant ego-subgraphs and environment labels. Note that the adopted GNNs including \(\mathrm{GNN^{M}}\), \(\mathrm{GNN^{C}}\), and \(\mathrm{GNN^{I}}\) for all ego-graphs are shared, following References [34, 78]. At the testing stage, we directly adopt the optimized \(f\) to conduct predictions. In Algorithm 1, “Epoch” means the overall number of epochs for optimizing the proposed method, and “Epoch_Cluster” denotes the number of epochs for clustering to infer environments in each training epoch. The setting details of the hyper-parameters can be found in Section 4.1.3.

For better understanding our proposed method intuitively, we present a linear toy example and the corresponding theoretical analysis inspired by Reference [78] to show why our method can achieve out-of-distribution generalization by learning node representations based on invariant ego-subgraph \(G_v^I\) (i.e., invariant node features \(X_v^I\) and structures \(A_v^I\)).

For simplification, in this toy example, we consider the ego-graph \(G_v\) (and \(\mathcal {N}_v\)) only contains the centered node \(v\) and its 1-hop neighbors (i.e., \(L=1\)), which can be split into invariant ego-subgraph \(G_v^I\) (and \(\mathcal {N}_v^I\)) and variant ego-subgraph \(G_v^S\) (and \(\mathcal {N}_v^S\)). And we consider the dimensionality of node features \(F=2\), including one-dimensional invariant node feature \(x_v^I\) and variant node feature \(x_v^S\), i.e., \(x_v = [x_v^I, x_v^S]\) for each node \(v\). The illustration of ego-graph \(G_v\) is shown in Figure 2. The dependence among variables in the toy example is shown in Figure 3. We do not distinguish the notation of random variables and of their particular instances when there is no risk of confusion in this toy example.

Considering the representation learning function \(g^*\) that averages the node representations in invariant ego-subgraph \(G_v^I\) to produce the centered node representations and classifier \(w^*\) is identity mappings in Assumption 1, the node label can be determined by the invariant node features and structures:where \(\epsilon _1\) is standard normal noise. And we assume that the variant node feature \(x_v^S\) is generated by identity mapping given the input of the node’s label \(y_v\) and environment \(e_v\), which can be denoted as:where \(\epsilon _2\) is standard normal noise. \(e_v\) denotes the node \(v\)’s environment, following normal distribution whose mean and variance are dependent on node environment. Besides, we assume the variant structures are also dependent on the node environment and the environments of nodes in \(\mathcal {N}_v^S\) is \(e_v\). For example, in citation networks, the invariant node features and structures can be the paper published avenues and citations among them that determine the subject topics (i.e., labels), while the variant node features and structures can be the citation indexes and edges between papers with high citations in some publication periods (i.e., environments).

Therefore, given the invariant and variant ego-subgraph, we consider the following predictor model:Note that the ideal solution for the predictor model is \(\theta = [\theta _1, \theta _2, \theta _3, \theta _4] = [1, 0, 0, 0]\), indicating that the predictor accurately identifies the sufficiently predictive and invariant node features and structures for making OOD generalized predictions. However, the following proposition shows that we cannot obtain this ideal solution if only using standard empirical risk minimization (ERM):

The proof is in Appendix A.1. Proposition 3 indicates directly optimizing with ERM will inevitably make the predictor model heavily rely on spurious correlations, since \(\theta _2, \theta _3, \theta _4\) is not constant zero, so that the model performs poorly under distribution shifts with multiple latent environments. Next, we show that our objective in Equation (8) can mitigate this issue.

The proof is in Appendix A.2. Proposition 4 indicates our method can get rid of spurious correlations and learn OOD generalized node representations under distribution shifts with multiple latent environments by generating node representations based on the identified invariant ego-subgraph \(G_v^I\).

Intuitively, Proposition 3 shows that the optimal solution under standard empirical risk minimization (ERM) in this toy example (as shown in Figure 2) consists of non-zero coefficients of the predictor model for variant ego-subgraph, which means that the predictions rely on variant environment information, e.g., different species that the proteins come from in protein-protein interaction graphs and the publication time of papers in citation networks. Therefore, the OOD generalization performance is poor. However, Proposition 4 shows that the optimal solution using the proposed method in this toy example only includes non-zero coefficients of the predictor model for invariant ego-subgraph, demonstrating that our method can make predictions only based on the invariant information and is not affected by variant spurious correlations, leading to strong OOD generalization ability.

The experimental results are shown in Table 5, from which we have the following observations: Our proposed INL consistently and significantly outperforms the baselines and achieves the best performance in all settings. The results demonstrate the effectiveness of our proposed method in handling distribution shifts, which has a remarkable out-of-distribution generalization ability. The general invariant learning methods, e.g., IRM, GroupDRO, V-REx, only have slight improvements to ERM. EERM is a recently proposed invariant method specifically designed for learning node representations but assumes a single environment is shared for all the nodes. EERM outputs competitive results in some settings but fails to obtain consistent improvements, indicating modeling multiple latent environments is crucial for handling distribution shifts in graph. GIL achieves promising gains over the other baselines, but the proposed method still performs better than it.

In addition, when \(r_{train} = 1/3\), i.e., no distribution shifts between training and testing data, our proposed method also achieves the best results, meaning that learning invariant ego-subgraphs for nodes is also beneficial. As \(r_{train}\) grows larger, the performance of all the methods tends to decrease, since there exists a larger degree of distribution shift. Nevertheless, our proposed method is able to maintain the most relatively stable performance. In fact, the performance gap between INL and the best results of baselines becomes more significant as the degree of distribution shift increases. For example, the accuracy improvements against the strongest baselines increases from 4.77% to 9.91% when \(r_{train}\) changes from \(1/3\) to 0.7 on Citeseer, indicating the powerful OOD generalization ability of our method under various complex distribution shifts.

To further analyze whether our method can accurately capture the invariant ego-subgraphs under distribution shifts, we compare the output invariant node features and structures with the ground-truth on the synthetic dataset Citeseer. The evaluation metric is ROC-AUC. The results in Figure 4 show that the ROC-AUC for discovering invariant node features and structures is around 70% and 80%, respectively, which is significantly higher than random selection (50%). It demonstrates our INL can discover the truly predictive invariant ego-subgraphs and further make OOD generalized predictions.

We further evaluate the effectiveness of our method on two real-world graph datasets, i.e., OGB-Arxiv and OGB-Proteins from OGB [33]. The properties of citation networks can change significantly in different time ranges. So, the node distribution shifts on OGB-Arxiv are introduced by selecting papers published before 2011 as training set, within 2011–2014 as validation set, and within 2014–2016/2016–2018/2018–2020 as testing sets. For OGB-Proteins dataset, since the interactions between proteins can vary from different species that the proteins come from, we split the protein nodes into training/validation/test sets according to their species. We assume the test nodes are strictly unseen during training stage, which is more common in practice and more challenging than the default setting of OGB [33].

The experimental results are presented in Table 6. Our proposed method consistently achieves the best performance, indicating that INL can well handle distribution shifts existing in real-world scenarios. For example, INL increases the classification accuracy by 3.41% on OGB-Arxiv (tested on 2016–2018 with GraphSAGE backbone) and ROC-AUC by 7.54% on OGB-Proteins (tested on species-1 with GAT backbone) against the strongest baselines, respectively. Besides, different datasets have different distribution shifts, and none of the baselines can consistently achieve promising OOD generalized performance as our method. Therefore, the results show that our proposed method can well handle diverse types of distribution shifts in real graph datasets.

Besides the quantitative evaluation, we plot a showcase from the OGB-Arxiv to intuitively validate the effectiveness of our method. Figure 5 shows that the learned invariant ego-subgraph \(G_v^I\) (denoted by solid lines) and variant ego-subgraph \(G_v^S\) (denoted by dashed lines) of one node \(v\) (ID: 139,332). We plot the top-five selected edges by the masks for simplicity. It can be observed that the invariant ego-subgraph \(G_v^I\) learned by our method accurately corresponds to the neighbors in the ego-graph from the same subject area (i.e., artificial intelligence), which have truly predictive and invariant relations with the centered node. However, the variant ego-subgraph \(G_v^S\) highlights the neighbors that are from different subject areas that are published in the same year with the centered node and have a high citation index (spurious feature). Besides, there is another paper \(u\) whose subject area is information retrieval (IR) that also cites those papers with high citation indexes, meaning that the node \(u\) has similar variant patterns with node \(v\) so they are in the same environment. We can observe that these nodes form clear cluster structures based on the variant ego-subgraphs, demonstrating the effectiveness of the proposed graph clustering algorithm in inferring latent environments.

In our proposed model, all components are jointly optimized. To show that the node environment inference module and invariance regularization module can mutually promote each other, we record the test accuracy, the modularity, which is a measurement for the quality of graph cluster, and the normalized mutual information (NMI) [41], which is another metric (falling within the range \([0, 1]\)) for evaluating the clustering accuracy, as the model is trained. The results on Citeseer (\(r_{train} = 0.7\)) are shown in Figure 6. We can observe that the test accuracy and the modularity (clustering properties) improve synchronously over training. The results show that, as the training stage progresses, the invariant ego-subgraph generator is optimized so it can generate more informative invariant ego-subgraphs and therefore improve the performance on the testing set. However, accurately discovering invariant ego-subgraphs can also promote identifying variant ego-subgraphs, which capture the environment-discriminate features and better infer the latent environments. In addition, we observe that the test accuracy and the NMI (clustering accuracy) also improve collectively over training. Notice that INL achieves such results without needing any ground-truth environment label.

These empirical results well support the following points: (1) The invariant and variant patterns widely exist in real-world graphs, and our proposed INL can well identify invariant/ variant ego-subgraphs under distribution shifts with multiple latent environments. (2) The variant ego-subgraphs form clear clustering structures, and our INL can capture such patterns to accurately infer the environment labels of nodes. (3) Based on the inferred environments, our INL learns node representations by the invariant ego-subgraph for each node so it can achieve better OOD generalization performance. The environment inference and invariance regularization module can mutually enhance each other.

We perform ablation studies over the key components of the invariant ego-subgraph generator \(\Psi\), i.e., masks on node features and edges, to understand their functionalities more deeply. We compare INL with the following two ablated versions: (1) w/o node feature mask: It removes the node feature mask by setting both invariant and variant node features in the ego-graph \(G_v\) to \(X_v\), i.e., \(X_v^I=X_v^S=X_v\). (2) w/o edge mask: It removes the edge mask by setting both invariant and variant edges in the ego-graph \(G_v\) to \(A_v\), i.e., \(A_v^I=A_v^S=A_v\). The results of the two ablated versions drop compared with INL, as shown in Figure 7. The performance gaps between INL and the two ablated versions become more significant as the degree of distribution shift increases (i.e., \(r_{train}\) from 1/3 to 0.7), which demonstrates the significance of accurately identifying the invariant node features and edges by the learnable masks.

We can observe the convergence of our proposed method empirically, although the clustering objective in environment inference (i.e., Equation (4)) and invariance objective in invariance regularization (i.e., Equation (8)) are iteratively optimized. In Figures 8(a) and (b), we show the two objectives in the training process on Citeseer (\(r_{train}=0.7\)) and OGB-Arxiv, respectively. The loss converges before reaching the maximal training epoch, while the results on the other datasets show similar patterns.

In Figure 9, we also show the objective of the inner iteration in Algorithm 1, i.e., the training dynamics of the clustering objective in one epoch of the outer iteration. The epoch of the outer iteration is specified as 100 and 250 for Citeseer (\(r_{train}=0.7\)) and OGB-Arxiv, respectively, which is the middle of the whole training process, while the results in other epochs of the outer iteration show similar patterns.

The time complexity of the proposed INL is \(O(\left|E\right|d+\left|V\right|d^2)\), where \(\left|V\right|\) and \(\left|E\right|\) denote the number of nodes and edges, respectively, and \(d\) is the dimensionality of the node representations. Specifically, we adopt the message-passing GNN, which has a complexity of \(O(\left|E\right|d+\left|V\right|d^2),\) to instantiate the GNN components in INL, and the GNNs are shared for all ego-graphs. Since we only need to generate mask for the existing edges in graphs, the time complexity of generating invariant and variant ego-subgraphs and further obtaining their representations is \(O(\left|E\right|d+\left|V\right|d^2)\). The time complexity of calculating the modularity matrix \(B\) in environment inference is \(O(\left|E\right|(d+|\mathcal {E}_{infer}|)+\left|V\right|(d+|\mathcal {E}_{infer}|)^2)\), where \(|\mathcal {E}_{infer}|\) denotes the number of inferred environments. The time complexity of the invariance regularizer is \(O(|\mathcal {E}_{infer}|d^2)\), as the number of parameters for most GNNs is \(O(d^2)\). Since \(|\mathcal {E}_{infer}|\) are small constants, the overall time complexity of INL is \(O(\left|E\right|d+\left|V\right|d^2)\). In comparison, the time complexity of other GNN-based node representation methods is also \(O(\left|E\right|d+\left|V\right|d^2)\). Therefore, the time complexity of our proposed INL is on par with the existing methods.

In addition to the analysis of the time complexity, the empirical time cost of the proposed method and baselines are also tested. We show the results on Citeseer (\(r_{train}=0.7\)) in Figure 10 while the results on other datasets show similar patterns. The results indicate that INL does not introduce infeasible time cost for achieving the best performances in practice. Its time cost for each training epoch is comparable with the baselines and more efficient than some competitive methods, demonstrating the efficiency and effectiveness of our method.

We compare the output invariant node features and structures generated by the proposed INL and GNNExplainer [88] with the ground-truth on the synthetic dataset Citeseer. Specifically, we generate post hoc explanations from GNNExplainer as the identified invariant ego-subgraphs, where we use the models trained under ERM as the models to explain. The evaluation metric is ROC-AUC. The results in Table 7 show that the masks on invariant node features and edges generated by GNNExplainer can be easily affected by the spurious correlations. Moreover, even when spurious correlations do not exist, the ROC-AUC of masks on invariant node features and edges generated by our INL still outperforms the result of the explainability method GNNExplainer, showing the effectiveness of INL when identifying invariant patterns.

We investigate the sensitivity of hyper-parameters of our method, including the number of inferred environments \(|\mathcal {E}_{infer}|\), the invariance regularizer coefficient \(\lambda\), and the number of epochs for clustering to infer environments in each training epoch (i.e., Epoch_Cluster in Algorithm 1). For simplicity, we only report the results on Citeseer (\(r_{train}=0.7\)) and OGB-Arxiv (2016–2018 with GraphSage backbone) in Figures 11–13, while the results on other datasets show similar patterns.

First, the number of inferred environments has a slight impact on the model performance. For Citeseer, the performance reaches a peak when \(|\mathcal {E}_{infer}|=3\), showing that INL achieves the best result when the number of environments matches the ground truth. For OGB-Arxiv, the best number of environments is \(|\mathcal {E}_{infer}|=5\). A plausible reason is that OGB-Arxiv dataset consists of more nodes and edges, which form more environments than Citeseer. Second, we also find the coefficient \(\lambda\) has a slight influence on the performance, indicating that we need to properly balance the classification loss and the invariance regularizer term. Finally, a proper value of the hyper-parameter Epoch_Cluster is important. A small value may not be sufficient to infer the environments accurately, while a very large value is unnecessary and may affect the training efficiency. Although an appropriate choice of hyper-parameters can further improve the performance, our method is not very sensitive to hyper-parameters. Figures 11–13 show that INL can outperform the best baselines with a wide range of hyper-parameters choices.

Node representation learning on graphs has been extensively studied, such as random-walk based methods [19, 29, 63] and matrix factorization-based methods [10, 12, 62]. Recently, graph neural networks (GNNs) [28, 38, 75] have revolutionized the field of node representation learning [96]. They generally utilize a neighborhood aggregation (or message passing) paradigm to capture the structural information within nodes’ neighborhood. The message passing of the \(t\)th layer in GNNs is usually denoted as:where \(u\) is the neighbor of node \(v\). \(\mathbf {Z}_v^{(t)}\) represents the embedding of node \(v\) at the \(t\)th layer, and \(\mathbf {Z}_v^{(0)}\) is initialized with the input node feature. \(\mathbf {m}_v^{(t)}\) represents the aggregated message from the neighbors of node \(v\). \(\mathrm{COMBINE}^{(t)}(\cdot)\) and \(\mathrm{AGGREGATION}^{(t)}(\cdot)\) are the combination and aggregation functions of GNNs [89]. Many GNNs and their variants [30, 46, 53, 59, 90, 98] have been proposed, achieving state-of-the-art performance on various tasks and demonstrating profound successes in challenging applications, such as recommendation systems [9, 26, 31, 77, 83], information retrieval [17, 91, 95], drug discovery [18, 80], protein function prediction [33, 36], traffic forecasting [21, 37], and so on. However, most existing GNNs do not consider the out-of-distribution generalization ability, so their performances drop substantially on testing data with distribution shifts [33, 44, 80].

A few recent works begin to study the generalization ability of GNNs. The early works [27, 48, 66, 76] focus on the generalization bounds over the training distribution, i.e., in-distribution generalization, which is orthogonal to the OOD generalization and not suitable for the distribution shifts studied in this article. More recently, the OOD generalization ability of GNNs starts to receive research interest [7, 39, 43, 58, 79, 82, 87]. In particular, Bevilacqua et al. [7] learn size-invariant representations for tackling the distribution shifts that exist on graph size. DIR [79] is proposed to discover invariant rationales for GNNs. GIL [45] focuses on capturing the invariant relationships between predictive graph structural information and labels under distribution shifts for OOD generalization. These works mostly concentrate on graph-level tasks and largely ignore the more challenging node-level tasks with multiple latent environments. Some works [24, 54, 99] are proposed to deal with semi-supervised node classification under non-I.I.D. setting. They focus on the adaptation ability of GNNs under distribution shifts, i.e., transferring GNN models trained on the source domain (i.e., environment) to the related target domain with different distributions. For example, SR-GNN [99] is proposed to handle distribution shifts between the selected training and testing nodes by adopting CMD [93] and importance sampling. Reference [24] proposes to learn GNN models by considering agnostic label selection bias. However, these works assume that test data are available and will participate in the training process, which is not in the scope of the OOD generalization problem studied in this article. One exception is the very recent pioneering work EERM [78], which studies invariant node learning by assuming all nodes share a single environment. However, it ignores the more common and challenging situation that nodes are from multiple latent environments. We empirically show that our proposed method greatly outperforms EERM by effectively identifying and modeling multiple latent environments.

The studies on the explainability of GNNs aim to understand the predictions of black-box GNNs by providing the explanations [20, 72, 92]. They generally try to answer which nodes, edges, or features of the input graph are more important for predicting the labels. Several works are proposed to find a subgraph structure and a small subset of node features for the target nodes as the explanations for GNN’s predictions [49, 52, 88]. For example, GNNExplainer [88] learns the soft masks on edges and node features to explain the predictions with the mask optimization. PGExplainer [52] further learns the approximated discrete masks on edges to explain the predictions with a parameterized mask predictor. GraphMask [68] is a post hoc method for explaining the importance of edges in the graph convolution layer. A recent work [79] finds that these explainability works are very sensitive to distribution shifts as most GNN models and proposes discovering invariant explanations in graph-level classification tasks. However, these works focus on understanding the predictions of GNNs instead of learning node representations for better generalization ability under distribution shifts.

Invariant learning has received surging attentions to enable OOD generalization, aiming to generalize to unseen environments by exploiting the invariant relationships between features and labels across distribution shifts. Several works [2, 4, 11, 40, 42, 64] are proposed to learn invariant model and show guaranteed generalization under distribution shifts. However, most existing methods heavily rely on additional environment labels that have to be explicitly provided in the training dataset. Such annotations for the nodes on graph data are usually unavailable and prohibitively expensive to collect, so that these invariant learning methods are inapplicable. A few works study OOD generalization on latent environments in computer vision [16, 51, 56], which cannot be directly applied to graph data. In summary, how to learn invariant node representations under distribution shifts without explicit environment labels remains largely unexplored in the literature.

The Modularity is generally used to measure the divergence between the number of intra-cluster edges and the expected number of a random graph [60], where nodes \(v\) and \(u\) with degrees \(d_v\) and \(d_u\) are connected with probability \({d_v d_u}/{2m}\) and \(m\) is the edge number. By maximizing the modularity, the nodes are densely connected within each cluster [73]:where \(C\) is a cluster assignment matrix, and \(A\) is the adjacency matrix of the input graph for clustering. \(\mathbf {d}\) and \(m\) indicate the degree vector and the number of edges, respectively. However, there are two obstacles for directly adopting this classical modularity maximization method to learn cluster assignment as the inferred environments. The first is that the modularity maximization ignores the inter-cluster edges whose connecting probability should be minimized in the meantime. The second is that we should use the variant patterns \((X^S, A^S)\) of the input graph for clustering rather than use the whole input graph \((X, A)\). Since the invariant patterns capture the invariant relationships between predictive node features and graph structures with the node labels, the variant patterns in turn capture variant spurious correlations under different distributions.