Use What You Know: Causal Foundation Models with Partial Graphs
cs.LG
Submitted: 2026-02-16
Updated: 2026-09-27
Journal ref: ICML 2026
Code: https://github.com/ArikReuter/Graphs4CausalFoundationModels
License: http://creativecommons.org/licenses/by/4.0/
The gist: Estimating causal quantities traditionally relies on bespoke estimators tailored to specific assumptions.
Terminology
Abstract
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.
Sources
- CausalPFN: Amortized Causal Effect Estimation via In-Context Learning
- Convolutional Networks on Graphs for Learning Molecular Fingerprints
- TabArena: A Living Benchmark for Machine Learning on Tabular Data
- Neural Processes
- TabPFN-2.5: Advancing the State of the Art in Tabular Foundation Models
- Gaussian Error Linear Units (GELUs)
- TabPFN: A Transformer That Solves Small Tabular Classification Problems in a Second
- Adam: A Method for Stochastic Optimization
- Foundation Models for Causal Inference via Prior-Data Fitted Networks
- A detailed treatment of Doob's theorem
- Transformers Can Do Bayesian Inference
- RealCause: Realistic Causal Inference Benchmarking
- Scalable Diffusion Models with Transformers
- TabICLv2: A better, faster, scalable, and open tabular foundation model
- Do-PFN: In-Context Learning for Causal Effect Estimation
- ACTIVA: Amortized Causal Effect Estimation via Transformer-based Variational Autoencoder
- Estimating individual treatment effect: generalization bounds and algorithms
- GLU Variants Improve Transformer
- Qwen2-VL: Enhancing Vision-Language Model's Perception of the World at Any Resolution
- LimiX: Unleashing Structured-Data Modeling Capability for Generalist Intelligence
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks