CGRL: Causal-Guided Representation Learning for Node-Level Out-of-Distribution Generalization

Bowen Lu; Lianqiang Yang; Teng Li; Kun Zhang

stat.MLMar 25, 2026

CGRL: Causal-Guided Representation Learning for Node-Level Out-of-Distribution Generalization

Authors: Bowen Lu, Lianqiang Yang, Teng Li, Kun Zhang

Organizations: Anhui University, Hefei, China · Carnegie Mellon University, Pittsburgh, Commonwealth of Pennsylvania, USA

Abstract

Graph Neural Networks (GNNs) deliver strong performance on graph tasks, but their accuracy drops significantly under out-of-distribution (OOD) scenarios. Under distribution shifts, GNNs often fit environmental noise and spurious correlations instead of stable causal mechanisms, leading to weak OOD robustness and unstable predictive representations. Existing solutions based on environment invariance or causal reasoning are insufficient for node classification, as they do not explicitly model the fine-grained latent geometry required by the task. We further observe a training instability named Info-Jitter, where the mutual information between predictive representations and ground-truth labels fluctuates throughout training. To address these issues, we construct a node-classification-specific causal graph derived from the task's geometric objective. Using do-calculus to block non-causal paths caused by environmental noise, we derive a deconfounded interventional objective and a variational lower bound to disentangle representations into intra-class and inter-class components. We then propose Causal-Guided Representation Learning (CGRL), a framework with two core modules. First, a multi-branch re-weighted representation learning (RRL) module learns a causal modulation matrix to amplify causal signals and suppress environmental noise during message passing. Second, an optimization strategy combining intra-class aggregation, inter-class separation, energy-based reconstruction and supervised prediction regularizes the latent space for robust node-level generalization. Experiments on multiple benchmark datasets show that CGRL outperforms strong baselines across various distribution shifts and effectively mitigates the Info-Jitter phenomenon.

Figures & tables

Figure 1. Illustration of the Info-Jitter phenomenon under covariate shifts. We track the mutual information (MI) between predictive representations and ground-truth labels across training epochs on the Cora dataset. Covariate shifts are induced by degree-based partitioning, where high-degree nodes are used for training and low-degree nodes for OOD data. Under identical distributions (left panel of (a) and left panel of (b)), both the conventional GCN and the causal baseline CaNet exhibit smooth and stable MI convergence after an initial rise. In contrast, under covariate shifts (right panel of (a) and right panel of (b)), both models display pronounced MI oscillations throughout training, indicating unstable representation-label alignment and revealing their vulnerability to environment-dependent spurious correlations.

Notation	Meaning	Notation	Meaning
$\mathcal{G}$	Input graph	$G$	Causal graph
$\mathcal{V}$	Node set of graph $\mathcal{G}$	$\phi$	Encoder parameters
$\mathcal{E}$	Edge set of graph $\mathcal{G}$	$v$	Node of $\mathcal{V}$
$E$	Environment of $\mathcal{G}$	$u$	Neighbor nodes of $v$
$\mathbf{Y}$	Ground-truth label	$\mathcal{G}_{v}$	Ego-graph of $v$
$N$	Number of nodes	$\theta$	Classifier parameters

Table 1. Notations and corresponding meanings

Panel of Figure 2 — Figure 2. (a) A causal graph G consisting of 8 variables. (b) $\mathbf{Z}_{a}$ within the ring and $\mathbf{Z}_{r}$ outside the ring respectively represent intra-class and inter-class nodes, where nodes of different colors and types represent different categories.

Figure 3. CGRL consists of two components: causal representation learning and optimization strategy . The former is implemented by a multi-branch re-weighted representation learning (RRL) module, which learns a causal modulation matrix $\Gamma$ to adaptively reweight the predictive representation $\mathbf{Z}_{p}$ . The latter contains four objectives: intra-class aggregation, inter-class separation, energy-based reconstruction, and prediction loss. Given the adjacency matrix $A$ and node feature $\mathbf{X}$ , a GNN encoder first produces the latent representation $\mathbf{Z}$ , which parameterizes the graph prior $P(\mathcal{G}_{v})$ . During training, $\mathbf{Z}$ is further processed by the Gumbel-softmax, the RRL module, and the energy-based reconstruction module to obtain $\mathbf{Z}_{p}$ , which is then fed into a classifier for prediction and supervision via $\mathcal{L}_{sup}$ . In parallel, two geometric surrogate objectives are imposed on the latent space, yielding the intra-class loss $\mathcal{L}_{intra}$ and inter-class loss $\mathcal{L}_{inter}$ . Together with the reconstruction loss $\mathcal{L}_{rec}$ , these components jointly encourage deconfounded representation learning and robust OOD generalization.

Datasets	Nodes	Edges	Classes	Features	Shift types
Cora	2708	5429	7	1433	Feature
Citeseer	3327	4732	6	3703	Feature
Pubmed	19717	44338	3	500	Feature
Arxiv	169343	1166243	40	128	Temporal
Twitch	34120	892346	2	2545	Spatial

Table 2. Datasets with feature and spatiotemporal shifts.

Shifts Type	Datasets	Weight decay	Dropout	Learning rate	Branches	Hidden dimension	Layer of model
Feature	Cora	1e-5	0.4	0.001	3	128	4
	Citeseer	5e-5	0.4	0.001	3	128	3
	Pumbed	5e-5	0.3	0.01	2	128	2
Sptiotemporal	Arxiv	5e-5	0.2	0.01	2	32	2
Sptiotemporal	Twitch	3e-5	0.1	0.01	4	64	5
Covariate	GOODCora	1e-3	0	1e-3	3	256	4

Table 3. Parameter settings of CGRL

Backbone	Method	Cora	Citeseer	Pubmed	Arxiv			Twitch
Backbone	Method	Cora	Citeseer	Pubmed	2014-2015	2016-2017	2018-2020	ES	FR	EN
GCN	ERM	74.30 ±2.66	74.93 ±2.39	81.36 ±1.78	56.33 ±0.17	53.53 ±0.44	45.83 ±0.47	66.07 ±0.14	52.62 ±0.01	63.15 ±0.08
	IRM	74.19 ±2.60	75.34 ±1.61	81.14 ±1.72	55.92 ±0.24	53.25 ±0.49	45.66 ±0.83	66.95 ±0.27	52.53 ±0.02	62.91 ±0.08
	Coral	74.26 ±2.28	74.97 ±2.53	81.56 ±2.35	56.42 ±0.26	53.53 ±0.54	45.92 ±0.52	66.15 ±0.14	52.67 ±0.02	63.18 ±0.03
	DANN	73.09 ±3.25	74.74 ±2.78	80.77 ±1.43	56.35 ±0.11	53.81 ±0.33	45.89 ±0.37	66.15 ±0.13	52.66 ±0.02	63.20 ±0.06
	GroupDRO	74.25 ±2.61	75.02 ±2.05	81.07 ±1.89	56.52 ±0.27	53.40 ±0.29	45.76 ±0.59	66.82 ±0.26	52.69 ±0.02	62.95 ±0.11

Table 4. Test (mean±standard deviation) Accuracy (%) and ROC-AUC (%) on OOD data of Arxiv and Twitch datasets, respectively. OOM denotes out-of-memory. Bold numbers indicate the best result, and the underlined numbers indicate the second-best result.

Backbone	Method	Covariate				Concept
		GOODCora		GOODCBAS	GOODWebKB	GOODCora		GOODCBAS	GOODWebKB
		degree	word	color	university	degree	word	color	university
GCN	ERM	55.78 ±0.52	64.76 ±0.30	78.57 ±2.02	16.14 ±1.35	60.24 ±0.40	64.32 ±0.15	82.14 ±1.17	27.52 ±0.75
	IRM	55.77 ±0.66	64.81 ±0.33	78.57 ±1.17	13.75 ±4.91	61.23 ±0.08	64.42 ±0.18	81.67 ±0.89	27.52 ±0.75
	Coral	56.03 ±0.37	64.75 ±0.26	78.09 ±0.67	11.90 ±1.72	60.41 ±0.27	64.34 ±0.17	82.86 ±0.58	26.61 ±0.75
	DANN	56.10 ±0.59	64.77 ±0.42	77.57 ±2.86	15.08 ±0.37	60.78 ±0.38	64.51 ±0.19	82.50 ±0.72	26.91 ±0.36

Table 5. Test (mean±standard deviation) Accuracy (%) on OOD data of GOODCora, GOODCBAS and GOOODWebKB datasets under the different shifts, respectively. OOM denotes out-of-memory. Bold numbers indicate the best result, and the underlined numbers indicate the second-best result.

Figure 4. Mutual information between predictive representations and ground-truth labels on the GOODCora dataset with degree-domain covariate shifts. Subfigures (a)–(d) correspond to GAT-based ERM, CANet, CIA-LRA, and CGRL, respectively.

Figure 5. Mutual information between predictive representations and ground-truth labels on the GOODCora dataset with degree-domain concept shifts. Subfigures (a)–(d) correspond to GCN-based ERM, CANet, CIA-LRA, and CGRL, respectively.

Figure 6. Mutual information between predictive representations and ground-truth labels on the GOODCora dataset with word-domain concept shifts. Subfigures (a)–(d) correspond to GAT-based ERM, CANet, CIA-LRA, and CGRL, respectively.

Backbone	Method	Cora		Citeseer		Pubmed		Arxiv		Twitch
Backbone	Method	Jitter	TailMI	Jitter	TailMI	Jitter	TailMI	Jitter	TailMI	Jitter	TailMI
GCN	ERM	2.21 $\pm$ 0.08	2.47 $\pm$ 0.02	2.27 $\pm$ 0.21	2.33 $\pm$ 0.03	5.07 $\pm$ 0.29	1.17 $\pm$ 0.01	10.16 $\pm$ 0.86	1.71 $\pm$ 0.04	6.35 $\pm$ 0.53	0.31 $\pm$ 0.02
	IRM	4.17 $\pm$ 1.26	2.32 $\pm$ 0.01	7.22 $\pm$ 1.75	2.11 $\pm$ 0.03	9.44 $\pm$ 0.43	1.13 $\pm$ 0.01	13.39 $\pm$ 3.01	1.85 $\pm$ 0.01	5.73 $\pm$ 0.77	0.09 $\pm$ 0.00
	Coral	1.77 $\pm$ 0.21	1.70 $\pm$ 0.59	1.60 $\pm$ 0.28	1.57 $\pm$ 0.53	2.57 $\pm$ 0.28	0.85 $\pm$ 0.19	4.39 $\pm$ 2.08	0.92 $\pm$ 0.35	0.83 $\pm$ 1.02	0.07 $\pm$ 0.08
	DANN	1.87 $\pm$ 0.27	1.71 $\pm$ 0.62	1.74 $\pm$ 0.24	1.69 $\pm$ 0.52	2.47 $\pm$ 0.15	1.09 $\pm$ 0.02	4.37 $\pm$ 3.35	0.99 $\pm$ 0.44	0.66 $\pm$ 1.15	0.07 $\pm$ 0.04
	GroupDRO	1.85 $\pm$ 0.23	1.69 $\pm$ 0.58	1.77 $\pm$ 0.21	1.51 $\pm$ 0.49	2.87 $\pm$ 0.12	0.81 $\pm$ 0.18	5.53 $\pm$ 1.37	1.08 $\pm$ 0.45	0.70 $\pm$ 0.07	0.06 $\pm$ 0.03

Table 6. Comparison of mutual information fluctuation metrics on datasets under the sptiotemporal shifts and different backbones.

Figure 7. Softmax confidence for Cora and GOODCora datasets.

Figure 8. Mutual information on ID and OOD data.

Methods	Cora	Citeseer	Pubmed	Arxiv			Twitch
				2014-2015	2016-2017	2018-2020	ES	FR	EN
w/o $\mathcal{L}_{sup}$	70.34 ±1.17	69.89 ±2.31	78.52 ±2.17	46.23 ±0.27	45.14 ±0.28	42.17 ±0.36	65.47 ±0.08	49.48 ±0.12	58.78 ±0.04
w/o $\mathcal{L}_{rec}$	91.78 ±0.32	87.12 ±0.19	85.73 ±0.41	49.82 ±0.18	48.43 ±0.21	46.77 ±0.67	66.56 ±0.27	52.67 ±0.19	62.76 ±0.08
w/o $\mathcal{L}_{intra}$ / $\mathcal{L}^{sc}_{intra}$	96.64 ±0.23	94.14 ±0.21	90.81 ±0.59	58.63 ±0.44	58.64 ±0.63	56.12 ±1.02	66.89 ±0.18	52.72 ±0.34	62.35 ±0.13
w/o $\mathcal{L}_{inter}$ / $\mathcal{L}^{sc}_{inter}$	97.57 ±0.21	95.40 ±0.18	92.47 ±0.31	58.83 ±0.24	58.22 ±0.28	57.37 ±0.98	67.05 ±0.32	53.01 ±0.29	63.10 ±0.15
CGRL	99.14 ±0.16	97.96 ±0.23	96.90 ±0.44	61.57 ±0.32	60.12 ±0.56	59.45 ±0.82	68.02 ±0.29	54.98 ±0.23	64.89 ±0.12

Table 7. Ablation results under feature and spatiotemporal shifts across datasets.

Figure 9. Ablation results under covariate and concept shifts across datasets.

Figure 10. The sensitivity of branches K to all datasets.

Figure 11. The sensitivity of

\lambda_{1}

and

\lambda_{2}

to all datasets.

Explore similar work

Feb 20, 2026

Causal Neighbourhood Learning for Invariant Graph Representations

Graph data often contain noisy and spurious correlations that mask the true causal relationships, which are essential for enabling graph models to make predictions based on the underlying causal structure of the data. Dependence on spurious connections makes it challenging for traditional Graph Neural Networks (GNNs) to generalize effectively across different graphs. Furthermore, traditional aggregation methods tend to amplify these spurious patterns, limiting model robustness under distribution shifts. To address these issues, we propose Causal Neighbourhood Learning with Graph Neural Networks (CNL-GNN), a novel framework that performs causal interventions on graph structure. CNL-GNN effectively identifies and preserves causally relevant connections and reduces spurious influences through the generation of counterfactual neighbourhoods and adaptive edge perturbation guided by learnable importance masking and an attention-based mechanism. In addition, by combining structural-level interventions with the disentanglement of causal features from confounding factors, the model learns invariant node representations that are robust and generalize well across different graph structures. Our approach improves causal graph learning beyond traditional feature-based methods, resulting in a robust classification model. Extensive experiments on four publicly available datasets, including multiple domain variants of one dataset, demonstrate that CNL-GNN outperforms state-of-the-art GNN models.

Simi Job, Xiaohui Tao, Taotao Cai +2

School of Science, Engineering and Digital Technologies, University of Southern Queensland, Australia. · School of Data Science, Lingnan University, Hong Kong.

Graph Representation Learning Causal Representation Learning OOD Generalization Graph Structure Learning Causal Inference

Ask PDF

Oct 6, 2026

Complementary Supervised and Self-Supervised Representations for Out-of-Distribution Graph Learning

Out-of-distribution (OOD) generalization remains challenging for graph neural networks (GNNs), as graph distributions can vary substantially across time and domains. Supervised and self-supervised graph representation learning are guided by distinct objectives and offer different perspectives on graph representations. In this work, we study whether self-supervised representations (SSL) can provide complementary signals to improve supervised OOD node classification. We develop two backbone-agnostic frameworks that exploit such information at different stages of learning and prediction. Co-Train jointly learns supervised and SSL representations and adaptively integrates them during training, while Dual-Space Retrieval performs non-parametric prediction in the two representation spaces and combines their predictions through confidence-aware fusion at inference time. The supervised and SSL encoders are separately parameterized and need not share the same GNN architecture. We evaluate multiple GNN backbones and two distinct SSL objectives, DGI and GRACE, on four graph benchmarks spanning temporal and cross-domain distribution shifts. Extensive experiments show that Co-Train consistently outperforms strong supervised OOD baselines, while Dual-Space Retrieval achieves competitive performance as a flexible non-parametric alternative. Results across different backbones and SSL objectives, together with representation analyses and ablations, demonstrate that SSL representations provide complementary information to supervised representations and can improve OOD node classification across diverse settings.

Qingying Hao, Zikang Chen, Chuxuan Hu +4

ShanghaiTech University Shanghai, China · University of Illinois Urbana-Champaign Urbana, Illinois, USA · The Pennsylvania State University University Park, Pennsylvania, USA

Graph Neural Networks Graph Representation Learning Graph SSL OOD Generalization

Ask PDF

Mar 2, 2026

Towards OOD Generalization in Dynamic Graphs via Causal Invariant Learning

Although dynamic graph neural networks (DyGNNs) have demonstrated promising capabilities, most existing methods ignore out-of-distribution (OOD) shifts that commonly exist in dynamic graphs. Dynamic graph OOD generalization is non-trivial due to the following challenges: 1) Identifying invariant and variant patterns amid complex graph evolution, 2) Capturing the intrinsic evolution rationale from these patterns, and 3) Ensuring model generalization across diverse OOD shifts despite limited data distribution observations. Although several attempts have been made to tackle these challenges, none has successfully addressed all three simultaneously, and they face various limitations in complex OOD scenarios. To solve these issues, we propose a Dynamic graph Causal Invariant Learning (DyCIL) model for OOD generalization via exploiting invariant spatio-temporal patterns from a causal view. Specifically, we first develop a dynamic causal subgraph generator to identify causal dynamic subgraphs explicitly. Next, we design a causal-aware spatio-temporal attention module to extract the intrinsic evolution rationale behind invariant patterns. Finally, we further introduce an adaptive environment generator to capture the underlying dynamics of distributional shifts. Extensive experiments on both real-world and synthetic dynamic graph datasets demonstrate the superiority of our model over state-of-the-art baselines in handling OOD shifts.

Xinxun Zhang, Pengfei Jiao, Mengzhou Gao +2

School of Cyberspace, Hangzhou Dianzi University, China · State Key Laboratory of AI Safety, Beijing · School of Computer Science, Wuhan University, China +1

OOD Generalization Dynamic Graph Learning

Ask PDF