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

Causal Neighbourhood Learning for Invariant Graph Representations
Feb 20, 2026Simi Job, Xiaohui Tao, Taotao Cai +2Graph Representation Learning Causal Representation Learning
School of Science, Engineering and Digital Technologies, University of Southern Queensland, Australia. · School of Data Science, Lingnan University, Hong Kong.
Complementary Supervised and Self-Supervised Representations for Out-of-Distribution Graph Learning
Oct 6, 2026Qingying Hao, Zikang Chen, Chuxuan Hu +4Graph Neural Networks Graph Representation Learning
ShanghaiTech University Shanghai, China · University of Illinois Urbana-Champaign Urbana, Illinois, USA · The Pennsylvania State University University Park, Pennsylvania, USA
Towards OOD Generalization in Dynamic Graphs via Causal Invariant Learning
Mar 2, 2026Xinxun Zhang, Pengfei Jiao, Mengzhou Gao +2OOD Generalization Dynamic Graph Learning
School of Cyberspace, Hangzhou Dianzi University, China · State Key Laboratory of AI Safety, Beijing · School of Computer Science, Wuhan University, China +1

Feb 20, 2026J/K move · Enter open · S save

Causal Neighbourhood Learning for Invariant Graph Representations

Simi Job, Xiaohui Tao, Taotao Cai, Haoran Xie+1

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Causal Neighbourhood Learning for Invariant Graph Representations

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

Towards OOD Generalization in Dynamic Graphs via Causal Invariant Learning