← 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.
Read the original paper →