cs.LGMay 14, 2026

Causal Foundation Models with Continuous Treatments

Authors: Christopher StithMedha BarathVahid BalazadehJesse C. CresswellRahul G. Krishnan

Organizations: Layer 6 AI · University of Toronto · Vector Institute

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. Inference code (including trained model weights) can be found at https://github.com/layer6ai-labs/CCPFN-inference .

Explore similar work

CardsList
  1. Causal Foundation Models

    Sep 2, 2026Christopher Stith, Hossein Rahmani, Jesse C. CresswellCausal InferencesCausal Foundation Models