Causal Foundation Models with Continuous Treatments

Estimating causal effects from observational data is a fundamental challenge in many disciplines. In this paper, we present the first causal foundation model for the continuous treatment setting. This setting is even more challenging as models need to represent effects across a continuum of treatment values. Our model meta-learns the ability to predict causal effects across a wide variety of unseen tasks without additional training or fine-tuning. We design a novel prior over data-generating processes with continuous treatments, 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 compared to causal models which are trained for each task.