← Back to Machine Learning
cs.LG

One model estimates causal effects across any continuous treatment

Christopher Stith, Medha Barath, Vahid Balazadeh, Jesse C. Cresswell, Rahul G. Krishnan

May 14, 2026

Estimating how outcomes change across a continuous range of interventions — drug dosages, policy intensities, pricing — is harder than binary treatment comparisons and has seen far less methodological development. This model uses in-context learning: a transformer is pre-trained on a rich synthetic corpus of causal data-generating processes, then at inference time reconstructs individual treatment-response curves directly from observational data, with no fine-tuning required. By amortizing Bayesian posterior inference into the forward pass, the approach sidesteps the expensive per-task optimization that competing methods require. On individual treatment-response benchmarks, it matches or exceeds models trained specifically for each task.
Published as Causal Foundation Models with Continuous Treatments arXiv:2605.15133
Read the original paper →