A Causal Foundation Model for Structure and Outcome Prediction

# A Causal Foundation Model for Structure and Outcome Prediction
Source: [https://arxiv.org/html/2606.26467](https://arxiv.org/html/2606.26467)
###### Abstract

We introduce TabPFN\-CFM, a causal foundation model that can handle multiple causal problems\. TabPFN\-CFM predicts both causal structure and outcomes from observational data, supports queries on all three levels of Pearl’s Causal Hierarchy and uses known graph structure when available to improve predictions\. TabPFN\-CFM is trained on synthetic datasets, and generalises to real datasets, demonstrating improved performance over both structural and outcome prediction baselines\.

Machine Learning, ICML

## 1Introduction

Causal relationships between processes determine how intervening on a system impacts its behaviour\. However, determining the causal structure and causal outcomes from observational data is a challenging problem\(Pearl,[2000](https://arxiv.org/html/2606.26467#bib.bib11); Imbens and Rubin,[2015](https://arxiv.org/html/2606.26467#bib.bib24)\)\. Recently, machine learning models have been developed to make causal predictions on the outcome of interventions efficiently with minimal data\(Künzelet al\.,[2019](https://arxiv.org/html/2606.26467#bib.bib18); Hollmannet al\.,[2023](https://arxiv.org/html/2606.26467#bib.bib15); Robertsonet al\.,[2025](https://arxiv.org/html/2606.26467#bib.bib16); Balazadehet al\.,[2025](https://arxiv.org/html/2606.26467#bib.bib17)\), as well as structural predictions\(Lorchet al\.,[2022](https://arxiv.org/html/2606.26467#bib.bib14); Keet al\.,[2023](https://arxiv.org/html/2606.26467#bib.bib25)\)\. We train a causal foundation model that predicts both causal structure and outcomes from observational data\. Compared to existing methods, TabPFN\-Causal Foundation Model \(TabPFN\-CFM\) improves accuracy, supports all three levels of Pearl’s Causal Hierarchy\(Pearl,[2000](https://arxiv.org/html/2606.26467#bib.bib11)\), and leverages known graph structure when available\. This is further accompanied by a refined training procedure that improves training efficiency almost 4 times\.

### 1\.1Problem Setup

We assume the underlying system follows an SCMψ=\{U,V,F\}\\psi=\\\{U,V,F\\\}, consisting of unobserved variablesUU, observed variablesVV, and unknown structural equationsFF\. Unknown variables include unobserved covariates and sources of noise\. The observed variablesVVare split into covariatesXX, a binary treatmentTTand an outcomeYY\. The structural equationsFFdefine each variable’s parents and form a graph𝒢\\mathcal\{G\}\. We have access to an observational dataset𝒟obs=\{𝐱i,yi,ti\}i=1n\\mathcal\{D\}^\{\\mathrm\{obs\}\}=\\\{\\mathbf\{x\}\_\{i\},y\_\{i\},t\_\{i\}\\\}\_\{i=1\}^\{n\}, consisting ofnni\.i\.d\. samples drawn fromψ\\psi, and optionally, prior knowledge of the causal graph𝒢est\\mathcal\{G\}^\{\\mathrm\{est\}\}\. If there is no prior then𝒢est=∅\\mathcal\{G\}^\{\\mathrm\{est\}\}=\\varnothing\.

Our model targets two objectives\. Firstly, if the true causal graph is unknown, our model can make an estimate of the underlying causal graph structure\. Second is outcome predictiton under three types of causal queries\. Given a samples covariates𝐱∗\\mathbf\{x\}^\{\*\}, natural treatmentt∗t^\{\*\}and outcomey∗y^\{\*\}, the Observational Query,P​\(y∗\|𝐱∗,T=t∗,𝒟obs,𝒢est\)P\(y^\{\*\}\|\\mathbf\{x\}^\{\*\},T=t^\{\*\},\\mathcal\{D\}^\{\\mathrm\{obs\}\},\\mathcal\{G\}^\{\\mathrm\{est\}\}\), predicts the outcome when the treatment takes its observed value\. The Interventional Query,P​\(y∗\|𝐱∗,do​\(T=1−t∗\),𝒟obs,𝒢est\)P\(y^\{\*\}\|\\mathbf\{x\}^\{\*\},\\mathrm\{do\}\(T=1\-t^\{\*\}\),\\mathcal\{D\}^\{\\mathrm\{obs\}\},\\mathcal\{G\}^\{\\mathrm\{est\}\}\), predicts the outcome under an externally assigned treatment which differs from the observed treatment\. The Counterfactual Query,P​\(y∗\|𝐱∗,do​\(T=1−t∗\),y=yt∗,𝒟obs,𝒢est\)P\(y^\{\*\}\|\\mathbf\{x\}^\{\*\},\\mathrm\{do\}\(T=1\-t^\{\*\}\),y=y\_\{t^\{\*\}\},\\mathcal\{D\}^\{\\mathrm\{obs\}\},\\mathcal\{G\}^\{\\mathrm\{est\}\}\), predicts the what\-if outcome under a different treatment, conditional on the outcome observed under the original treatment\. By training a single model jointly on all tasks, the model gains a broader understanding of causal processes and learns more efficiently than models trained on a single task\.

oTo approach this problem, we follow the Bayesian PFN framework\(Robertsonet al\.,[2025](https://arxiv.org/html/2606.26467#bib.bib16); Balazadehet al\.,[2025](https://arxiv.org/html/2606.26467#bib.bib17)\)\. A prior over SCMs,p​\(ψ\)p\(\\psi\), generates a true SCM,ψtrue∼p​\(ψ\)\\psi\_\{\\mathrm\{true\}\}\\sim p\(\\psi\)\. Observational data yields the posterior,P​\(ψ\|𝒟obs,𝒢est\)P\(\\psi\|\\mathcal\{D\}^\{\\mathrm\{obs\}\},\\mathcal\{G\}^\{\\mathrm\{est\}\}\)\. Given an SCM, the outcome can be estimated asP​\(y\|𝐱∗,do​\(T=t∗\),ψ\)P\(y\|\\mathbf\{x\}^\{\*\},\\mathrm\{do\}\(T=t^\{\*\}\),\\psi\)\. The target distribution for a causal query, say an Interventional Query, is

P​\(y∗\|𝐱∗,do​\(T=1−t∗\),𝒟obs,𝒢est\)=\\displaystyle P\(y^\{\*\}\|\\mathbf\{x\}^\{\*\},\\mathrm\{do\}\(T=1\-t^\{\*\}\),\\mathcal\{D\}^\{\\mathrm\{obs\}\},\\mathcal\{G\}^\{\\mathrm\{est\}\}\)=\(1\)∫P​\(y∗\|𝐱∗,do​\(T=1−t∗\),ψ\)​P​\(ψ\|𝒟obs,𝒢est\)​𝑑ψ\.\\displaystyle\\int P\(y^\{\*\}\|\\mathbf\{x\}^\{\*\},\\mathrm\{do\}\(T=1\-t^\{\*\}\),\\psi\)P\(\\psi\|\\mathcal\{D\}^\{\\mathrm\{obs\}\},\\mathcal\{G\}^\{\\mathrm\{est\}\}\)d\\psi\.\(2\)If the graph structure𝒢\\mathcal\{G\}is unknown, its posterior is

P​\(𝒢∗\|𝒟obs\)=∫P​\(𝒢∗\|ψ\)​P​\(ψ\|𝒟obs\)​𝑑ψ\.P\(\\mathcal\{G\}^\{\*\}\|\\mathcal\{D\}^\{\\mathrm\{obs\}\}\)=\\int P\(\\mathcal\{G\}^\{\*\}\|\\psi\)P\(\\psi\|\\mathcal\{D\}^\{\\mathrm\{obs\}\}\)d\\psi\.\(3\)
The optional graph allows the posterior to be estimated more accurately, by providing information about the causal relationships between variables\. In Appendix[A](https://arxiv.org/html/2606.26467#A1), we prove that including the graph as input can only improve posterior accuracy, with no improvement occurring if the graph does not change the posterior distribution ofψ\\psi\(whereyyhas support\)\.

In practice, instead of explicitly estimatingP​\(ψ\|𝒟obs\)P\(\\psi\|\\mathcal\{D\}^\{\\mathrm\{obs\}\}\), we train a model to directly estimateP​\(y∗\|𝐱∗,do​\(T=t∗\),𝒟obs\)P\(y^\{\*\}\|\\mathbf\{x\}^\{\*\},\\mathrm\{do\}\(T=t^\{\*\}\),\\mathcal\{D\}^\{\\mathrm\{obs\}\}\)andP​\(𝒢\|𝒟obs\)P\(\\mathcal\{G\}\|\\mathcal\{D\}^\{\\mathrm\{obs\}\}\), by drawing samples from the priorp​\(ψ\)p\(\\psi\)and optimising log likelihoods conditioned on observations,L=−E​\[log⁡p^θ​\(y∗\|𝐱∗,do​\(T=1−t∗\),𝒟obs,𝒢\)\]L=\-E\[\\log\\hat\{p\}\_\{\\theta\}\(y^\{\*\}\|\\mathbf\{x\}^\{\*\},\\mathrm\{do\}\(T=1\-t^\{\*\}\),\\mathcal\{D\}^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\]\. We show this loss is equivalent to optimising the KL divergence between the model’s estimated distribution and the true posterior distribution in Appendix[B](https://arxiv.org/html/2606.26467#A2)\.

Existing deep learning methods for causal structure learning do not predict confounding variables caused by unobserved factors\. Our method addresses this by using Acyclic Directed Mixed Graphs \(ADMGs\), which represent unobserved confounding via bidirected edges; details are provided in Appendix[C](https://arxiv.org/html/2606.26467#A3)\. Confounding may lead to unidentifiable graphs, soP​\(ψ\|𝒟obs,𝒢est\)P\(\\psi\|\\mathcal\{D\}^\{\\mathrm\{obs\}\},\\mathcal\{G\}^\{\\mathrm\{est\}\}\)is uncertain, which is reflected in the estimated distributions\. By predicting confounding, our model not only sheds light on the underlying causal structure, but also allows the user to account for confounding when interpreting results\.

## 2Data generation

Our model is trained on synthetic data generated from a prior distribution of SCMs to approximate Bayesian inference on this prior\. First, we sample SCMs with random graph structures, missing nodes, noise distributions, and structural equations\. Random functions are generated using neural networks with random weights and nonlinearities, along with random noise distributions\.

To generate a single training sample, for each SCM, we sample observational data𝒟obs\\mathcal\{D\}^\{\\mathrm\{obs\}\}from the SCM and generate another observational datapoint for prediction,𝒟pred=\{𝐱∗,t∗,y∗\}\\mathcal\{D\}^\{\\mathrm\{pred\}\}=\\\{\\mathbf\{x\}^\{\*\},t^\{\*\},y^\{\*\}\\\}\. Since the true causal graph is known, we generate a counterfactual \(with fixed noise\),𝒟causal=\{𝐱t∗,t=do\(1−t∗\),yt∗\}\\mathcal\{D\}^\{\\mathrm\{causal\}\}=\\\{\\mathbf\{x\}^\{\*\}\_\{t\},t=\\mathrm\{do\}\(1\-t^\{\*\}\),y^\{\*\}\_\{t\}\\\}\. Finally, we have the ADMG corresponding to the SCM structure,𝒢est=\{A,C\}\\mathcal\{G\}^\{\\textrm\{est\}\}=\\\{A,C\\\}, with adjacency matrixAAand bidirected confounding matrixCC\.

Our setup extends the priors used inRobertsonet al\.\([2025](https://arxiv.org/html/2606.26467#bib.bib16)\)\. This diversity in the prior ensures our model learns to perform inference on a wide variety of causal systems, improving its generalization to real data\. A full description of the data generation process is given in Appendix[D](https://arxiv.org/html/2606.26467#A4)\.

## 3Model architecture

### 3\.1Architecture overview

FitGGPredRow AttnCol AttnMHAMLPRow AttnCol AttnMLPGG\-EmbKVKVKVRow MHACol MHAGraphMLP×n\\times nblocksMatrix Decodery^\\hat\{y\}DecodercontextgraphtargetFigure 1:Model architecture diagram\.Our model extends the Do\-PFN architecture to support interventional and counterfactual prediction with explicit graph conditioning\. Details are given in Appendix[E](https://arxiv.org/html/2606.26467#A5)\.

The encoder embeds the fit set𝒟fit\\mathcal\{D\}^\{\\mathrm\{fit\}\}, the prediction query𝒟pred\\mathcal\{D\}^\{\\mathrm\{pred\}\}, and the prior ADMG structure\. Covariates, outcomes and treatments are embedded with randomized column positional embeddings as in TabPFNv2, along with the query type\. The ADMG is encoded through the adjacencyAAand bidirected matrixCCas well as the derived ancestral matrix, giving the model direct access to parent, child, ancestor, and descendant relationship\(Keet al\.,[2023](https://arxiv.org/html/2606.26467#bib.bib25)\)\.

The main transformer consists of row\-wise attention, column\-wise attention and feed\-forward layers\. The predictions additionally attend to the graph embeddings, enabling model predictions to incorporate the graph prior\.

The outputs are computed from the final hidden state\. The final hidden state corresponding to the outcome variable is passed through an MLP head to produce logits over discretized outcome buckets, yielding the predictive distributionp^θ​\(y\)\\hat\{p\}\_\{\\theta\}\(y\)\. The graph structure predictions are generated using the final hidden states corresponding to each feature variable\. A decoder computes elementwise predictions for every possible edge in the graph\. Predictions are made for the directed adjacency matrixA^\\hat\{A\}, the bidirected correlation matrixC^\\hat\{C\}and the ancestral matrixR^\\hat\{R\}\.

In our architecture, attention flows from𝒟fit\\mathcal\{D\}^\{\\mathrm\{fit\}\}and the graph to𝒟pred\\mathcal\{D\}^\{\\mathrm\{pred\}\}, but not the other way around\. In addition to allowing efficient inference from the KV cache, as in TabPFN, this makes the predicted graph structure independent of the prediction query, including the input graph prior\.

Compared with TabPFNv2 and Do\-PFN, we omit feature grouping, ensembling, random augmentations, and random feature products because the graph explicitly indexes the variables and would be difficult to preserve under feature mixing\. Furthermore, we also introduce several training and backbone improvements for stability and efficiency which together improve training speed 3\-4x, with significantly lower final loss\. These changes are evaluated in the ablation study in Appendix[F](https://arxiv.org/html/2606.26467#A6)\.

### 3\.2Training procedure

Each training datapoint consists of datasets𝒟fit\\mathcal\{D\}^\{\\mathrm\{fit\}\},𝒟pred=\{𝐱∗,t∗,y∗\}\\mathcal\{D\}^\{\\mathrm\{pred\}\}=\\\{\\mathbf\{x\}^\{\*\},t^\{\*\},y^\{\*\}\\\},𝒟causal=\{𝐱t∗,t=do\(1−t∗\),yt∗\}\\mathcal\{D\}^\{\\mathrm\{causal\}\}=\\\{\\mathbf\{x\}^\{\*\}\_\{t\},t=do\(1\-t^\{\*\}\),y^\{\*\}\_\{t\}\\\}, and𝒢est\\mathcal\{G\}^\{\\mathrm\{est\}\}\. Each datapoint is used to generate samples for all three query types\. In the observational query the model is given𝐱∗\\mathbf\{x\}^\{\*\}andt∗t^\{\*\}, and the target isy∗y^\{\*\}\. In the interventional query the model is given𝐱∗\\mathbf\{x\}^\{\*\}andd​o​\(1−t∗\)do\(1\-t^\{\*\}\), and the target isyt∗y^\{\*\}\_\{t\}\. In the counterfactual query, the model is given𝐱t∗\\mathbf\{x\}\_\{t\}^\{\*\},d​o​\(1−t∗\)do\(1\-t^\{\*\}\)andy∗y^\{\*\}, and the target isyt∗y^\{\*\}\_\{t\}\. In each case, letytargy\_\{\\mathrm\{targ\}\}andp^θ​\(y\)\\hat\{p\}\_\{\\theta\}\(y\)be the target and predicted distribution, respectively\. The prediction loss is

ℒpred=C​E​\(ytarg,p^θ​\(y\)\)\.\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{pred\}\}=CE\(y\_\{\\mathrm\{targ\}\},\\hat\{p\}\_\{\\theta\}\(y\)\)\.\(4\)This is the standard prediction loss for PFNs\(Hollmannet al\.,[2023](https://arxiv.org/html/2606.26467#bib.bib15); Balazadehet al\.,[2025](https://arxiv.org/html/2606.26467#bib.bib17); Robertsonet al\.,[2025](https://arxiv.org/html/2606.26467#bib.bib16)\)\. The model is always given𝒟fit\\mathcal\{D\}^\{\\mathrm\{fit\}\}\. In order to allow the model to learn to make predictions with and without the true causal graph,𝒢est\\mathcal\{G\}^\{\\mathrm\{est\}\}is set to zero half the time\.

The structural prediction losses between predictionsA^,R^,C^\\hat\{A\},\\hat\{R\},\\hat\{C\}and true matrices are the elementwise binary cross entropy,

ℒgraph=1\(k\+2\)2​∑i=1k\+2∑j=1k\+2B​C​E​\(Mi,j,M^i,j\),\\displaystyle\\mathcal\{L\}\_\{\\mathrm\{graph\}\}=\\frac\{1\}\{\(k\+2\)^\{2\}\}\\sum\_\{i=1\}^\{k\+2\}\\sum\_\{j=1\}^\{k\+2\}BCE\(M\_\{i,j\},\\hat\{M\}\_\{i,j\}\),\(5\)M∈\{A,R,C\},M^∈\{A^,R^,C^\}\.\\displaystyle M\\in\\\{A,R,C\\\},\\hat\{M\}\\in\\\{\\hat\{A\},\\hat\{R\},\\hat\{C\}\\\}\.\(6\)Since the structural predictions are independent of the query, this loss is computed only once per datapoint\. The overall loss is

ℒ=ℒpred\+λ⋅ℒgraph\.\\displaystyle\\mathcal\{L\}=\\mathcal\{L\}\_\{\\mathrm\{pred\}\}\+\\lambda\\cdot\\mathcal\{L\}\_\{\\mathrm\{graph\}\}\.\(7\)Finally, the model is trained with dummy features from batch padding\. Although samples are processed independently, the model can attend to these dummy features; empirically, adding them at inference improves performance, suggesting they may support intermediate computations\.

## 4Evaluations

### 4\.1Synthetic toy examples

First, we showcase TabPFN\-CFM on the Instrumental Variable \(IV\) problem with SEM:Z→T→Y,U→T,U→YZ\\rightarrow T\\rightarrow Y,U\\rightarrow T,U\\rightarrow Y\. A linear IV SEM is generated with known parameters and noise distributions, which allows the exact observational, interventional and counterfactual distributions to be computed\. A sample is drawn from this SEM, and we evaluate the model predictions on this sample, both with and without the graph structure as input\. See Appendix[G\.1](https://arxiv.org/html/2606.26467#A7.SS1)for details\. Figures[2](https://arxiv.org/html/2606.26467#S4.F2)and[5](https://arxiv.org/html/2606.26467#A7.F5)show the predicted and exact distributions for the three query types, without and with the prior structure respectively\. The model predictions align well in all cases\.

![Refer to caption](https://arxiv.org/html/2606.26467v1/x1.png)

![Refer to caption](https://arxiv.org/html/2606.26467v1/x2.png)

![Refer to caption](https://arxiv.org/html/2606.26467v1/x3.png)

Figure 2:Observational \(left\), Interventional \(center\), and Counterfactual \(right\) distributions for the IV example\. Exact solutions are in blue, model predictions in orange for whenUUis observed, and in green whenUUis unobserved\. Model is not given the graph structure\.Table[3](https://arxiv.org/html/2606.26467#A0.T3)shows the predicted adjacency matrix\. The predicted adjacency matrix closely matches the true adjacency matrix, indicating that the model has effectively learned the underlying causal structure of the SEM\.

A more complex nonlinear SEM is investigated in Appendix[G\.2](https://arxiv.org/html/2606.26467#A7.SS2)and Figure[7](https://arxiv.org/html/2606.26467#A7.F7)shows how predictions improve as the number of fit samples increases\.

### 4\.2Synthetic SEMs

Our model is evaluated on a large number of synthetic SEMs drawn from our prior dataset distribution as well as out of distribution \(OOD\) priors\. For consistency, 500 datasets are drawn and saved for each setting used for evaluation\. The OOD dataset is generated with random Fourier functions as nonlinearities instead of neural networks, see Appendix[H](https://arxiv.org/html/2606.26467#A8)\. This tests our model’s ability to generalise to unseen functional forms\. The model is compared both with and without the true graph structure as input\.

We evaluate the model’s prediction mean squared error \(MSE\) for observational, interventional and counterfactual queries, and accuracy and area under the receiver operating characteristic curve \(AUC\) of the graph structure prediction\.

For outcome prediction, we compare against Do\-PFN and established meta\-learners, the S\-learner, T\-learner, X\-learner and Doubly Robust \(DR\) learner, all using TabPFNv2\.5 as the base model\(Künzelet al\.,[2019](https://arxiv.org/html/2606.26467#bib.bib18); Battocchiet al\.,[2019](https://arxiv.org/html/2606.26467#bib.bib19)\)\. The baseline learners are unable to make counterfactual predictions\. For graph prediction, we compare against AVICI\(Lorchet al\.,[2022](https://arxiv.org/html/2606.26467#bib.bib14)\), a strong deep\-learning graph prediction model as well as established methods FCI\(Spirteset al\.,[1995](https://arxiv.org/html/2606.26467#bib.bib20)\), GES\(Chickering,[2002](https://arxiv.org/html/2606.26467#bib.bib21)\), LiNGAM\(Shimizuet al\.,[2006](https://arxiv.org/html/2606.26467#bib.bib22)\), and PC\(Spirteset al\.,[2000](https://arxiv.org/html/2606.26467#bib.bib23)\)\. Only FCI can predict confounding\.

Table 1:Predictions on OOD SEMs withn=1024n=1024fit samples, compared by mean squared error \(↓\\downarrow\)\. The baseline methods are unable to make counterfactual predictions\.Tables[12](https://arxiv.org/html/2606.26467#A8.T12)and[1](https://arxiv.org/html/2606.26467#S4.T1)summarise outcome prediction results for in\-distribution and out\-of\-distribution SEMs, respectively\. Our model generalises well to out of distribution SEMs, indicating it is not overfit to our prior\. In the observational setting, TabPFN\-CFM is comparable to the baselines, likely due to the strong TabPFNv2\.5 base model used by the Meta\-Learners\. Our model outperforms all the baselines on the interventional setting, even without the graph structure\. Giving the model the true graph structure improves predictions in the causal interventional and counterfactual settings but not the observational setting, as causal direction does not affect the observational distribution\.

Table 2:Structural predictions on synthetic OOD SEMs withn=1024n=1024fit samples, evaluated on predicted adjacency, ancestral and confounding matrices, averaged over the dataset\. Main values tracks the average AUROC \(↑\\uparrow\), with accuracy \(%\)\(↑\\uparrow\) in brackets\.Tables[13](https://arxiv.org/html/2606.26467#A8.T13)and[2](https://arxiv.org/html/2606.26467#S4.T2)summarise the structure prediction results where our model outperforms the baselines\. As well as directly comparing the predicted graphs, we also use the predicted adjacency matrix to generate an ancestral matrix using monte\-carlo sampling in Table[8](https://arxiv.org/html/2606.26467#A0.T8)\. This is significantly less accurate than directly using the predicted ancestral matrix, showing the advantage of directly predicting the ancestral structure rather than relying on the adjacency matrix\. Full results with varying fit samples are shown in Table[14](https://arxiv.org/html/2606.26467#A8.T14),[15](https://arxiv.org/html/2606.26467#A8.T15),[16](https://arxiv.org/html/2606.26467#A8.T16)and[17](https://arxiv.org/html/2606.26467#A8.T17)\.

### 4\.3Real Data

Our model is evaluated on two real datasets, the Amazon Sales dataset\(Blöbaumet al\.,[2024](https://arxiv.org/html/2606.26467#bib.bib9)\)and the Law School Admissions dataset\(Wightman,[1998](https://arxiv.org/html/2606.26467#bib.bib10)\)\. These datasets have established causal graphs and counterfactuals generated with the DoWhy framework, taken fromRobertsonet al\.\([2025](https://arxiv.org/html/2606.26467#bib.bib16)\)\. The results are shown in Table[4](https://arxiv.org/html/2606.26467#A0.T4)and Table[5](https://arxiv.org/html/2606.26467#A0.T5)\. While all models are similar in their observational performance, our model excels in the interventional settings\. In the Law School Admissions dataset, our model is able to use the counterfactual information to significantly improve its accuracy\.

## 5Discussion

We have introduced TabPFN\-CFM, a causal foundation model that handles multiple causal tasks, including both causal structure and outcome prediction\. We believe causal PFNs can continue to be extended to more complex settings, such as time series and cyclic causal settings\.

## References

- V\. Balazadeh, H\. Kamkari, V\. Thomas, B\. Li, J\. Ma, J\. C\. Cresswell, and R\. G\. Krishnan \(2025\)CausalPFN: amortized causal effect estimation via in\-context learning\.InAdvances in Neural Information Processing Systems,Cited by:[§1\.1](https://arxiv.org/html/2606.26467#S1.SS1.p3.4),[§1](https://arxiv.org/html/2606.26467#S1.p1.1),[§3\.2](https://arxiv.org/html/2606.26467#S3.SS2.p1.18)\.
- K\. Battocchi, E\. Dillon, M\. Hei, G\. Lewis, P\. Oka, M\. Oprescu, and V\. Syrgkanis \(2019\)EconML a python package for ml\-based heterogeneous treatment effects estimation\.External Links:[Link](https://github.com/py-why/EconML)Cited by:[§4\.2](https://arxiv.org/html/2606.26467#S4.SS2.p3.1)\.
- P\. Blöbaum, P\. Götz, K\. Budhathoki, A\. A\. Mastakouri, and D\. Janzing \(2024\)DoWhy\-gcm: an extension of dowhy for causal inference in graphical causal models\.Cited by:[§4\.3](https://arxiv.org/html/2606.26467#S4.SS3.p1.1)\.
- D\.M\. Chickering \(2002\)Optimal structure identification with greedy search\. journal of machine learning research\.Cited by:[§4\.2](https://arxiv.org/html/2606.26467#S4.SS2.p3.1)\.
- X\. Glorot and Y\. Bengio \(2010\)Understanding the difficulty of training deep feedforward neural networks\.InInternational Conference on Artificial Intelligence,Cited by:[§D\.1](https://arxiv.org/html/2606.26467#A4.SS1.p2.2)\.
- A\. Henry, P\. R\. Dachapally, S\. S\. Pawar, and Y\. Chen \(2020\)Query\-key normalization for transformers\.InFindings of the Association for Computational Linguistics: EMNLP 2020,Cited by:[8th item](https://arxiv.org/html/2606.26467#A5.I1.i8.p1.1)\.
- N\. Hollmann, S\. Müller, K\. Eggensperger, and F\. Hutter \(2023\)TabPFN: a transformer that solves small tabular classification problems in a second\.InInternational Conference on Artificial Intelligence,Cited by:[Appendix B](https://arxiv.org/html/2606.26467#A2.p1.3),[§1](https://arxiv.org/html/2606.26467#S1.p1.1),[§3\.2](https://arxiv.org/html/2606.26467#S3.SS2.p1.18)\.
- G\. W\. Imbens and D\. B\. Rubin \(2015\)Causal inference for statistics, social, and biomedical sciences: an introduction\.Cambridge University Press\.Cited by:[§1](https://arxiv.org/html/2606.26467#S1.p1.1)\.
- K\. Jordan, J\. Bernstein, B\. Rappazzo, @fernbear\.bsky\.social, B\. Vlado, Y\. Jiacheng, F\. Cesista, B\. Koszarsky, and @Grad62304977 \(2024a\)Modded\-nanogpt: speedrunning the nanogpt baseline\.Cited by:[§E\.5](https://arxiv.org/html/2606.26467#A5.SS5.p1.1)\.
- K\. Jordan, Y\. Jin, V\. Boza, Y\. Jiacheng, F\. Cesista, L\. Newhouse, and J\. Bernstein \(2024b\)Muon: an optimizer for hidden layers in neural networks\.Cited by:[7th item](https://arxiv.org/html/2606.26467#A5.I1.i7.p1.1)\.
- N\. R\. Ke, S\. Chiappa, J\. X\. Wang, J\. Bornschein, A\. Goyal, M\. Rey, T\. Weber, M\. Botvinick, M\. C\. Mozer, and D\. J\. Rezende \(2023\)Learning to induce causal structure\.InInternational Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2606.26467#S1.p1.1),[§3\.1](https://arxiv.org/html/2606.26467#S3.SS1.p2.4)\.
- S\. R\. Künzel, J\. S\. Sekhon, P\. J\. Bickel, and B\. Yu \(2019\)Metalearners for estimating heterogeneous treatment effects using machine learning\.Cited by:[§1](https://arxiv.org/html/2606.26467#S1.p1.1),[§4\.2](https://arxiv.org/html/2606.26467#S4.SS2.p3.1)\.
- J\. Liu, J\. Su, X\. Yao, Z\. Jiang, G\. Lai, Y\. Du, Y\. Qin, W\. Xu, E\. Lu, J\. Yan, Y\. Chen, H\. Zheng, Y\. Liu, S\. Liu, B\. Yin, W\. He, H\. Zhu, Y\. Wang, J\. Wang, M\. Dong, Z\. Zhang, Y\. Kang, H\. Zhang, X\. Xu, Y\. Zhang, Y\. Wu, X\. Zhou, and Z\. Yang \(2025\)Muon is scalable for llm training\.Cited by:[7th item](https://arxiv.org/html/2606.26467#A5.I1.i7.p1.1)\.
- L\. Lorch, S\. Sussex, J\. Rothfuss, A\. Krause, and B\. Schölkopf \(2022\)Amortized inference for causal structure learning\.InAdvances in Neural Information Processing Systems,Cited by:[§E\.2](https://arxiv.org/html/2606.26467#A5.SS2.p2.12),[§1](https://arxiv.org/html/2606.26467#S1.p1.1),[§4\.2](https://arxiv.org/html/2606.26467#S4.SS2.p3.1)\.
- J\. Pearl \(2000\)Causality : models, reasoning, and inference\.Cambridge University Press\.Cited by:[Appendix C](https://arxiv.org/html/2606.26467#A3.p1.1),[§1](https://arxiv.org/html/2606.26467#S1.p1.1)\.
- T\. Richardson \(2003\)Markov properties for acyclic directed mixed graphs\.Cited by:[Appendix C](https://arxiv.org/html/2606.26467#A3.p3.7)\.
- J\. Robertson, A\. Reuter, S\. Guo, N\. Hollmann, F\. Hutter, and B\. Schölkopf \(2025\)Do\-PFN: in\-context learning for causal effect estimation\.InAdvances in Neural Information Processing Systems,Cited by:[Appendix B](https://arxiv.org/html/2606.26467#A2.p1.3),[§1\.1](https://arxiv.org/html/2606.26467#S1.SS1.p3.4),[§1](https://arxiv.org/html/2606.26467#S1.p1.1),[§2](https://arxiv.org/html/2606.26467#S2.p3.1),[§3\.2](https://arxiv.org/html/2606.26467#S3.SS2.p1.18),[§4\.3](https://arxiv.org/html/2606.26467#S4.SS3.p1.1)\.
- S\. Shimizu, P\. O\. Hoyer, A\. Hyvärinen, and A\. Kerminen \(2006\)A linear non\-gaussian acyclic model for causal discovery\.Cited by:[§4\.2](https://arxiv.org/html/2606.26467#S4.SS2.p3.1)\.
- P\. Spirtes, C\. Meek, and T\. Richardson \(1995\)Causal inference in the presence of latent variables and selection bias\.Cited by:[§4\.2](https://arxiv.org/html/2606.26467#S4.SS2.p3.1)\.
- P\. Spirtes, C\. N\. Glymour, R\. Scheines, and D\. Heckerman \(2000\)Causation, prediction, and search\.Cited by:[§4\.2](https://arxiv.org/html/2606.26467#S4.SS2.p3.1)\.
- L\. F\. Wightman \(1998\)LSAC national longitudinal bar passage study\.LSAC Research Report Series\.Cited by:[§4\.3](https://arxiv.org/html/2606.26467#S4.SS3.p1.1)\.
- R\. Xiong, Y\. Yang, D\. He, K\. Zheng, S\. Zheng, H\. Zhang, Y\. Lan, L\. Wang, and T\. Liu \(2020\)On layer normalization in the transformer architecture\.Cited by:[5th item](https://arxiv.org/html/2606.26467#A5.I1.i5.p1.1)\.
- Z\. Zhang, Y\. Song, G\. Yu, X\. Han, Y\. Lin, C\. Xiao, C\. Song, Z\. Liu, Z\. Mi, and M\. Sun \(2024\)ReLU2wins: discovering efficient activation functions for sparse llms\.Cited by:[3rd item](https://arxiv.org/html/2606.26467#A5.I1.i3.p1.1)\.
- Q\. Zhao \(2025\)On statistical and causal models associated with acyclic directed mixed graphs\.Cited by:[Appendix C](https://arxiv.org/html/2606.26467#A3.p4.2)\.

Table 3:Predicted adjacency matrices for the linear SEM experiment on the IV graph\. Values represent predicted edge probabilities\. The true edges areZ→TZ\\rightarrow T,T→YT\\rightarrow Y,U→TU\\rightarrow TandU→YU\\rightarrow Y\.Table 4:Comparison of our model with baselines on the Amazon Sales dataset, compared by mean squared error \(↓\\downarrow\)\. The baseline methods are unable to use the true causal graph or make counterfactual predictions\.Table 5:Comparison of our model with baselines on the Law School Admissions dataset, compared by mean squared error \(↓\\downarrow\)\. The baseline methods are unable to make counterfactual predictions\.Table 6:Structural predictions on the Amazon Sales dataset, evaluated on predicted adjacency matrix and ancestral matrix\. Main values tracks the AUROC \(↑\\uparrow\), with accuracy \(↑\\uparrow\)in brackets\.Table 7:Structural predictions on the Law School Admissions dataset, evaluated on predicted adjacency matrix and ancestral matrix\. Main values tracks the AUROC \(↑\\uparrow\), with accuracy \(↑\\uparrow\)in brackets\.Table 8:Comparison of direct ancestral matrix prediction and ancestral matrix generated from the predicted adjacency matrix using Monte\-Carlo sampling for out\-of\-distribution SEMs, across different numbers of fit data points \(n\-fit\)\. Values represent AUC \(↑\\uparrow\)\.## Appendix AComparing estimation with and without graph input

Including the graph input can improve model predictions\. For simplicity, ignorecchere\. With and without𝒢\\mathcal\{G\}, the target distributions are

p0​\(y\|Dobs\)\\displaystyle p\_\{0\}\(y\|D^\{\\mathrm\{obs\}\}\)=∫p​\(y\|ψ\)​p​\(ψ\|Dobs\)​𝑑ψ\\displaystyle=\\int p\(y\|\\psi\)p\(\\psi\|D^\{\\mathrm\{obs\}\}\)\\ d\\psi\(8\)pG​\(y\|Dobs,𝒢\)\\displaystyle p\_\{G\}\(y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)=∫p​\(y\|ψ\)​p​\(ψ\|Dobs,𝒢\)​𝑑ψ\\displaystyle=\\int p\(y\|\\psi\)p\(\\psi\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ d\\psi\(9\)Comparing their losses, we have:

Ey\|Dobs,𝒢​\[−log⁡pG​\(y\|Dobs,𝒢\)−\(−log⁡p0​\(y\|Dobs\)\)\]\\displaystyle E\_\{y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\}\\left\[\-\\log p\_\{G\}\(y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\-\(\-\\log p\_\{0\}\(y\|D^\{\\mathrm\{obs\}\}\)\)\\right\]=∫−pG​\(y\|Dobs,𝒢\)​log⁡pG​\(y\|Dobs,𝒢\)p0​\(y\|Dobs\)​d​y\\displaystyle=\\int\-p\_\{G\}\(y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\log\\frac\{p\_\{G\}\(y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\}\{p\_\{0\}\(y\|D^\{\\mathrm\{obs\}\}\)\}\\ dy\(11\)=−KL\(pG\(y\|Dobs,𝒢\)\|\|p0\(y\|Dobs\)\)≤0\\displaystyle=\-KL\(p\_\{G\}\(y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\|\|p\_\{0\}\(y\|D^\{\\mathrm\{obs\}\}\)\)\\leq 0\(12\)Now, taking expectation overDobsD^\{\\mathrm\{obs\}\}and𝒢\\mathcal\{G\}, we have

LG−L0\\displaystyle L\_\{G\}\-L\_\{0\}=Ey,Dobs,𝒢​\[−log⁡pG​\(y\|Dobs,𝒢\)−\(−log⁡p0​\(y\|Dobs\)\)\]\\displaystyle=E\_\{y,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\}\\left\[\-\\log p\_\{G\}\(y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\-\(\-\\log p\_\{0\}\(y\|D^\{\\mathrm\{obs\}\}\)\)\\right\]\(13\)=−Ey,Dobs,𝒢\[KL\(pG\(y\|Dobs,𝒢\)\|\|p0\(y\|Dobs\)\)\]≤0\\displaystyle=\-E\_\{y,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\}\\left\[KL\(p\_\{G\}\(y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\|\|p\_\{0\}\(y\|D^\{\\mathrm\{obs\}\}\)\)\\right\]\\leq 0\(14\)Therefore, including the graph input can only improve the model predictions, with no improvement occurring ifpG​\(y\|Dobs,𝒢\)=p0​\(y\|Dobs\)p\_\{G\}\(y\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)=p\_\{0\}\(y\|D^\{\\mathrm\{obs\}\}\)\. Equivalently, in terms of the posterior ofψ\\psi, there is an improvement if

∫p​\(y\|ψ\)​p​\(ψ\|Dobs,𝒢\)​𝑑ψ≠∫p​\(y\|ψ\)​p​\(ψ\|Dobs\)​𝑑ψ\\displaystyle\\int p\(y\|\\psi\)p\(\\psi\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ d\\psi\\neq\\int p\(y\|\\psi\)p\(\\psi\|D^\{\\mathrm\{obs\}\}\)\\ d\\psi\(15\)∫p​\(y\|ψ\)​\[p​\(ψ\|Dobs,𝒢\)−p​\(ψ\|Dobs\)\]​𝑑ψ≠0\\displaystyle\\int p\(y\|\\psi\)\\left\[p\(\\psi\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\-p\(\\psi\|D^\{\\mathrm\{obs\}\}\)\\right\]\\ d\\psi\\neq 0\(16\)p​\(ψ\|Dobs,𝒢\)≠p​\(ψ\|Dobs\)\\displaystyle p\(\\psi\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\neq p\(\\psi\|D^\{\\mathrm\{obs\}\}\)\(17\)So an improvement occurs if including𝒢\\mathcal\{G\}changes the posterior distribution ofψ\\psi\(whereyyhas support\)\. The notable case where this does not occur is when the graph structure is uniquely identified by the observational data,p​\(ψ\|Dobs\)p\(\\psi\|D^\{\\mathrm\{obs\}\}\)is a delta distribution on the true SCM, andp​\(ψ\|Dobs,𝒢\)p\(\\psi\|D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)is the same delta distribution, the SEM is identifiable\. In the finite data regime, even if the SEM is identifiable, the posterior may not be perfect, so including the graph input can still improve the model predictions\.

## Appendix BAppendix: Proof of equivalence between log probability and KL divergence

Proof minimising log probability \(and cross\-entropy\) equals minimising KL divergence\. This proof extends the proof given in TabPFN\(Hollmannet al\.,[2023](https://arxiv.org/html/2606.26467#bib.bib15)\)and Do\-PFN\(Robertsonet al\.,[2025](https://arxiv.org/html/2606.26467#bib.bib16)\)by including the conditioning on the graph input\. We have the following definitions and properties:

ψ\\displaystyle\\psi∼p​\(ψ\)\\displaystyle\\sim p\(\\psi\)\(19\)Dobs\\displaystyle D^\{\\mathrm\{obs\}\}∼p​\(Dobs\|ψ\)\\displaystyle\\sim p\(D^\{\\mathrm\{obs\}\}\|\\psi\)\(20\)𝒢\\displaystyle\\mathcal\{G\}∼p​\(𝒢\|ψ\)\\displaystyle\\sim p\(\\mathcal\{G\}\|\\psi\)\(21\)c\\displaystyle c∼p​\(c\|ψ\)\\displaystyle\\sim p\(c\|\\psi\)\(22\)y\\displaystyle y∼p​\(y\|c,ψ\)\\displaystyle\\sim p\(y\|c,\\psi\)\(23\)L\\displaystyle L=−E​\[log⁡p^θ​\(y\|c,Dobs,𝒢\)\]\\displaystyle=\-E\[\\log\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\]\(24\)\{c,y\}\\displaystyle\\\{c,y\\\}⟂⟂Dobs\|ψ\\displaystyle\\mathrel\{\\perp\\mspace\{\-10\.0mu\}\\perp\}D^\{\\mathrm\{obs\}\}\|\\psi\(25\)\{c,y\}\\displaystyle\\\{c,y\\\}⟂⟂𝒢\|ψ\\displaystyle\\mathrel\{\\perp\\mspace\{\-10\.0mu\}\\perp\}\\mathcal\{G\}\|\\psi\(26\)Optimising the loss functionLLis equivalent to optimising the KL divergence between the true distribution and the model distribution\. We have,

L\\displaystyle L=−E​\[log⁡p^θ​\(y\|c,Dobs,𝒢\)\]\\displaystyle=\-E\[\\log\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\]\(28\)=−∫∫∫log⁡p^θ​\(y\|c,Dobs,𝒢\)​p​\(y,c,Dobs,𝒢\)​𝑑y​𝑑c​𝑑Dobs​𝑑𝒢\\displaystyle=\-\\int\\int\\int\\log\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)p\(y,c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ dy\\ dc\\ dD^\{\\mathrm\{obs\}\}d\\mathcal\{G\}\(29\)=−∫∫∫log⁡p^θ​\(y\|c,Dobs,𝒢\)​p​\(y\|c,Dobs,𝒢\)​p​\(c,Dobs,𝒢\)​𝑑y​𝑑c​𝑑Dobs​𝑑𝒢\\displaystyle=\-\\int\\int\\int\\log\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)p\(c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ dy\\ dc\\ dD^\{\\mathrm\{obs\}\}d\\mathcal\{G\}\(30\)=−Ec,Dobs,𝒢​\[∫p​\(y\|c,Dobs,𝒢\)​log⁡p^θ​\(y\|c,Dobs,𝒢\)​𝑑y\]\\displaystyle=\-E\_\{c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\}\\left\[\\int p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\log\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ dy\\right\]\(31\)=−Ec,Dobs,𝒢​\[∫p​\(y\|c,Dobs,𝒢\)​\[log⁡p^θ​\(y\|c,Dobs,𝒢\)p​\(y\|c,Dobs,𝒢\)\+log⁡p​\(y\|c,Dobs,𝒢\)\]​𝑑y\]\\displaystyle=\-E\_\{c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\}\\left\[\\int p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\left\[\\log\\frac\{\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\}\{p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\}\+\\log p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\right\]\\ dy\\right\]\(32\)=Ec,Dobs,𝒢\[KL\(p\(y\|c,Dobs,𝒢\)\|\|p^θ\(y\|c,Dobs,𝒢\)\)\+C\)\]\\displaystyle=E\_\{c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\}\\left\[KL\(p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\|\|\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\)\+C\)\\right\]\(33\)hence minimising the log probability is equivalent to minimising the KL divergence between the true distribution and the model distribution\. This can also be written withψ\\psimore explicitly\. Using the independence property, we have

p​\(y\|c,Dobs,𝒢\)\\displaystyle p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)=∫p​\(y\|c,Dobs,𝒢,ψ\)​p​\(ψ\|c,Dobs,𝒢\)​𝑑ψ\\displaystyle=\\int p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\},\\psi\)p\(\\psi\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ d\\psi\(35\)=∫p​\(y\|c,ψ\)​p​\(ψ\|c,Dobs,𝒢\)​𝑑ψ\.\\displaystyle=\\int p\(y\|c,\\psi\)p\(\\psi\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ d\\psi\.\(36\)Hence,

L\\displaystyle L=−Ec,Dobs,𝒢​\[∫p​\(y\|c,Dobs,𝒢\)​log⁡p^θ​\(y\|c,Dobs,𝒢\)​𝑑y\]\\displaystyle=\-E\_\{c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\}\\left\[\\int p\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\log\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ dy\\right\]\(38\)=−Ec,Dobs,𝒢​\[∫∫p​\(y\|c,ψ\)​p​\(ψ\|c,Dobs,𝒢\)​log⁡p^θ​\(y\|c,Dobs,𝒢\)​𝑑y​𝑑ψ\]\\displaystyle=\-E\_\{c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\}\\left\[\\int\\int p\(y\|c,\\psi\)p\(\\psi\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\log\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ dy\\ d\\psi\\right\]\(39\)=−Ec,Dobs,𝒢,ψ​\[∫p​\(y\|c,ψ\)​log⁡p^θ​\(y\|c,Dobs,𝒢\)​𝑑y\]\\displaystyle=\-E\_\{c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\},\\psi\}\\left\[\\int p\(y\|c,\\psi\)\\log\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\\ dy\\ \\right\]\(40\)=Ec,Dobs,𝒢,ψ\[KL\(p\(y\|c,ψ\)\|\|p^θ\(y\|c,Dobs,𝒢\)\)\+C\)\]\\displaystyle=E\_\{c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\},\\psi\}\\left\[KL\(p\(y\|c,\\psi\)\|\|\\hat\{p\}\_\{\\theta\}\(y\|c,D^\{\\mathrm\{obs\}\},\\mathcal\{G\}\)\)\+C\)\\right\]\(41\)

## Appendix CCausal Graphs and Structural Causal Models

A*nonparametric structural equation model with independent errors*\(NPSEM\-IE\) satisfies the following conditions\(Pearl,[2000](https://arxiv.org/html/2606.26467#bib.bib11)\):

1. 1\.For eachi=1,…,di=1,\\dots,d, there exists a set of parent variablesPai⊆V∖\{Vi\}\\mathrm\{Pa\}\_\{i\}\\subseteq V\\setminus\\\{V\_\{i\}\\\}such that Vi=fi​\(Pai,Ui\)\.V\_\{i\}=f\_\{i\}\\bigl\(\\mathrm\{Pa\}\_\{i\},U\_\{i\}\\bigr\)\.
2. 2\.The graph with vertex setVVand directed edgesVj→ViV\_\{j\}\\to V\_\{i\}wheneverVj∈PaiV\_\{j\}\\in\\mathrm\{Pa\}\_\{i\}is acyclic\.
3. 3\.The unobserved variables are mutually independent and exogenous: Ui⟂⟂Ujfor​i≠j\.U\_\{i\}\\mathrel\{\\perp\\mspace\{\-10\.0mu\}\\perp\}U\_\{j\}\\quad\\text\{for \}i\\neq j\.

If our causal system satisfies an NPSEM\-IE \(or can be trivially reduced to one\), corresponding to only ”noise” variables being unobserved, then the causal graph corresponding to observed variablesVVwill be a DAG structure\. Each variable is independent of its non\-descendants given its parents \(the local Markov property\), with the direction of edges indicating the direction of causality\.

However, in the more general case of a DAG over variables\(U,V\)\(U,V\), the NPSEM\-IE assumptions need not hold\. The unobserved variables inUUmay be endogenous \(i\.e\., have parents inVVor inUU\), a single unobserved variable may have multiple children inVV, and the unobserved variables may be statistically dependent\. After marginalizing outUU, the causal relationships among the observed variablesVVcan be represented by an acyclic directed mixed graph \(ADMG\)\(Richardson,[2003](https://arxiv.org/html/2606.26467#bib.bib12)\), which contains both directed edges \(indicating directed causal influence\) and bidirected edges \(indicating the presence of unobserved common causes, i\.e\. unobserved confounding\)\. Such cases are common in real\-world scenarios, where not all relevant variables can be observed or measured\.

There are many models referred to as ADMGs in the literature\(Zhao,[2025](https://arxiv.org/html/2606.26467#bib.bib13)\); here we will refer to ADMGs as DAGs with unobserved variables marginalized out\. An ADMG graph consists of an adjacency matrix,AA, and correlation matrix,CC\. The adjacency matrix contains directed causal relationships, while the symmetric correlation matrix contains bidirected non\-causal correlations caused by unobserved confounding\.

Finally, we note there is a more general case, where the underlying structure is not a DAG and instead contains cycles\. Such structures can arise in systems with temporal feedback loops or cyclic causality\. However, in this work we focus on the acyclic case \(DAGs and ADMGs\) and leave cyclic structures for future work\.

## Appendix DDataset generation

### D\.1Generating random SCMs

Firstly, the DAG structure is generated\. The DAG size, average connectivity degree and type of DAG are sampled\. The type of DAG is either Erdos\-Renyi or Scale\-Free\. For Scale\-Free graphs, the attractiveness parameter is also sampled\. Given these parameters, a random DAG is generated\. A treatment node and target node are sampled, with bias towards selecting the treatment as a causal ancestor of the target to ensure meaningful interventions\.

The structural equations are sampled for every node in the DAG\. First, the nonlinearity is sampled from a set of possible functions, including whether to use an MLP \(Table[9](https://arxiv.org/html/2606.26467#A4.T9)\)\. Then, the type of noise, position of noise and noise distribution are sampled \(Table[10](https://arxiv.org/html/2606.26467#A4.T10)\)\. Relevant distribution parameters \(e\.g\. variance\) are also sampled randomly\. Finally, the node’s weight vector𝐖\\mathbf\{W\}and bias scalarbbare drawn from the Xavier Normal distribution\(Glorot and Bengio,[2010](https://arxiv.org/html/2606.26467#bib.bib1)\)\.

The node’s structural equation takes as input the concatenation of its parent values,𝐱pa\\mathbf\{x\}\_\{\\mathrm\{pa\}\}, and computes the node’s value in the most general form as:

xout=g​\(𝐖2​fnonlinear​\(𝐖1⋅g​\(𝐱pa,ϵpre\)\+b\),ϵpost\)x\_\{\\mathrm\{out\}\}=g\(\\mathbf\{W\}\_\{2\}f\_\{\\mathrm\{nonlinear\}\}\(\\mathbf\{W\}\_\{1\}\\cdot g\(\\mathbf\{x\}\_\{\\mathrm\{pa\}\},\\epsilon\_\{\\mathrm\{pre\}\}\)\+b\),\\epsilon\_\{\\mathrm\{post\}\}\)\(43\)Whereffis the elementwise nonlinearity,ggis the noise composition function \(either elementwise additive or multiplicative\),ϵpre\\epsilon\_\{\\mathrm\{pre\}\}andϵpost\\epsilon\_\{\\mathrm\{post\}\}are noise variables sampled from the pre and post noise distribution\. If the node has no parents,𝐱pa\\mathbf\{x\}\_\{\\mathrm\{pa\}\}is empty and only noise is used as input to the node\.

The treatment node is generated differently, since it is a binary variable\. The treatment node’s structural equation is:

x~out=𝐖⋅\(𝐱pa\+ϵpre\)\\displaystyle\\tilde\{x\}\_\{\\mathrm\{out\}\}=\\mathbf\{W\}\\cdot\(\\mathbf\{x\}\_\{\\mathrm\{pa\}\}\+\\epsilon\_\{\\mathrm\{pre\}\}\)\(44\)xout=a⋅𝟙​\[x~out\>xthresh\]\+b\\displaystyle x\_\{\\mathrm\{out\}\}=a\\cdot\\mathbbm\{1\}\[\\tilde\{x\}\_\{\\mathrm\{out\}\}\>x\_\{\\mathrm\{thresh\}\}\]\+b\(45\)wherea,ba,bare random scalars andxthreshx\_\{\\mathrm\{thresh\}\}is set so the observational distribution ofx~out\\tilde\{x\}\_\{\\mathrm\{out\}\}has a randomly sampled class balance\. This ensures the treatment always contains positive and negative samples\.

Finally, a subset of nodes in the DAG are selected as observed variablesVV, with the remaining nodes treated as unobserved variablesUU\. The selection is random, with the constraint that the treatment and target nodes are always observed\. Unobserved variables are marginalized out when generating datasets, creating an observed acyclic adjacency matrix,AA, and bidirected confounding matrix,CC, between observed nodes together forming the ADMG structure\. If the sample is trivial \(e\.g\. all treatments the same or no causal relations\), the sample is discarded\. This completes the specification of a single SCM,ψ\\psi\. Our synthetic prior distribution of SCMs,P​\(ψ\)P\(\\psi\), is the distribution induced by the above generative process with all parameters sampled randomly\.

### D\.2Sampling causal datasets from SCMs

Now, we describe how each sample dataset is generated\. A SCMψ\\psiis sampled from the priorP​\(ψ\)P\(\\psi\)\. Then, for each sampleiiin the dataset, the exogenous noise variablesUUare sampled from their respective distributions\. The observed variablesVVare generated by evaluating the structural equations in topological order\. This process is repeatednntimes to generate the observational dataset withnnentries,𝒟fit=\{𝐱i,ti,yi\}i=1n\\mathcal\{D\}^\{\\mathrm\{fit\}\}=\\\{\\mathbf\{x\}\_\{i\},t\_\{i\},y\_\{i\}\\\}\_\{i=1\}^\{n\}\. The underlying graph is stored inGest=\{A,C\}G^\{\\mathrm\{est\}\}=\\\{A,C\\\}\.

Then, we generate the prediction dataset\. Another sample is drawn from the observational dataset as above,𝒟pred=\{𝐱∗,t∗,y∗\}\\mathcal\{D\}^\{\\mathrm\{pred\}\}=\\\{\\mathbf\{x\}^\{\*\},t^\{\*\},y^\{\*\}\\\}\. Since we know the true SCM, the counterfactual sample can also be generated\. However, we must be careful to correctly simulate the counterfactual since nodes betweenTTandYY\(observed or unobserved\) may be descendants ofTT; these nodes should not be fixed\. To account for this, we fix every node and noise variable that may be topologically ordered beforeTT\. That is, we split the nodes into descendants ofTT,NdescN\_\{\\mathrm\{desc\}\}, and non\-descendants,NndN\_\{\\mathrm\{nd\}\}\. All of the noise variables and nodes inNndN\_\{\\mathrm\{nd\}\}are kept fixed to their sampled values from the observational sample\. The treatment node is set to its interventional valueT=do​\(1−t∗\)T=\\mathrm\{do\}\(1\-t^\{\*\}\)\. Nodes inNdescN\_\{\\mathrm\{desc\}\}are recomputed with interventionalTT\. This procedure generates the ”what\-if” counterfactual outcome\. Nodes inVVbetweenTTandYYwill change under intervention, but their counterfactual values will depend on their observed value through fixed parent noise variables, allowing the model to use this information to make more accurate predictions\. We denote this new sample as𝒟causal=\{𝐱t∗,t=do\(1−t∗\),yt∗\}\\mathcal\{D\}^\{\\mathrm\{causal\}\}=\\\{\\mathbf\{x\}^\{\*\}\_\{t\},t=\\mathrm\{do\}\(1\-t^\{\*\}\),y^\{\*\}\_\{t\}\\\}\. We note there are alternative methods of generating causal queries, such as fixing only ancestors ofTTcorresponding to interventional rather than counterfactual distribution\. These are not used in this paper, but we allow users to easily generate and train new models on these alternative distributions with our codebase by changing a configuration flag\.

Our model is trained with up ton=1600n=1600fit samples andk=33k=33covariates \(35 nodes in total\)\.

Table 9:Nonliearities\. Outputs are scaled to be approxmately standardized\. Quantize and Square nonlinearities modified for stability\.Table 10:Noise generation parameters

## Appendix EModel architecture

Our architecture is based on TabPFNv2 and Do\-PFN, with several upgrades to handle the additional objectives and improve training efficiency\. The model consists of four components, the encoder module, the observation module, the graph module and the prediction module\.

### E\.1Encoder module

The encoder module embeds the inputs\. For notational simplicity, we define / describe the shapes as follows,𝐱fit∈ℝn×k\\mathbf\{x\}^\{\\mathrm\{fit\}\}\\in\\mathbb\{R\}^\{n\\times k\},𝐱∗∈ℝ1×k\\mathbf\{x\}^\{\*\}\\in\\mathbb\{R\}^\{1\\times k\},tfit∈\{0,1\}nt^\{\\mathrm\{fit\}\}\\in\\\{0,1\\\}^\{n\},t∗∈\{0,1\}×\{0,1\}t^\{\*\}\\in\\\{0,1\\\}\\times\\\{0,1\\\},yfit∈ℝny^\{\\mathrm\{fit\}\}\\in\\mathbb\{R\}^\{n\},G∈\{0,1\}\(k\+2\)×\(k\+2\)G\\in\\\{0,1\\\}^\{\(k\+2\)\\times\(k\+2\)\}, wherekkis the number of covariates, andnnis the number of observations\. The hidden dimension of the transformer ishh\. Both𝐱fit\\mathbf\{x\}^\{\\mathrm\{fit\}\}and𝐱∗\\mathbf\{x\}^\{\*\}are encoded with the same linear transform, followed by a column\-wise positional embedding:

hax\\displaystyle h\_\{a\}^\{x\}=Linearx\(𝐱\|\|CDF\(𝐱\)\)∈ℝn×k×h\\displaystyle=\\text\{Linear\}\_\{x\}\(\\mathbf\{x\}\|\|\\text\{CDF\}\(\\mathbf\{x\}\)\)\\in\\mathbb\{R\}^\{n\\times k\\times h\}\(46\)hx\\displaystyle h^\{x\}=hax\+PosEmbcol∈ℝn×k×h\.\\displaystyle=h\_\{a\}^\{x\}\+\\text\{PosEmb\}\_\{\\mathrm\{col\}\}\\in\\mathbb\{R\}^\{n\\times k\\times h\}\.\(47\)Each covariate is concatenated with its CDF transform, which helps embedding extreme outliers\. Each of thekkfeature columns of the dataset is given a positional embedding generated by the subspace method from TabPFNv2\. For each columnjj

PosEmbcol​j\\displaystyle\\text\{PosEmb\}\_\{\\mathrm\{col\}\\ j\}=Wcol​ϵ\\displaystyle=W\_\{\\mathrm\{col\}\}\\epsilon\(48\)ϵ\\displaystyle\\epsilon∼𝒩​\(0,Ih/4\)\\displaystyle\\sim\\mathcal\{N\}\(0,I\_\{h/4\}\)\(49\)Wcol\\displaystyle W\_\{\\mathrm\{col\}\}∈ℝh×h/4\.\\displaystyle\\in\\mathbb\{R\}^\{h\\times h/4\}\.\(50\)One embedding is generated for each of thekkcolumns and concatenated together\. A new positional embedding is generated for each forward pass, meaning the positional embeddings are randomised, for both training and inference\. The targetsyfity^\{\\mathrm\{fit\}\}andy∗y^\{\*\}are embedded with a linear transform,

hy\\displaystyle h^\{y\}=Lineary​\(y​‖CDF​\(y\)‖​Itype\)∈ℝn×hy∈\{yfit,y∗\}\\displaystyle=\\text\{Linear\}\_\{y\}\(y\|\|\\mathrm\{CDF\}\(y\)\|\|I\_\{\\mathrm\{type\}\}\)\\in\\mathbb\{R\}^\{n\\times h\}\\qquad y\\in\\\{y^\{\\mathrm\{fit\}\},y^\{\*\}\\\}\(51\)whereItype∈\{0,1,2\}I\_\{\\mathrm\{type\}\}\\in\\\{0,1,2\\\}is an indicator for the type ofyy, with value0foryfity^\{\\mathrm\{fit\}\},11for counterfactual queries wherey∗y^\{\*\}is observed, and22for counterfactual queries wherey∗y^\{\*\}is unobserved\. The CDF embedding is used here again\. Next, we embed the treatment variables,

ht\\displaystyle h^\{t\}=embt​\(tfit\)∈ℝn×1×h\\displaystyle=\\text\{emb\}\_\{t\}\(t^\{\\mathrm\{fit\}\}\)\\in\\mathbb\{R\}^\{n\\times 1\\times h\}\(52\)ht⁣∗\\displaystyle h^\{t\*\}=embt⁣∗​\(t∗\)∈ℝ1×1×h\.\\displaystyle=\\text\{emb\}\_\{t\*\}\(t^\{\*\}\)\\in\\mathbb\{R\}^\{1\\times 1\\times h\}\.\(53\)Note the embedding fort∗t^\{\*\}has 4 values corresponding to if the query is interventional or observational\. These are concatenated to generate the full input embeddings,

h0fit\\displaystyle h^\{\\mathrm\{fit\}\}\_\{0\}=Concat​\(ht,hx,hy,dim=1\)∈ℝn×\(k\+2\)×h\\displaystyle=\\text\{Concat\}\(h^\{t\},h^\{x\},h^\{y\},\\text\{dim\}=1\)\\in\\mathbb\{R\}^\{n\\times\(k\+2\)\\times h\}\(54\)h0pred\\displaystyle h^\{\\mathrm\{pred\}\}\_\{0\}=Concat​\(ht⁣∗,hx⁣∗,𝟎,dim=1\)∈ℝ1×\(k\+2\)×h\.\\displaystyle=\\text\{Concat\}\(h^\{t\*\},h^\{x\*\},\\mathbf\{0\},\\text\{dim\}=1\)\\in\\mathbb\{R\}^\{1\\times\(k\+2\)\\times h\}\.\(55\)
Finally, the prior ADMG graph structure is embedded\. The ADMG is represented as a directed adjacency matrix,AA, and bidirected correlation matrixCC\. Both matrices are binary\. If the matrices are unknown, they are filled with zeros\. FromAA, we create the ancestral matrixRR, the matrix consisting of all ancestors for each node \(reachability or transitive closure\)\. The graph embedding is

hG\\displaystyle h^\{G\}=embA​\(A\)\+embR​\(R\)\+embC​\(C\)\+embA​T​\(AT\)\+embR​T​\(RT\)\+embC​T​\(CT\)\.\\displaystyle=\\text\{emb\}\_\{A\}\(A\)\+\\text\{emb\}\_\{R\}\(R\)\+\\text\{emb\}\_\{C\}\(C\)\+\\text\{emb\}\_\{AT\}\(A^\{T\}\)\+\\text\{emb\}\_\{RT\}\(R^\{T\}\)\+\\text\{emb\}\_\{CT\}\(C^\{T\}\)\.\(56\)whereemb​\(\)\\text\{emb\}\(\)is an elementwise learned binary embedding\. This representation is heavily redundant since the transpose matrices and ancestral matrices are constructed fromAAandCC\. This approach allows the model to easily identify every node’s parents and children as well as ancestors and descendants without multiple hops\.

Unlike TabPFNv2 \(and Do\-PFN\), we omit input augmentations: grouping features, ensembling, random augmentations and random feature products\. As well as simplifying the model, this is also required since the structure is explicitly represented in the input, so mixing features would not easily work\. However, the CDF augmentation is kept in our model, as a fixed part of the encoder\.

### E\.2Observation module

The observation module processesh0fith^\{\\mathrm\{fit\}\}\_\{0\}\. The embeddings are passed throughLLtabular transformer layers with self attention\. Each transformer block consists of a row\-wise attention, column\-wise attention and a feed forward layer\.

hrow\\displaystyle h\_\{\\mathrm\{row\}\}=MultiHeadAttention​\(h,d​i​m=1\)\\displaystyle=\\mathrm\{MultiHeadAttention\}\(h,dim=1\)\(57\)hcol\\displaystyle h\_\{\\mathrm\{col\}\}=MultiHeadAttention​\(h,d​i​m=0\)\\displaystyle=\\mathrm\{MultiHeadAttention\}\(h,dim=0\)\(58\)hmlp\\displaystyle h\_\{\\mathrm\{mlp\}\}=FFN​\(h\)\\displaystyle=\\mathrm\{FFN\}\(h\)\(59\)hlfit\\displaystyle h^\{\\mathrm\{fit\}\}\_\{l\}=Block​\(hl−1fit\)\\displaystyle=\\mathrm\{Block\}\(h^\{\\mathrm\{fit\}\}\_\{l\-1\}\)\(60\)whereBlock\\mathrm\{Block\}denotes a pre\-layer\-normalized transformer block with residual sublayers values usinghrowh\_\{\\mathrm\{row\}\},hcolh\_\{\\mathrm\{col\}\}, andhmlph\_\{\\mathrm\{mlp\}\}andhhis the post layer normalised residual input for each block\.

Lethl,ifith^\{\\mathrm\{fit\}\}\_\{l,i\}be the output of layerll, featureii\. The matrix decoder module uses the final embeddingshL,ifith^\{\\mathrm\{fit\}\}\_\{L,i\}to generate predicted graph structure matrices\. The embeddings are averaged over fit rows and projected using two linear layers with RMS normalisation to map each output to embeddingsui,vi∈ℛhu\_\{i\},v\_\{i\}\\in\\mathcal\{R\}^\{h\}\. The logit for each edge is

M^i,j\\displaystyle\\hat\{M\}\_\{i,j\}=α⋅gelu​\(uiT⋅W\)⋅vj\+β,M^∈\{A^,R^,C^\}\\displaystyle=\\alpha\\cdot\\text\{gelu\}\(u\_\{i\}^\{T\}\\cdot W\)\\cdot v\_\{j\}\+\\beta,\\qquad\\hat\{M\}\\in\\\{\\hat\{A\},\\hat\{R\},\\hat\{C\}\\\}\(61\)with learned parametersW∈ℛd×dW\\in\\mathcal\{R\}^\{d\\times d\}and scalarα\\alphaandβ\\beta\. Separate decoder parameters are used for each output matrix,A^,R^,C^\\hat\{A\},\\hat\{R\},\\hat\{C\}\. Entries inM^∈ℛn×n\\hat\{M\}\\in\\mathcal\{R\}^\{n\\times n\}are logits representing the probability of each edge occurring\. Empirically, we found this nonlinear decoder worked better than a direct dot product used in AVICI\(Lorchet al\.,[2022](https://arxiv.org/html/2606.26467#bib.bib14)\)\. The matrix decoder’s predicted graph depends only on𝒟fit\\mathcal\{D\}^\{\\mathrm\{fit\}\}, and not on𝒟pred\\mathcal\{D\}^\{\\mathrm\{pred\}\}or the input graph structure\.

### E\.3Graph module

The graph module takes the graph embeddings and the intermediate hidden states from the observation module to generate graph hidden states\. These hidden states use feature observation hidden states for both positional embeddings and to pass information\. The intermediate embeddingshlcolh\_\{l\}^\{\\mathrm\{col\}\}are averaged over the column dimension to create feature embeddings,h¯lcol∈ℛ\(k\+2\)×h\\bar\{h\}\_\{l\}^\{\\mathrm\{col\}\}\\in\\mathcal\{R\}^\{\(k\+2\)\\times h\}\. These embeddings are then transformed into source and destination embeddings using two separate linear layers\. The resulting representations are combined via an outer product to produce the feature matrixHlfeat∈ℛ\(k\+2\)×\(k\+2\)×hH^\{\\mathrm\{feat\}\}\_\{l\}\\in\\mathcal\{R\}^\{\(k\+2\)\\times\(k\+2\)\\times h\}\. The graph hidden states is computed as

HlG=LN​\(hG\+Hlfeat\)\.H^\{G\}\_\{l\}=\\mathrm\{LN\}\(h^\{G\}\+H^\{\\mathrm\{feat\}\}\_\{l\}\)\.\(62\)Note that this module only depends on𝒟fit\\mathcal\{D\}^\{\\mathrm\{fit\}\}andGG, not on𝐱∗\\mathbf\{x\}^\{\*\}\. Also,hGh^\{G\}is not updated, we always use the initialhGh^\{G\}from the encoder module\. Empirical testing found no improvements from updating these weights between each layer\.

### E\.4Prediction module

The prediction module generates the distributionp^θ​\(y\)\\hat\{p\}\_\{\\theta\}\(y\), using the previous embeddings and encoded prediction inputs\. This module follows the same architecture and parameterization as the observation module, with the addition of a cross\-attention layer over the graph module embeddings:

hrow∗\\displaystyle h^\{\*\}\_\{\\mathrm\{row\}\}=MultiHeadAttention​\(h∗,dim=1\),\\displaystyle=\\mathrm\{MultiHeadAttention\}\(h^\{\*\},\\mathrm\{dim\}=1\),\(63\)hcol∗\\displaystyle h^\{\*\}\_\{\\mathrm\{col\}\}=MultiHeadAttention​\(KV=hrow,Q=h∗,dim=0\),\\displaystyle=\\mathrm\{MultiHeadAttention\}\(\\mathrm\{KV\}=h\_\{\\mathrm\{row\}\},\\mathrm\{Q\}=h^\{\*\},\\mathrm\{dim\}=0\),\(64\)hG∗\\displaystyle h^\{\*\}\_\{G\}=MultiHeadAttention​\(KV=HG,Q=h∗,dim=0\),\\displaystyle=\\mathrm\{MultiHeadAttention\}\(\\mathrm\{KV\}=H^\{G\},\\mathrm\{Q\}=h^\{\*\},\\mathrm\{dim\}=0\),\(65\)hmlp\\displaystyle h\_\{\\mathrm\{mlp\}\}=FFN​\(hG∗\),\\displaystyle=\\mathrm\{FFN\}\(h^\{\*\}\_\{G\}\),\(66\)hl∗\\displaystyle h^\{\*\}\_\{l\}=Block​\(hl−1∗\)\.\\displaystyle=\\mathrm\{Block\}\(h^\{\*\}\_\{l\-1\}\)\.\(67\)Block\\mathrm\{Block\}is a pre\-layer norm residual block usinghrow∗,hcol∗,hmlph^\{\*\}\_\{\\mathrm\{row\}\},h^\{\*\}\_\{\\mathrm\{col\}\},h\_\{\\mathrm\{mlp\}\}residual sublayers, with residual streamh∗h^\{\*\}\. Cross attention to the observation embeddingshrowh\_\{\\mathrm\{row\}\}andHGH^\{G\}allows the prediction module to access the observational data and graph structure\. Finally, the predicted outcome is generated from the final embedding corresponding to the outcome feature,hL,y∗∈ℛhh^\{\*\}\_\{L,y\}\\in\\mathcal\{R\}^\{h\}using a MLP head, mapping to prediction bucket logits,

z\\displaystyle z=MLP​\(hL,y∗\)∈ℝb\\displaystyle=\\mathrm\{MLP\}\(h^\{\*\}\_\{L,y\}\)\\in\\mathbb\{R\}^\{b\}\(68\)p^θ​\(y\)\\displaystyle\\hat\{p\}\_\{\\theta\}\(y\)=softmax​\(z\)∈ℝb\\displaystyle=\\mathrm\{softmax\}\(z\)\\in\\mathbb\{R\}^\{b\}\(69\)wherezzare output logits,bbis the number of buckets in the discretized output distribution andp^θ​\(y\)\\hat\{p\}\_\{\\theta\}\(y\)is the predicted probability distribution over these buckets\.

### E\.5Architecture changes

Several additional architecture optimisations and training improvements were made compared to the transformer backbone used in TabPFNv2\. These improvements are aimed at improving training stability and efficiency, rather than adding specific model features\. Changes are largely inspired by the ”Modded NanoGPT” project\(Jordanet al\.,[2024a](https://arxiv.org/html/2606.26467#bib.bib6)\)\. The main changes found to work are:

- •Full attention tohrowh\_\{\\mathrm\{row\}\}in the prediction module\.
- •Cosine learning rate schedule with 1000 linear warmup steps\.
- •relu2\\text\{relu\}^\{2\}activation instead of relu/gelu\(Zhanget al\.,[2024](https://arxiv.org/html/2606.26467#bib.bib4)\)\.
- •Reduce positional embedding initialisation weight scale\.
- •Pre layer norm instead of post layer norm\(Xionget al\.,[2020](https://arxiv.org/html/2606.26467#bib.bib3)\)\.
- •Apply layer norm before prediction decoder MLP head\.
- •Muon optimizer for matrices instead of AdamW\(Jordanet al\.,[2024b](https://arxiv.org/html/2606.26467#bib.bib7)\), using the Kimi variant\(Liuet al\.,[2025](https://arxiv.org/html/2606.26467#bib.bib8)\)\.
- •QK norm with learnable temp for attention heads\(Henryet al\.,[2020](https://arxiv.org/html/2606.26467#bib.bib5)\)\.
- •Increased width from 192 to 288\. Total parameters increase from 10\.28M to 22\.97M\.

We test the effect of these changes in the ablation section[F](https://arxiv.org/html/2606.26467#A6), and find they significantly improve training stability and efficiency\.

## Appendix FArchitecture Ablations

This section conducts an ablation of the architecture changes described in Section[E\.5](https://arxiv.org/html/2606.26467#A5.SS5)\. The TabPFN\-CFM architecture consists of a modified Do\-PFN architecture to handle the additional structural and outcome prediction objectives\. Independently, we introduce multiple updates to the transformer backbone and training procedure to improve training efficiency and final performance\. No new functionality is added from these changes\.

To evaluate the impact of each of the changes, we conduct an ablation study on on smaller datasets with models trained for 30k steps\. The training and validation datasets are identical on all setups to isolate the effect of architecture changes\. The learning rate is 3e\-4 for Muon and 1e\-4 for AdamW, as AdamW required a lower learning rate for stability\. Results are shown in Figure[3](https://arxiv.org/html/2606.26467#A6.F3)\. We sequentially apply each of the architectural changes on top of previous changes and record the prediction and adjacency losses\. Both prediction and adjacency losses decrease with the changes\. Note changes interact with each other in nonlinear ways and some changes only show in longer training runs, so the impact of each change in the final model is more complex than shown in the plots here\.

![Refer to caption](https://arxiv.org/html/2606.26467v1/x4.png)Figure 3:Ablation results with 30k training steps for prediction loss and adjacency matrix loss\. Changes are applied sequentially from left to right\.Next, we compare our model with longer training runs with and without the changes, but with the base version scaled up to the same 23\.2M parameter count for fairness\. Models are trained for a much longer 150k steps\. Figure[4](https://arxiv.org/html/2606.26467#A6.F4)shows validation loss curves for both models\. Our optimised model reaches the same prediction loss 4 times faster \(37K vs 150K steps to reach \-0\.195 loss\) and 3 times faster adjacency loss \(50K vs 150K steps to reach 0\.283 loss\) compared to the unoptimised version, with significantly lower final loss\. The base model trained in 8\.52 hours while the final model trained in 10\.24 hours, a 20% increase in training time largely from Muon and the additional norm layers\. Despite a small increase in per step time, the final model is significantly more efficient\.

![Refer to caption](https://arxiv.org/html/2606.26467v1/x5.png)

![Refer to caption](https://arxiv.org/html/2606.26467v1/x6.png)

Figure 4:Loss curves comparing our and baseline architecture for 150k training steps, prediction loss \(left\), adjacency matrix loss \(right\)\.
## Appendix GSynthetic toy examples

### G\.1Linear Instrumental Variable example

Showcase experiment on the Instrumental Variable \(IV\) problem, with SEM:Z→T→Y,U→T,U→YZ\\rightarrow T\\rightarrow Y,U\\rightarrow T,U\\rightarrow Y\. To demonstrate, we examine predictions from a simple random synthetic linear SEM\. The SEM equations \(after normalizing with the observational sample statistics\) are:

ϵ\{u,z,t,y\}\\displaystyle\\epsilon\_\{\\\{u,z,t,y\\\}\}∼U​\(−0\.4,0\.4\)\\displaystyle\\sim U\(\-0\.4,0\.4\)U\\displaystyle U=4\.303​ϵy−0\.031\\displaystyle=4\.303\\epsilon\_\{y\}\-0\.031Z\\displaystyle Z=4\.380​ϵz−0\.007\\displaystyle=4\.380\\epsilon\_\{z\}\-0\.007T\\displaystyle T=binarize0\.5​\(−0\.105​z\+0\.104​u\+ϵt\)\\displaystyle=\\text\{binarize\}\_\{0\.5\}\(\-0\.105z\+0\.104u\+\\epsilon\_\{t\}\)Y\\displaystyle Y=−1\.818∗T−0\.108∗U\+1\.047∗ϵy\+0\.906\\displaystyle=\-1\.818\*T\-0\.108\*U\+1\.047\*\\epsilon\_\{y\}\+0\.906where the binarize function returns 1 if the input is in the top 50 percent of the overall distribution and 0 otherwise\. We draw a single sample from this SEM:

z\\displaystyle z=−0\.187\\displaystyle=\-0\.187u\\displaystyle u=0\.421\\displaystyle=0\.421t\\displaystyle t=1\\displaystyle=1y\\displaystyle y=−0\.568\\displaystyle=\-0\.568Now, we identify the exact distribution assuming everything is observed, for observational, interventional, and counterfactual queries\. The observational distribution ofyyis

Y\\displaystyle Y=−0\.957\+1\.047∗ϵy,\\displaystyle=\-0\.957\+1\.047\*\\epsilon\_\{y\},the interventional distribution withd​o​\(T=0\)do\(T=0\)is

Y\\displaystyle Y=0\.861\+1\.047∗ϵy\\displaystyle=0\.861\+1\.047\*\\epsilon\_\{y\}and the counterfactual distribution, which assumesyyis observed, is

Y\\displaystyle Y=1\.250\\displaystyle=1\.250sinceϵy=0\.372\\epsilon\_\{y\}=0\.372can be determined by the observation ofyy, which allows for the exact counterfactual value to be computed with no uncertainty\.

![Refer to caption](https://arxiv.org/html/2606.26467v1/x7.png)

![Refer to caption](https://arxiv.org/html/2606.26467v1/x8.png)

![Refer to caption](https://arxiv.org/html/2606.26467v1/x9.png)

Figure 5:Interventional \(left\), Observational \(center\), and Counterfactual \(right\) distributions\. Exact solutions are in blue, model predictions in orange for whenUUis observed, and in green whenUUis unobserved\. Model is given the graph structure\.
### G\.2Nonlinear SEM

A more complex nonlinear SEM is used to test the model\. The graph structure is fixed asZ→T→V→Y,U→V,U→W→YZ\\rightarrow T\\rightarrow V\\rightarrow Y,U\\rightarrow V,U\\rightarrow W\\rightarrow YwhereUUis unobserved\. The graph contains latent variables and covariates that are descendants ofTT\. The SEM equations are drawn from our nonlinear prior distribution\. Since the exact posterior for interventional and observational distributions is intractable, we compare the model predictions to a single sample drawn from the true distribution\. Results are shown forn=512n=512\. Figure[6](https://arxiv.org/html/2606.26467#A7.F6)shows the model predictions align well with the samples, even without the graph structure\. Giving the graph structure improves predictions\. Table[11](https://arxiv.org/html/2606.26467#A7.T11)shows the true and predicted adjacency matrices\. The predicted adjacency matrix generally matches the true adjacency matrix, though there is some uncertainty in the children ofTTlikely due toV,W,YV,W,Yall being correlated\. In Figure[7](https://arxiv.org/html/2606.26467#A7.F7), we show the predicted distribution shifts closer towards the true counterfactual asnnincreases \(without𝒢\\mathcal\{G\}\)\. For smallnn, the counterfactual is biased towards the observedy∗y^\{\*\}, which vanishes asnnincreases and the model is able to learn the underlying SEM better\.

![Refer to caption](https://arxiv.org/html/2606.26467v1/x10.png)

![Refer to caption](https://arxiv.org/html/2606.26467v1/x11.png)

![Refer to caption](https://arxiv.org/html/2606.26467v1/x12.png)

Figure 6:Interventional \(left\), Observational \(center\), and Counterfactual \(right\) distributions\. Predictions are compared with and without the true graph structure𝒢\\mathcal\{G\}\. The Observational and Interventional predictions are compared to a single sample drawn from the true distribution, while the Counterfactual predictions are compared to the true counterfactual value\.![Refer to caption](https://arxiv.org/html/2606.26467v1/x13.png)Figure 7:Counterfactual distribution predictions for the nonlinear SEM with different𝒟fit\\mathcal\{D\}^\{\\mathrm\{fit\}\}sample sizes\.Table 11:True and predicted adjacency matrices for the nonlinear SEM experiment\. Left: True adjacency matrix\. Right: Predicted adjacency matrix \(values represent predicted edge probabilities\)\.TZVWYT00100Z10000V00001W00001Y00000
TZVWYT0\.0000\.0020\.9920\.9380\.552Z0\.9820\.0000\.3180\.5220\.391V0\.0000\.0000\.0000\.5970\.758W0\.0060\.0050\.5730\.0000\.871Y0\.0000\.0000\.0010\.0050\.000

## Appendix HOOD data generation

We test our model on datasets drawn using random Fourier functions as nonlinearities\. This has been used as part of the prior in other works, but our models are not trained on these functions, so they are OOD\. The distribution of SEM functions is constructed as follows\. First, parameters are sampled,

w∼U​\[0\.1,0\.5\]\\displaystyle w\\sim U\[0\.1,0\.5\]\(70\)c∼U​\[1\.41,5\.65\]\\displaystyle c\\sim U\[1\.41,5\.65\]\(71\)b∼U​\[−0\.5,0\.5\]\\displaystyle b\\sim U\[\-0\.5,0\.5\]\(72\)Wi,j∼Cauchy​\(0,w\)\\displaystyle W\_\{i,j\}\\sim\\text\{Cauchy\}\(0,w\)\(73\)ϕi∼N​\(0,2​π\)\\displaystyle\\phi\_\{i\}\\sim N\(0,2\\pi\)\(74\)ki∼N​\(0,c​\(nin\)−0\.5\)\.\\displaystyle k\_\{i\}\\sim N\(0,c\(n\_\{\\mathrm\{in\}\}\)^\{\-0\.5\}\)\.\(75\)The output of each node is

y=b\+∑i=132ki​cos⁡\(\(x​W⊤\)i\+ϕi\)\\displaystyle y=b\+\\sum\_\{i=1\}^\{32\}k\_\{i\}\\cos\\\!\\bigl\(\(xW^\{\\top\}\)\_\{i\}\+\\phi\_\{i\}\\bigr\)\(76\)wherexxis a concatenation of the node’s parent values and noise variables, and parametersW,k,ϕW,k,\\phiare randomly sampled frequency, magnitude and phase for each of the 32 coefficients\. This process generates random functions that exhibit non\-trivial nonlinear behaviors over a range of frequencies\.

Table 12:Predictions on in distribution SEMs withn=1024n=1024fit samples, compared by mean squared error \(↓\\downarrow\)\. The baseline methods are unable to make counterfactual predictions\.Table 13:Structural predictions on in distribution SEMs withn=1024n=1024fit samples, evaluated on predicted adjacency matrix, ancestral matrix and confounding matrix\. Main values tracks the AUROC \(↑\\uparrow\), with accuracy \(↑\\uparrow\)in brackets\.Table 14:Predictions on in distribution SEMs with varying fit samples \(second header\), compared by mean squared error \(↓\\downarrow\)\. The baseline methods are unable to make counterfactual predictions\.Table 15:Structural predictions on in distribution SEMs with varying fit samples\. Main values track AUROC \(↑\\uparrow\), with accuracy \(↑\\uparrow\) in brackets\.
Table 16:Predictions on out of distribution SEMs with varying fit samples \(second header\), compared by mean squared error \(↓\\downarrow\)\. The baseline methods are unable to make counterfactual predictions\.Table 17:Structural predictions on synthetic OOD SEMs with varying fit samples\. Main values track AUROC \(↑\\uparrow\), with accuracy \(↑\\uparrow\) in brackets\.