# Use What You Know: Causal Foundation Models with Partial Graphs

arXiv: https://arxiv.org/abs/2602.14972 · HTML: https://mlsanity.com/paper/2602.14972/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 34 of its 58 references that are in the map.
- First posted on arXiv 16 Feb 2026.

## Builds on (most cited first, 30 of 34)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolution](https://mlsanity.com/paper/2409.12191/) (2024)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Convolutional Networks on Graphs for Learning Molecular Fingerprints](https://mlsanity.com/paper/1509.09292/) (2015)
- [TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second](https://mlsanity.com/paper/2207.01848/) (2022)
- [Self-Normalizing Neural Networks](https://mlsanity.com/paper/1706.02515/) (2017)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Conditional Neural Processes](https://mlsanity.com/paper/1807.01613/) (2018)
- [Transformers Can Do Bayesian Inference](https://mlsanity.com/paper/2112.10510/) (2021)
- [Neural Processes](https://mlsanity.com/paper/1807.01622/) (2018)
- [TabICL: A Tabular Foundation Model for In-Context Learning on Large Data](https://mlsanity.com/paper/2502.05564/) (2025)
- [TabArena: A Living Benchmark for Machine Learning on Tabular Data](https://mlsanity.com/paper/2506.16791/) (2025)
- [TabPFN-2.5: Advancing the State of the Art in Tabular Foundation Models](https://mlsanity.com/paper/2511.08667/) (2025)
- [Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Modeling](https://mlsanity.com/paper/2207.04179/) (2022)
- [LimiX: Unleashing Structured-Data Modeling Capability for Generalist Intelligence](https://mlsanity.com/paper/2509.03505/) (2025)
- [Mitra: Mixed Synthetic Priors for Enhancing Tabular Foundation Models](https://mlsanity.com/paper/2510.21204/) (2025)
- [Statistical Foundations of Prior-Data Fitted Networks](https://mlsanity.com/paper/2305.11097/) (2023)
- [Do-PFN: In-Context Learning for Causal Effect Estimation](https://mlsanity.com/paper/2506.06039/) (2025)
- [CausalPFN: Amortized Causal Effect Estimation via In-Context Learning](https://mlsanity.com/paper/2506.07918/) (2025)
- [RealCause: Realistic Causal Inference Benchmarking](https://mlsanity.com/paper/2011.15007/) (2020)
- [Foundation Models for Causal Inference via Prior-Data Fitted Networks](https://mlsanity.com/paper/2506.10914/) (2025)
- [A Meta-Learning Approach to Bayesian Causal Discovery](https://mlsanity.com/paper/2412.16577/) (2024)

## Built on by (most cited first, 7 of 7)

- [SurvPFN: Towards Foundation Models for Survival Predictions](https://mlsanity.com/paper/2606.04564/) (2026)
- [FoundCause: Causal Discovery with Latent Confounders from Observational Data](https://mlsanity.com/paper/2606.17516/) (2026)
- [Foundation Models for Partial Causal Identification](https://mlsanity.com/paper/2608.20841/) (2026)
- [Causal Foundation Models](https://mlsanity.com/paper/2609.03003/) (2026)
- [CAFE: Counterfactual Prediction via Fast Posterior Estimation](https://mlsanity.com/paper/2609.32167/) (2026)
- [CIDER-FM: Foundation Models for Causal Inference from Diverse Experimental Regimes](https://mlsanity.com/paper/2609.39523/) (2026)
- [Market Impact Estimation via Causal Foundation Models](https://mlsanity.com/paper/2610.12537/) (2026)

## Abstract (arXiv)

Estimating causal quantities traditionally relies on bespoke estimators tailored to specific assumptions. Recently proposed Causal Foundation Models (CFMs) promise a more unified approach by amortising causal discovery and inference in a single step. However, in their current state, they do not allow for the incorporation of any domain knowledge, which can lead to suboptimal predictions. We bridge this gap by introducing methods to condition CFMs on causal information, such as the causal graph or more readily available ancestral information. When access to complete causal graph information is too strict a requirement, our approach also effectively leverages partial causal information. We systematically evaluate conditioning strategies and find that injecting learnable biases into the attention mechanism, together with a graph-convolutional encoder, is a highly effective method to utilise full and partial causal information. Our experiments show that this conditioning allows a general-purpose CFM to match the performance of specialised models trained on specific causal structures. Overall, our approach addresses a central hurdle on the path towards all-in-one causal foundation models: the capability to answer causal queries in a data-driven manner while effectively leveraging any amount of domain expertise.