Abstract:Causal inference, estimating causal effects from observational data, is a fundamental tool in many disciplines. Of particular importance across a variety of domains is the continuous treatment setting, where the variable of intervention has a continuous range. This setting is far less explored and represents a substantial shift from the binary treatment setting, with models needing to represent effects across a continuum of treatment values. In this paper, we present the first causal foundation model for the continuous treatment setting. Our model meta-learns the ability to predict causal effects across a wide variety of unseen tasks without additional training or fine-tuning. First, we design a novel prior over data-generating processes with continuous treatment variables in order to generate a rich causal training corpus. We then train a transformer to reconstruct individual treatment-response curves given only observational data, leveraging in-context learning to amortize expensive Bayesian posterior inference. Our model achieves state-of-the-art performance on individual treatment-response curve reconstruction tasks compared to causal models which are trained specifically for those tasks.
| Comments: | 22 pages, 9 figures |
| Subjects: | Machine Learning (cs.LG) |
| Cite as: | arXiv:2605.15133 [cs.LG] |
| (or arXiv:2605.15133v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2605.15133 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Christopher Stith [view email]
[v1]
Thu, 14 May 2026 17:40:28 UTC (2,385 KB)