A Knowledge-Informed Pretrained Model for Causal Discovery

Wenbo Xu; Yue He; Yunhai Wang; Xingxuan Zhang; Kun Kuang; Yueguo Chen; Peng Cui

cs.LGMar 21, 2026

A Knowledge-Informed Pretrained Model for Causal Discovery

Authors: Wenbo Xu, Yue He, Yunhai Wang, Xingxuan Zhang, Kun Kuang, Yueguo Chen, Peng Cui

Organizations: Renmin University of China, China · Tsinghua University, China · Zhejiang University, China

Abstract

Causal discovery has been widely studied, yet many existing methods rely on strong assumptions or fall into two extremes: either depending on costly interventional signals or partial ground truth as strong priors, or adopting purely data driven paradigms with limited guidance, which hinders practical deployment. Motivated by real-world scenarios where only coarse domain knowledge is available, we propose a knowledge-informed pretrained model for causal discovery that integrates weak prior knowledge as a principled middle ground. Our model adopts a dual source encoder-decoder architecture to process observational data in a knowledge-informed way. We design a diverse pretraining dataset and a curriculum learning strategy that smoothly adapts the model to varying prior strengths across mechanisms, graph densities, and variable scales. Extensive experiments on in-distribution, out-of distribution, and real-world datasets demonstrate consistent improvements over existing baselines, with strong robustness and practical applicability.

Figures & tables

Figure 1: Simple causality in chemistry. (A) Expert knowledge is structured as graph. (B) The corresponding ground-truth causal graph.

Figure 2: An illustration of knowledge encoding. (A) A symbolic form of Fig. 1 . (B) The corresponding adjacency matrix

\mathbf{A}

. (C) The reachability matrix

\mathbf{R}

, which further encodes indirect relations (e.g.,

R_{A,C}=1

). (D) The resulting prior knowledge

G_{P}

, where infeasible relations are excluded and self-loops are removed.

Figure 3: (A) The overall pipeline of Kode, illustrating the global framework that integrates raw data and prior knowledge to infer the causal graph; (B) The dual-source encoder, which achieves alignment and feature compression from hybrid inputs; (C) The graph decoder, responsible for mapping the latent representations back into the final causal structure.

		Linear		Neural Networks
Variables	Edges	$\text{noise}\sim\mathcal{N}(0,\,d^{2})$		$\text{Latent}\sim\mathcal{U}(-1,\,1)$	$\text{Latent}\sim\Gamma(1,\,1)$	$\text{Latent}\sim{N}(0,\,d^{2})$
		$d\sim\Gamma(2.5,\,2.5)$	$d\sim\Gamma(1,\,1)$	$\text{Latent}\sim\mathcal{U}(-1,\,1)$	$\text{Latent}\sim\Gamma(1,\,1)$	$d\sim\Gamma(1,\,1)$
20	$\|\mathbf{E}\|\in[20,\,60]$	20,000	30,000	16,667	16,667	16,667
30	$\|\mathbf{E}\|\in[30,\,90]$	20,000	30,000	16,667	16,667	16,667
40	$\|\mathbf{E}\|\in[40,\,120]$	20,000	30,000	16,667	16,667	16,667

Table 1: Training dataset configurations under different generative mechanisms.

Prior Type		Probability
Reachable Matrix	(0, 1, -1)	0.45
Reachable Matrix	(0, -1)	0.45
Ground Truth	(0, 1, -1)	0.05
Ground Truth	(0, -1)	0.05

Table 2: Sampling probabilities for different prior types.

Epochs	Proportion	Sparse Threshold
[0, 5)	No Sparse	/
[5, 8)	8 : 2	(0.34, 0.40]
[8, 10)	6 : 4	(0.30, 0.34]
[10, 15)	4 : 6	(0.20, 0.30]
[15, 50]	2 : 8	0.20

Table 3: Illustration how the sampling proportions and the sparsity threshold evolve over training epochs.

Panel of Figure 4 — Figure 4: Results on 20 nodes, covering i.d. linear cases and both i.d. & o.o.d nonlinear mechanisms.

Figure 5: Performance on Gaussian Process ER(3) across scales: Kode maintains a clear advantage; GES (

4^{th}

bar) collapses owing to excessive density at

N=30

and

40

, so its results are omitted.

Panel of Figure 6 — Figure 6: Sachs Experiment

	Food		TEP
	SHD	F1	SHD	F1
Kode (with topo-priors)	80.1	0.40	31.26	0.44
Kode (without priors)	88.35	0 .35	3 5.47	0 .38
LiNGAM	89.66	0.08	48.82	0.14
PC	95.23	0.19	68.57	0.14
GES	97.24	0.32	124.35	0.09

Table 4: Industry Experiments

Causal Discovery

Causal discovery infers cause-and-effect structures from observational and interventional data, distinguishing causal relationships from correlation.

See all 221 papers

TabCausal: Pretraining Across Causal Environments for Tabular Causal Discovery
May 29, 2026Zi-Rong Li, Si-Yang Liu, Tian-Zuo Wang +1Causal Foundation Models Synthetic Data Pretraining
Nanjing University
CDFM: Towards a General-Purpose Causal Discovery Foundation Model
Jul 13, 2026Jie Qiao, Ruichu Cai, Zijian Li +6Causal Foundation Models Causal Discovery
School of Computer Science, Guangdong University of Technology, Guangzhou, China · College of Mathematics and Computer, Shantou University, Shantou, China · Department of Computer Science and Technology, Tsinghua University, Beijing, China
Use What You Know: Causal Foundation Models with Partial Graphs
Feb 16, 2026Arik Reuter, Anish Dhir, Cristiana Diaconu +6Causal Foundation Models Causal Discovery
University of Cambridge, Cambridge, United Kingdom · Max Planck Institute for Intelligent Systems, Tübingen, Germany · Gatsby Computational Neuroscience Unit, University College London +5

May 29, 2026J/K move · Enter open · S save

TabCausal: Pretraining Across Causal Environments for Tabular Causal Discovery

Zi-Rong Li, Si-Yang Liu, Tian-Zuo Wang, Han-Jia Ye

Nanjing University

Causal Foundation Models Synthetic Data Pretraining Structural Causal Models Causal Discovery

Open paper AskPDF arXiv

Causal discovery aims to recover directed causal relations from observational and interventional data, providing a basis for mechanistic understanding and reliable decision-making. Causal discovery foundation models (CDFMs) seek to amortize this problem by mapping a dataset directly to a causal graph in a single forward pass, avoiding per-dataset testing, search, or optimization. However, existing CDFMs remain limited, often failing to consistently match strong classical methods, and we find that a key bottleneck is how causal pretraining tasks are constructed. Based on this observation, we propose TabCausal, a data-driven CDFM trained with broad causal pretraining over diverse graph priors, structural mechanisms, noise models, dimensions, sample sizes, and intervention regimes. A dynamic task construction strategy composes these causal environments into varied discovery tasks, enabling more transferable structural learning from observational and mixed-interventional data. On large-scale synthetic benchmarks, TabCausal achieves better macro-averaged performance than a diverse set of causal discovery baselines. To further bridge abstract synthetic generators and realistic causal reasoning scenarios, we introduce a protocol-guided and LLM-audited semantic causal environment benchmark, where domain-grounded SCMs generate interpretable observational and interventional datasets for out-of-distribution analysis. Across both synthetic and semantic environments, TabCausal demonstrates robust structure recovery, especially under interventional evidence, highlighting broad causal pretraining as a key ingredient for transferable amortized causal discovery.

Abstract

Causal Discovery

Explore similar work

TabCausal: Pretraining Across Causal Environments for Tabular Causal Discovery

CDFM: Towards a General-Purpose Causal Discovery Foundation Model

Use What You Know: Causal Foundation Models with Partial Graphs