The recent works on causal discovery have followed a similar trend of learning partial ancestral graphs (PAGs) since observational data constrain the true causal directed acyclic graph (DAG) only up to a Markov equivalence class. This limits their application in the majority of downstream tasks, as uncertainty in causal relations remains unresolved. We propose a new refinement framework, CausalSAGE, for converting PAGs to DAGs while respecting the underlying causal relations. The framework expands discrete variables into state-level representations, constrains the search space using structural knowledge and soft priors, and applies a unified differentiable objective for joint optimization. The final DAG is obtained by aggregating the optimized structures and enforcing acyclicity when necessary. Our experimental evaluations show that the obtained DAGs preserve the underlying causal relations while also being efficient to obtain.
Figures & tables
Figure 1: Why PAGs are not enough. (a) Most downstream tasks require a fully specified causal DAG. (b) FCI handles latent confounders but outputs partially oriented graphs (PAGs). (c) Ambiguous parent sets (e.g., Pa(X) ) prevent computing interventional effects.
Figure 2: CausalSAGE . The framework expands discrete variables into state representations, constrains the search space using random/LLM priors, and finally applies a unified differentiable objective for joint optimization to output a DAG.
Figure 3: SHD reduction via refinement (Q1). Bars show Δ SHD = (Baseline − Refined), where the positive values indicate improvement.
Figure 4: Elimination of Orientation Ambiguity. FCI/RFCI leaves a large fraction of edges unresolved (46%–86%) across datasets. CausalSAGE reduces this ratio to 0% in all cases, effectively breaking Markov equivalence and producing fully directed DAGs.
Figure 5: Comparison with standard DAG learners (10k samples). SHD measures structural error, while F1 measures directional accuracy. Points closer to the top-left corner indicate stronger overall performance.
Figure 6: Effect of sample size on SHD. As the number of observational samples increases, structural error consistently decreases across datasets.
Dataset
Variables
Sample Size
Runtime (s)
Sachs
11
855
2.8
Child
20
5000
2.6
Insurance
27
10000
38
Alarm
37
10000
32
Hailfinder
56
10000
59
Win95pts
76
10000
71
Table 1: CausalSAGE is efficient. The runtime as a function of graph size shows the efficiency of our method.
Appendix figures & tables10 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
Nodes
Edges
Samples
Type
Sachs
11
17
855
Continuous
Child
20
25
5000
Discrete
Insurance
27
52
10000
Discrete
Alarm
37
46
10000
Discrete
Hailfinder
56
66
10000
Discrete
Win95pts
76
112
10000
Discrete
Appendix
Table 2: Benchmark datasets from the bnlearn repository.
Massachusetts Institute of Technology · This work was conducted while the author was a PhD student at Purdue University, independently of the author’s dissertation research.