Causal Foundation Models with Continuous Treatments

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Stith, Christopher, Barath, Medha, Balazadeh, Vahid, Cresswell, Jesse C., Krishnan, Rahul G.
Format: Preprint
Veröffentlicht: 2026
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866911685917278208
author Stith, Christopher
Barath, Medha
Balazadeh, Vahid
Cresswell, Jesse C.
Krishnan, Rahul G.
author_facet Stith, Christopher
Barath, Medha
Balazadeh, Vahid
Cresswell, Jesse C.
Krishnan, Rahul G.
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2605_15133
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Causal Foundation Models with Continuous Treatments
Stith, Christopher
Barath, Medha
Balazadeh, Vahid
Cresswell, Jesse C.
Krishnan, Rahul G.
Machine Learning
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.
title Causal Foundation Models with Continuous Treatments
topic Machine Learning
url https://arxiv.org/abs/2605.15133