Causal Foundation Models with Continuous Treatments

2026-05-14

AI Summary

This work addresses the challenging problem of causal effect estimation with continuous treatment variables, moving beyond the limitations of conventional binary treatment frameworks. It introduces the first foundation model tailored for continuous treatments, integrating Bayesian data-generating priors, a Transformer architecture, and in-context learning to enable zero-shot generalization to unseen intervention scenarios through meta-learning. The model achieves end-to-end reconstruction of individualized dose–response curves without requiring task-specific fine-tuning. Experimental results demonstrate that the proposed approach substantially outperforms existing specialized causal models in reconstruction accuracy, establishing state-of-the-art performance on this task.

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.

Problem

• causal inference
• continuous treatments
• treatment-response curve
• observational data
• individual treatment effect

Innovation

• causal foundation model
• continuous treatments
• in-context learning
• treatment-response curve
• amortized inference