Causal Foundation Models with Continuous Treatments

TL;DR

This paper introduces the first causal foundation model for continuous treatments, capable of predicting causal effects across various tasks without additional training, using a transformer trained on a rich causal corpus.

Contribution

It presents a novel prior for data generation and a transformer-based model that reconstructs treatment-response curves from observational data, enabling zero-shot causal inference.

Findings

Achieves state-of-the-art performance in treatment-response curve reconstruction.

Can generalize to unseen tasks without additional training.

Uses in-context learning to efficiently infer causal effects.

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.