Causal Foundation Models with Continuous Treatments
arXiv:2605.15133v1 [cs.LG] 14 May 2026
Christopher Stith∗ Layer 6 AI [email protected]
Medha Barath∗ University of Toronto [email protected]
Vahid Balazadeh University of Toronto Vector Institute [email protected] Jesse C. Cresswell Layer 6 AI [email protected]
Rahul G. Krishnan University of Toronto Vector Institute [email protected]
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.
Figure 1: Estimating causal effects for continuous treatments (right) is much more challenging than for binary treatments (left), as multiple treatment-response curves fit the observed data equally well. ∗ Equal Contribution
Preprint.
1
Introduction
Causal inference is a central task for decision-making across many domains, including precision medicine [1, 43], econometric policy-making [2, 9], and algorithmic marketing [6, 14]. Estimating the effect of an intervention from observational data alone is complicated, as the presence of confounders can bias naive estimators of potential outcomes. The causal inference community has built a rich library of estimators under the framework of ignorability, which assumes no unobserved confounding exists [34]. However, the end-to-end implementation of these estimators involves considerable time and effort: for any given task, a domain expert must inspect the data, propose an underlying mechanism to model it, choose an estimator that fits this mechanism, and only then train their model. In addition to this approach, there has been recent work at the intersection of causal inference and meta-learning [8]. The goal here is to train a model to perform a wide variety of causal inference tasks. A particularly promising framework has been Bayesian inference and in-context learning (ICL) [40, 3, 32]. Here, a model is trained over a diverse set of causal data-generating processes (DGPs) drawn from a prior π on possible DGPs, learning how to approximate the posterior-predictive distribution (PPD) for any new dataset. At inference, observational data for a given unseen task is passed to the model as context, from which the model learns the PPD for the causal estimand of interest, amortizing the cost of posterior inference. This pipeline turns the typically expensive and manual approach to causal inference into a completely data-driven process. Most previous work in Bayesian causal inference focuses on the binary treatment setting, where units can be split into control and treatment groups. However, numerous applications deal with the continuous treatment setting, in which interventional variables have a continuous range. In marketing and economics, one is interested in the sensitivity of economic indicators to prices or rates; in pharmacology, dose-response curves are frequently used to characterize the influence of medication; and in medical contexts, one is interested not just in whether or not to treat a patient, but how much of a medication to give. Transitioning from binary to continuous treatments introduces significant theoretical and practical hurdles. Whereas in the binary case the goal is to estimate a single number per unit or per dataset (such as the Average Treatment Effect), continuous settings often require the estimation of a full treatmentresponse curve. This is a non-trivial task, as illustrated in Figure 1. Many curves are consistent with the observed data, and without additional structural assumptions the treatment-response function is not identified from observational data alone. Designing a prior that captures a suitably diverse set of continuous-valued DGPs is also highly non-trivial. The model must not only learn how to represent scalar effects, but also a wide spectrum of functions while remaining robust to the issues arising from the higher dimensionality of continuous treatment spaces. Finally, within a transformer-based architecture, representing continuous-valued treatments requires careful consideration. Standard tokenization or simple linear projections may fail to capture the high-frequency variations or local smoothness required for accurate causal discovery. We introduce CCPFN (Continuous Causal Prior-Fitted Network), a causal foundation model that directly learns to reconstruct continuous individual treatment-response curves via in-context learning, and that can be applied off-the-shelf to any causal inference problem without fine-tuning. Central to this work is the construction of a novel prior over DGPs which interweaves three neural networks to generate covariates features, treatments, and outcomes. This prior directly encodes the assumption of identifiability via strong ignorability. To summarize, our key contributions are: • The first causal foundation model for continuous treatments, achieving top performance on treatment-response curve reconstruction tasks. • A novel prior over data-generating processes which generates a rich set of synthetic causal data with continuous treatments. • A (semi-)synthetic causal scenario generation method to facilitate model validation in causal inference with continuous treatments.
2
Background
We work in the potential outcomes framework of Neyman-Rubin [41] in the continuous treatment regime. We let capital letters denote random variables, while lower case letters denote their real2
izations. Let T ⊆ R be an interval; this represents the continuous treatment space. We let T ∈ T denote the observed treatment and X ∈ X the observed covariates. For each t ∈ T , we denote the potential outcome under treatment t by Y t . We let Y ∈ Y denote the observed (factual) outcome, so that Y = Y T . Finally, we let P = P (X, T, {Y t }t∈T , Y ) denote the data-generating process (DGP), which is the interventional distribution. Its marginal Pobs = P (X, T, Y ) is the observed distribution. Suppose we have an i.i.d. sample of units indexed by n = 1, 2, . . . , N . For each unit n we observe the covariate vector xn , the applied treatment tn , and the observed outcome yn . We are primarily concerned with recovering the individual treatment-response curve (ITRC) which represents the causal effect of applying any of the continuous treatment levels on a unit with covariates x. We define this curve using conditional expected potential outcomes (CEPOs). For a given covariate vector x and a treatment level t, the CEPO is defined as µt (x) := E[Y t | X = x],
∀t ∈ T .
(1)
The ITRC for an individual with covariates x is simply the function T → Y traced out by the CEPO across all treatment levels: t 7→ µt (x), ∀t ∈ T . (2) The CEPO, and hence the ITRC, is identifiable when it can be written as a function of the observational distribution Pobs [34]. We make the following assumptions to ensure this is the case: 1. Unconfoundedness. Conditional on the covariates X, the potential outcomes Y t are independent of the treatment assignment T , Y t ⊥ T | X ∀t ∈ T . 2. Positivity/Overlap. There exists a constant c > 0 such that for all x ∈ X and all t ∈ T , pT |X (t | x) ≥ c. Under Assumptions 1–2, the CEPO (and hence the ITRC) is identifiable from the observational distribution [19, 34]: µt (x) = E[Y t | X = x] = E[Y | X = x, T = t]. Bayesian Causal Inference. A Bayesian formulation of causal inference places a prior π(ψ) over DGPs ψ, each indexing a joint distribution P ψ (X, T, {Y t }t∈T , Y ). Given i.i.d. observations ψ Dobs = {(x(n) , t(n) , y (n) )}N n=1 sampled from the observational distribution Pobs , Bayes’ rule yields the posterior distribution π(ψ | Dobs ). The posterior-predictive distribution (PPD) of any causal estimand g(ψ) is then defined as Z π g (· | Dobs ) := B 7→ I(g(ψ) ∈ B) π(ψ | Dobs ) dψ , B ∈ B, (3) where B is the Borel σ-algebra over R. This paradigm provides a unified framework for point estimation [41, 37, 29]. In practice, computing the posterior is intractable, motivating the amortized approach to estimating the PPD described below. Prior-Fitted Networks. Recent works [21, 31, 38, 3, 40, 32] on PFNs amortize posterior-predictive inference using a single transformer-based [47] neural network qθ , trained to directly approximate the PPD via the data-prior loss ℓ(θ) := Eφ∼π, D∪{x,y}∼P φ − log qθ (y | x, D) . (4) Training requires only prior samples, while no posterior sampling is needed at inference time. As in [3] our target is the PPD of the CEPO (i.e. g(ψ) = µt (x; ψ) in Equation (3)) where our CEPO is parameterised by the DGP index ψ. Balazadeh et al. [3] proved that a.s.
Eµ∼πµt (·|x,Dobs ) [µ] −−→ µt (x ; ψ ⋆ ),
∀t ∈ T , and almost all x ∈ X ,
provided the prior π is CEPO-identifiable. Following their framework, we train a transformer qθ to approximate the full predictive distribution π µt using a Causal Data-Prior Loss. For any t ∈ T , this loss is defined as Lt (θ) := Eψ∼π, Dobs ∪{x}∼P ψ [− log qθ (µt (x ; ψ) | x, t, Dobs )] . obs
3
(5)
During training, we randomly sample t ∼ Uniform(T ) in order to cover the full treatment range T . Moreover, in practice, we use regression-as-classification and quantize the outcome variable into L = 1024 bins from [−10, 10] (after z-standardization). The model then outputs qθ as a discrete histogram distribution on each of these L bins. To approximate the true CEPO-PPD in a tractable manner we assume it is given by N (µt (x), σ 2 ) for σ ≪ 1. The histogram loss which we train with is then defined as LHL (θ | x, t) = −
L X
N (µt (x), σ 2 )[l] log qθ (l | x, t, Dobs ),
(6)
l=1
where N (µt (x), σ 2 )[l] denotes the probability of a N (µt (x), σ 2 )-random variable lying in bin l, and qθ (l | x, t, Dobs ) denotes the model’s output probability for bin l. Related Work. While ours is the first model to leverage in-context learning for the continuous treatment regime in causal inference, numerous techniques have been developed over the past decade to address this regime due to its importance across a variety of domains. DRNet [42] addresses the challenge of multi-armed treatments, each with associated treatment parameters. For each treatment, DRNet breaks the treatment range into E ∈ N equally-sized subintervals and assigns a head to each. This ensures the effect of the treatment parameter persists throughout the network layers, but it can lead to discontinuities in the learned treatment-response curves. SCIGAN [5] leverages generative adversarial networks [13] to generate and learn counterfactual representations. VCNet [36] further addresses the continuous treatment space by training a varying coefficient network, where neural network parameters θ = θ(t) depend on the treatment value t. ADMIT [48] extends the theoretical framework by bounding the counterfactual loss of estimating treatment-response curves and proposes an algorithm that makes counterfactual estimations. Meanwhile, several groups have approached the binary treatment setting for causal inference via causal foundation modeling; see [40, 3, 32]. Training on synthetic data is a common practice in the tabular foundation model (TFM) space. In particular, TabPFN [21, 16] and TabICL [38, 39] both train on purely synthetic data generated from a carefully-designed prior. While not designed specifically for causal inference, their priors (like ours) are generated via synthetic structural causal models (SCMs) using randomly-generated neural networks. More recently, it has been shown that data augmentation during fine-tuning with synthetic, SCM-based causal data can improve TFM performance [7].
3
Method
Generating a Prior. Following the general approach of CausalPFN [3], we design a novel prior over potential DGPs arising from continuous treatment scenarios. As we are interested in the setting where the causal estimands are identifiable, we design these priors so that the assumptions of unconfoundedness and positivity hold (see Section 2). The other key requirement of a training prior is to generate a suitably diverse set of DGPs so as to represent a variety of real-world scenarios. This is complicated considerably when dealing with continuous treatments T , and so the design of our prior differs considerably from [3]. (a) Generate covariates
(b) Generate treatment
(c) Generate outcome
X1
X2
ηY
X1 X3
Te
X3
µT
X2 X2
ηT
T
Figure 2: A schematic of our 3-MLP prior. In practice all MLPs drop edges with a certain probability. Inspired by the construction of structural causal models (SCMs) via randomized MLPs in [39], our prior consists of three separate randomized MLPs (see Figure 2): • MLPX generates the covariates X, which can be divided into three disjoint subsets: the set XT which are direct causes of T only; the set XY which are direct causes of Y only; and the set of confounders Xconf which are direct causes of both T and Y . 4
• MLPT generates the observed treatment T from XT ∪ Xconf . • MLPY generates outcomes (both factual and counterfactual), given treatment t ∈ T and covariates XY ∪ Xconf . Each SCM produces samples
N xn , tn , yn , t′n , µt′n (xn ) n=1
where (tn , yn ) are the factual (observed) treatment and outcome, respectively, t′n is a counterfactual treatment, and µt′n (xn ) is the corresponding CEPO. We let N denote the number of samples throughout. We presently discuss each component MLP; further details can be found in Appendix B. The following hyperparameters are randomly sampled: • The number of layers LX , LT , and LY of each MLP; • The number of hidden units HX , HT , and HY of each MLP; • The densities dX , dT , and dY are uniformly sampled in [0.1, 1]; edges in each MLP(·) are dropped with probability 1 − d(·) ; • The degree of confounding ρ ∼ U(0, 1), which is the proportion of covariates X that are direct causes of both T and Y . The remaining covariates are randomly split into XT and XY . MLPX . In addition to LX , HX , and dX , a noise scale s > 0 and a natural number K > 0 (representing the number of covariates) are chosen. For each l ∈ [LX ], noise values ϵ(l) ∈ RN × RHX for each n ∈ [N ] are generated by randomly-chosen distributions in {N (0, 1), Laplace(0, 1), t3 } followed by a random shift and scaling. The input ϵ(0) is propagated through the MLP by randomly initialized weights: z (l) = σ (l) (W (l) z (l−1) ) + ϵ(l) , z (0) = ϵ(0) , (7) where each σ (l) represents a vector of random activation functions. The MLP weights come from 12 2 2 the normal distribution N (0, σw ), where σw = max(H∗p ; H is the hidden width of the MLP d ,1.0) while pd represents the probability of a given weight being kept by a random sparsity mask. We then choose K nodes among the z (l) to be covariates; this produces covariates differing in complexity and marginal distribution. In addition, to simulate realistic tabular data, we apply random “tabular corruption”. A random subset of covariate nodes is either binarized, quantized, or zero-inflated (across the N samples). Contrary to past works [21, 39], which only apply this transformation after the forward pass, we transform roughly 35% of covariate nodes during the forward pass. This allows the SCM to see and utilize the realistic tabular features during generation, adding realistic tabular diversity to the DGP itself, rather than solely post-hoc. Another roughly 65% of the remaining covariates are further corrupted after the forward pass is completed. The superiority of this method as opposed to solely post-hoc tabular corruption, is demonstrated in our ablations (see Table 4). This does not affect the theoretical validity of the SCM, as all tabular corruption occurs within MLPX alone; it merely enhances the diversity and realism of the prior. MLPT . The input to MLPT is XT ∪ Xconf , and the output is a scalar T̃ which is the expected treatment value given XT ∪ Xconf . A random node ηT in the hidden layers is chosen to act as a heteroscedastic noise scale parameter. To produce the final outcome T we compute T = T̃ + σ(T̃ ) · ηT · ϵ,
ϵ ∼ N (0, 1),
(8)
where σ(T̃ ) is the empirical standard deviation of the T̃ over the N samples. Rather than designating the output node as T itself, this method practically ensures that positivity p(T | X = x) > 0 holds, as T is centered at T̃ with strictly positive variance. Finally, the treatment T is min-max scaled to lie in [0, 1]. MLPY . The input to MLPY is XY ∪ Xconf as well as a treatment value t ∈ [0, 1]. This MLP can thus generate both factual and counterfactual outcomes. Similar to MLPT , the output is a scalar which is the CEPO µt (X), while a random node ηY in the hidden layers is chosen as noise scale parameter. During the factual forward pass, the output T of MLPT is passed as input, and the factual outcome Y is generated by Y = µT (X) + σ(µ) · ηY · ϵ, 5
ϵ ∼ N (0, 1),
(9)
where σ(µ) is the empirical standard deviation of the CEPOs over the N samples. To generate CEPO for counterfactual treatment values t ̸= T , one passes t as input and selects the output node µt (X) without adding the noise ηY . In summary, this method constructs a prior over possible DGPs which arise in the backdoor setting of causal inference, with causal graph as shown in Figure 3. We also designed two additional priors that employed polynomial basis sampling and value-based sampling from MLP-generated tables. These priors are further described in Appendix B.2, and we include them in ablation studies.
X
T
Y
Model Architecture. We use a similar architecture to Figure 3: Causal graph associCausalPFN [3], with a PFN-style transformer encoder that lever- ated with the backdoor setting. ages in-context learning to learn the parameterized CEPO-PPD qθ . The central new architectural piece of our method CCPFN is a separate encoder for treatments to ensure that the signal from the treatment variable is not lost in high-dimensional settings. Input tokens t are passed directly through a nonlinear T -encoder, and in addition are separately appended to x and passed through a linear encoder (reminiscent of S-Learners [27]); see Figure 4. Moreover, as we work in the continuous treatment setting, we z-standardize all outcomes together (rather than splitting into control/treatment groups) as well as all treatment values. Given a query unit x, a grid of treatment values can be passed to CCPFN as a batch to estimate the full range of the ITRC t 7→ µt (x) in a single forward pass. Training. At each step of training, a DGP ψ ∼ π is sampled to yield a SCM which is generated by the three MLPs described above. We generate a dataset N xn , tn , yn , t′n , µt′n (xn ) n=1
Transformer
Encoder
Encoder
Encoder
of both factual and counterfactual scenarios. Counterfactual treatments are sampled uniformly on [0, 1]. This Concat ensures that the loss in (5) is minimized over the whole T Y X treatment range T = [0, 1]. At training time, the dataset is randomly shuffled and a context/query split position M is chosen. The observational dataset Dobs = Figure 4: Illustration of the tri-encoder schematic used by CCPFN. Treatments {(xn , tn , yn )}M n=1 is passed as context, and the counter- T are additionally routed through a sep′ N factuals {(xn , tn )}n=M +1 are passed as the query. arate encoder to boost treatment signal.
4
Experimental Setup
The fundamental problem of causal inference—namely, the inability to observe counterfactual outcomes—poses a unique challenge to experiments and model evaluation in causal inference [20]. In the absence of randomized controlled trials (RCTs), which are often impractical or unethical in various domains, causal inference models are tested on synthetic or semi-synthetic data [24, 35, 48, 33]. We follow this approach, implementing several synthetic and semi-synthetic benchmarks for both our validation and test datasets. Scenario-Generation Pipeline. While in the binary treatment regime public benchmarks have become more available recently, there is a serious lack of a broad suite of benchmarks in the continuous treatment regime. Part of the difficulty lies in having to store a potentially infinite number of counterfactuals per individual. To address these issues, we created a (semi-)synthetic scenario-generation pipeline. Each scenario represents a causal inference scenario/context (for example, the administration of a particular drug to a particular set of patients). The base covariates Dcov = {xn }N n=1 are loaded in from either an existing synthetic generation mechanism or from realworld tabular data. We then construct propensity and outcome models for each scenario based on the context and covariates provided. For the validation datasets, we tasked an LLM agent with generating plausible propensity models and outcome models for the given covariates (see Appendix C.2). For the test datasets, we used or modified propensity and outcome models from previous benchmarks. This pipeline generates causal DGPs for the backdoor scenario (Figure 3). We ensured that approximately 50% of covariates were direct causes of both treatment and outcome to ensure a high degree of confounding. This confounding present in the DGPs ensures that validating against these scenarios is 6
a genuinely causal task (rather than e.g. simply correlation matching). To verify the LLM-produced scenarios, we manually inspected each scenario and sampled the ITRCs they produced, using domain expertise to check consistency and provide a sanity check. With this method, we generated 14 scenarios (8 validation, 6 test). These span a wide range of contexts, including drug dosage effects, algorithmic marketing, debt forgiveness strategies, and news article engagement. We include plots of example treatment-response functions from several scenarios in Figure 5. More details can be found in Appendices A.1 and C.2.
Figure 5: Example individual treatment-response curves (ITRCs) for four of our validation scenarios (ACIC2016, Criteo, Hilltrom, and Twins). Solid curves are ITRCs for randomly-selected individuals; circles are the corresponding observed (T, Y ). Note that the observations do not lie precisely on the ITRCs due to the presence of exogenous noise; the curves represent expectations (CEPOs), not exact counterfactual outcomes. Benchmarks. Our test data benchmarks are comprised of one fully synthetic dataset and five semi-synthetic datasets that are often used to study causal inference, and span the domains of medicine, finance, and user research. These datasets are not used for training or any model tuning. Further details about the benchmarks are provided in Appendix A.2.
Table 1: Summary of test datasets used as benchmarks. Dataset
# Samples
# Covariates
MVICU [42, 48] Debt [33] News [42, 48] NewsHet [48] TCGA [42, 48] Warfarin [23, 49, 24]
4,963 10,000 7,881 7,881 4,428 4,490
13 10 2,870 2,870 4,000 19
Validation Datasets. Our validation datasets, used to tune model hyperparameters during training, were created by a mix of manual specification and LLM-generated scenarios. The prompts we used to generate DGPs can be found in Appendix C.2. We specifically make use of covariates from the ACIC2016 [11], ACIC2018 [25], Criteo [12], Hillstrom [17], Lalonde [10, 22], Lenta [28, 44], Twins [35], and X5 [50, 44] datasets, which are commonly employed in causal literature. Metrics. We use two primary metrics to evaluate our model’s ability to perform two key tasks: 1. Reconstruct individual treatment-response curves; 2. Prescribe the optimal treatment level/dosage for each individual. To evaluate performance on the first task, we use the mean integrated squared error (MISE) [42] between the true treatment-response curve µt and the model’s predicted µ̂t , which is reported as the 7
mean of the CEPO-PPD qθ (t, x). We normalize by the treatment interval length to standardize. This is calculated over all N individuals in a given dataset and over the whole treatment range [a, b]: N
1 X 1 MISE = N n=1 b − a
Z b
2
µ̂t (xn ) − µt (xn ) dt,
(10)
a
where xn is the covariate vector of the nth individual. To evaluate performance on the second task, we use the mean dosage policy error (DPE) [42]. This is the error between the true CEPO at the true optimal dosage/treatment level t∗ and the true CEPO at the model’s estimated optimal dosage t̂∗ , averaged over the whole dataset: DPE =
N 2 1 X µt∗ (xn ) (xn ) − µt̂∗ (xn ) (xn ) , N n=1
(11)
where t∗ (x) is the optimal dosage for an individual with covariates x. We consider in this paper optima that are either global minima or maxima, depending on the scenario. The estimated optimal dosage t̂∗ (x) of a model is calculated by first reconstructing the model’s estimate of the full treatmentresponse curve t 7→ µt (x) on a finite mesh and choosing t̂∗ (x) to be the argmin/argmax. The two metrics are clearly related, while capturing different aspects of model performance: MISE measures the model’s ability to accurately reconstruct the shape of the entire treatment-response curve over the full range of possible treatments (Task 1), while DPE measure the model’s ability to accurately prescribe an optimal treatment level for a given individual (Task 2). Baseline Methods. We compare CCPFN to various causal inference models which are trained on each individual dataset. We use neural network-based methods, including ADMIT [48], SCIGAN [5], DRNet [42], and VCNet [36]. We follow [36] and apply target regularisation on DRNet and VCNet, and add a conditional density estimation head for DRNet. We also compare to statistical baselines such as GPS [18], which is implemented using the causal_curve package [26], and EBCT [46], with VCNet employed as its inference network to estimate the ITRC. Finally, we test against causal forest double machine learning (DML) [9] and non-parametric DML, which is available through the EconML package [4]. We also benchmark CCPFN against top tabular foundation models (TabDPT [31], TabPFN-2.6 [15], and TabICLv2 [39]) implemented as S-Learners [27], which are not causally fine-tuned per dataset, but exhibit competitive performance (see Table 2). Crucially, practitioners lack access to ground-truth CEPOs during deployment, rendering baseline performance on benchmark datasets with respect to metrics such as MISE and DPE unsuitable for practical hyperparameter tuning. Therefore, during hyperparameter optimization, we select hyperparameters that minimize the training loss. The final reported performance is evaluated via 5-fold cross-validation on each test dataset. Further details on hyperparameter optimisation and our evaluation protocol are available in Appendix A.3 and A.4. Task 1: Estimation of treatment-response curves. We first evaluate CCPFN and the baseline models on their ability to estimate a full individual treatment-response curve. The synthetic/semisynthetic nature of our benchmarks gives us access to the full DGP that generates outcomes, thus allowing us to retrieve the ground-truth ITRC as well. We use the MISE metric (10) to evaluate performance on this task. Task 2: Optimal policy prediction. We also evaluate each model’s ability to prescribe an optimal dosage/treatment level. Specifically, we want to test how well a model can predict the treatment t∗ = t∗ (x) that yields the optimal value of the treatment-response function µt (x). We use the DPE metric (11) to evaluate performance on this task. This metric is not applicable to all our benchmarks, specifically when the corresponding treatment-response functions are all monotonic, which occurs in the Debt dataset. Moreover, all tested models reported essentially zero DPE on MVICU. Thus we only report DPE for Warfarin, TCGA, News, and NewsHet. We perform ablations on the choice of prior and other design choices (such as the loss function, and whether or not to enforce positivity in the prior). We evaluate these in the context of both Task 1 and Task 2. 8
5
Results
CCPFN demonstrates superior performance on the treatment-response curve estimation task across multiple benchmark datasets. When evaluated by average rank relative to mean integrated squared error (MISE), CCPFN consistently outperforms all baselines, achieving top results on the Debt and Warfarin datasets and coming in first place overall (Table 2). Table 2: Comparative evaluation of mean MISE across benchmark test datasets. Columns correspond to different benchmark datasets; values represent mean MISE ± standard deviation as computed with 5-fold cross-validation. First place is bold, second place is underlined. Dashes (—) indicate no meaningful results were obtained. When evaluating TabPFN we apply PCA to reduce the dimension to 100, due to memory constraints and to match the dimensionality reduction used in CCPFN. DRNet, VCNet, and EBCT did not produce meaningful results for the MISE metric and hence are omitted. Mean MISE (↓ better) Method
MVICU (×103 )
Debt (×10−2 )
Warfarin
TCGA
News
NewsHet (×10−2 )
CCPFN (Ours)
1.45 ± 2.4
2.26 ± .19
40.4 ± 5.7
8.63 ± 1.8
1.58 ± .11
5.76 ± .39
2.8
S-Learner (TabDPT) S-Learner (TabPFN) S-Learner (TabICL)
1.53 ± 2.5 1.59 ± 2.4 1.49 ± 2.5
2.41 ± .22 5.80 ± .51 3.70 ± .59
47.4 ± 8.1 159 ± 37 54.8 ± 11
8.47 ± 1.3 5.61 ± .90 4.01 ± .94
1.55 ± .05 1.49 ± .06 1.68 ± .05
6.28 ± .16 6.33 ± .11 5.38 ± .14
3.3 4.3 3.5
ADMIT SCIGAN GPS CausalForest NonparamDML
.250 ± .06 2.61 ± 2.2 3.09 ± 3.0 70.7 ± .00 96.0 ± .00
2.64 ± .33 4.44 ± .07 16.8 ± .32 28.9 ± 0.0 39.4 ± .57
87.9 ± 60 711 ± 44 418 ± 49 894 ± 24 1230 ± 49
6.21 ± 1.7 3.75 ± .91 — 8.71 ± 7.9 493 ± 4.6
1.67 ± .10 1.88 ± .09 1.70 ± .09 1.78 ± .57 1.63 ± .39
8.12 ± .32 4.72 ± .41 6.36 ± .13 4.50 ± .05 7.19 ± .05
4.3 5.0 6.8 6.7 7.8
Avg. Rank
Similarly, our model remains competitive on the DPE task (Table 3). Notably, the performance of neural-network and statistical baselines varies significantly across datasets, suggesting a reliance on extensive hyperparameter tuning. On the other hand, CCPFN yields consistent results, despite the fact that the model does not see any evaluation data during pre-training. Table 3: Comparative evaluation of mean DPE across benchmark test datasets. Columns correspond to different benchmark datasets; values represent mean DPE ± standard deviation as computed with 5-fold cross-validation. First place is bold, second place is underlined. When evaluating TabPFN we apply PCA to reduce the dimension to 100, due to memory constraints and to match the dimensionality reduction used in CCPFN. Mean DPE (↓ better)
Method −3
Debt (×10 CCPFN (Ours) S-Learner (TabDPT) S-Learner (TabPFN) S-Learner (TabICL) ADMIT SCIGAN GPS CausalForest NonparamDML DRNet VCNet EBCT
0.29 ± .22 8.24 ± 3.2 72.5 ± 13 14.4 ± 6.2 0.26 ± .50 0.00 ± .00 678 ± 180 0.00 ± .00 28.2 ± 5.7 64.7 ± 70 207 ± 140 327 ± 240
)
Warfarin
TCGA
News
NewsHet (×10−3 )
Avg. Rank
2.91 ± 1.7 2.45 ± 1.0 1.92 ± .82 0.31 ± .21 0.10 ± .08 150 ± 150 80.9 ± 2.6 2010 ± 24 2000 ± 49 2.19 ± 1.5 1800 ± 1000 1400 ± 1300
38.9 ± 8.1 35.5 ± 8.2 35.5 ± 8.3 35.8 ± 6.7 24.6 ± 5.4 38.6 ± 89 — 33.9 ± 7.9 23.6 ± 4.6 29.5 ± 6.6 31.6 ± 12 45.0 ± 12
3.71 ± .60 3.76 ± .59 2.69 ± .30 3.78 ± .41 3.65 ± .67 3.88 ± .54 3.46 ± .38 4.37 ± .57 3.81 ± .39 6.42 ± 2.0 5.00 ± 1.1 6.42 ± 3.2
1.64 ± .57 0.74 ± .32 1.14 ± .48 .292 ± .06 0.63 ± .04 6.91 ± .57 1.04 ± 1.2 6.84 ± .47 6.84 ± .47 5.35 ± 4.6 6.84 ± .47 1.91 ± 2.8
6.0 5.0 4.8 4.6 2.2 7.6 6.3 7.2 7.0 6.8 8.6 9.8
ITRC Reconstruction To qualitatively assess the model’s ability to capture complex treatmentresponse dynamics, we visualize its predicted curves against the ground truth individual treatmentresponse curves. We focus on the Warfarin dataset, where the outcome represents the loss between a patient’s administered dose and the optimal dose t∗ , as determined by the IWPC pharmacogenetic dosing algorithm [23, 24]. The outcome is defined in terms of absolute distance from the optimal dosage t∗ for each individual, and a successful model must accurately recover its characteristic V-shape. Figure 6 illustrates the predictions of the five top-performing methods. While all methods capture the general V-shape, CCPFN (shown in blue) tracks the ground truth more closely, particularly demonstrating superior accuracy near the endpoints of the dosage range. Ablations. We perform ablations on important design choices related to our prior. We find a significant performance gain from including in-pass tabular “corruption” [45, 30] (e.g. binarizing, 9
(a) Individual A’s ITRC
(b) Individual B’s ITRC
Figure 6: Predicted individual treatment-response curves (ITRCs) and true ITRC for two randomlyselected individuals.
quantizing, or zero-inflating randomly-selected nodes) as opposed to solely post-hoc (Table 4). Next, we reconsider the necessity of imposing positivity in the prior for DGP generation. Positivity is one of the assumptions needed for identifiability of the treatment-response curve in theory. Indeed, we find that positivity is a necessary factor, and the prior without it generates data that is less effective for training CCPFN (Table 5). Further ablations can be found in Appendix D. Table 4: Ablation on applying tabular “corruption” during the data-generation process vs. only afterwards (post-hoc). Reported metrics are mean MISE and mean DPE (lower is better). Debt (×10−2 ) Method
MVICU (×10−3 )
Warfarin
TCGA
NewsHet (×10−2 )
News
MISE DPE MISE DPE MISE DPE MISE DPE MISE DPE MISE DPE
Ours 2.22 0.011 1.45 Post-Hoc Corruption 3.75 0.311 1.53
– –
36.6 54
2.62 2.66
8.63 9.11
38.9 36.5
1.57 1.53
3.80 3.22
5.58 0.152 6.24 0.252
Table 5: Ablation on enforcing positivity in the prior on DGPs. No positivity enforced entails having MLPT output the observed T directly, with no noise node ηT added. Reported metrics are mean MISE and mean DPE (lower is better). Debt (×10−2 ) Method
MVICU (×10−3 )
TCGA
News
NewsHet (×10−2 )
MISE DPE MISE DPE MISE DPE MISE DPE MISE DPE MISE DPE
Ours 2.22 0.011 1.45 No Positivity Enforced 9.98 88.6 1.48
6
Warfarin
– –
36.6 39.7
2.62 3.34
8.63 7.11
38.9 36.0
1.57 1.68
3.80 4.07
5.58 0.152 6.97 0.316
Conclusion, Limitations, and Future Work
We introduce CCPFN, a causal foundation model for the continuous treatment setting. It demonstrates superior ability to reconstruct continuous individual treatment-response curves without any further fine-tuning on unseen datasets. A central contribution is our novel 3-MLP prior which naturally generates causal DGPs satisfying unconfoundedness and positivity. We are limited by the lack of broad, large-scale causal evaluation datasets for the continuous treatment setting. Moreover, CCPFN relies on the unconfoundedness and positivity assumptions. Unconfoundedness is impossible to verify in practice, and domain experts must be consulted to assess whether or not it is valid. Meanwhile, in the continuous treatment setting, it is impossible for positivity to hold exactly, as the possible treatment range is uncountably infinite. We are also limited by our embedding dimension, as datasets with large features are compressed into lower dimensional representations. There are many possible directions of future research. In regions of the treatment range T where observational data is scarce, we expect the model to have larger epistemic uncertainty regarding the 10
ITRC values. As CCPFN already predicts the full CEPO-PPD, this is a natural next step. It would also be of interest to extend to the case of multi-arm treatments as in [42].
11
References [1] Ahmed M. Alaa and Mihaela van der Schaar. Bayesian inference of individualized treatment effects using multi-task gaussian processes. In Proceedings of the 31st International Conference on Neural Information Processing Systems, NIPS’17, page 3427–3435, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964. [2] Susan Athey and Guido W. Imbens. The state of applied econometrics: Causality and policy evaluation. Journal of Economic Perspectives, 31(2):3–32, May 2017. doi: 10.1257/jep.31.2.3. URL https://www.aeaweb.org/articles?id=10.1257/jep.31.2.3. [3] Vahid Balazadeh, Hamidreza Kamkari, Valentin Thomas, Benson Li, Junwei Ma, Jesse C. Cresswell, and Rahul G. Krishnan. CausalPFN: Amortized Causal Effect Estimation via In-Context Learning. In Advances in Neural Information Processing Systems, volume 38, 2025. [4] Keith Battocchi, Eleanor Dillon, Maggie Hei, Greg Lewis, Paul Oka, Miruna Oprescu, and Vasilis Syrgkanis. EconML: A Python Package for ML-Based Heterogeneous Treatment Effects Estimation. https://github.com/py-why/EconML, 2019. Version 0.15.0. [5] Ioana Bica, James Jordon, and Mihaela van der Schaar. Estimating the effects of continuousvalued interventions using generative adversarial networks. In Advances in Neural Information Processing Systems, volume 33, pages 16434–16445, 2020. [6] Léon Bottou, Jonas Peters, Joaquin Quiñonero-Candela, Denis X. Charles, D. Max Chickering, Elon Portugaly, Dipankar Ray, Patrice Simard, and Ed Snelson. Counterfactual reasoning and learning systems: The example of computational advertising. Journal of Machine Learning Research, 14(101):3207–3260, 2013. URL http://jmlr.org/papers/v14/bottou13a. html. [7] Magnus Bühler, Lennart Purucker, and Frank Hutter. Causal data augmentation for robust fine-tuning of tabular foundation models. arXiv:2601.04110, 2026. [8] Lucius E. J. Bynum, Aahlad Manas Puli, Diego Herrero-Quevedo, Nhi Nguyen, Carlos Fernandez-Granda, Kyunghyun Cho, and Rajesh Ranganath. Black Box Causal Inference: Effect Estimation via Meta Prediction. arXiv:2503.05985, 2025. [9] Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. Double/debiased machine learning for treatment and structural parameters. The Econometrics Journal, 21(1), 2018. doi: 10.1111/ectj.12097. [10] Rajeev H. Dehejia and Sadek Wahba. Causal effects in nonexperimental studies: Reevaluating the evaluation of training programs. Journal of the American Statistical Association, 94(448): 1053–1062, 1999. [11] Vincent Dorie, Jennifer Hill, Uri Shalit, Marc Scott, and Dan Cervone. Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition. Statistical Science, 34(1):43–68, 2019. [12] Eustache Diemert, Artem Betlei, Christophe Renaudin, and Amini Massih-Reza. A large scale benchmark for uplift modeling. In Proceedings of the AdKDD and TargetAd Workshop, KDD, London,United Kingdom, August, 20, 2018. ACM, 2018. [13] Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems, volume 27, 2014. [14] Brett R. Gordon, Florian Zettelmeyer, Neha Bhargava, and Dan Chapsky. A comparison of approaches to advertising measurement: Evidence from big field experiments at facebook. Marketing Science, 38(2):193–225, 2019. [15] Léo Grinsztajn, Klemens Flöge, Oscar Key, Felix Birkel, Philipp Jund, Brendan Roof, Benjamin Jäger, Dominik Safaric, Simone Alessi, Adrian Hayler, Mihir Manium, Rosen Yu, Felix Jablonski, Shi Bin Hoo, Anurag Garg, Jake Robertson, Magnus Bühler, Vladyslav Moroshan, Lennart Purucker, Clara Cornu, Lilly Charlotte Wehrhahn, Alessandro Bonetto, Bernhard Schölkopf, 12
Sauraj Gambhir, Noah Hollmann, and Frank Hutter. TabPFN-2.5: Advancing the state of the art in tabular foundation models. arXiv: 2511.08667, 2025. [16] Léo Grinsztajn, Klemens Flöge, Oscar Key, Felix Birkel, Philipp Jund, Brendan Roof, Benjamin Jäger, Dominik Safaric, Simone Alessi, Adrian Hayler, Mihir Manium, Rosen Yu, Felix Jablonski, Shi Bin Hoo, Anurag Garg, Jake Robertson, Magnus Bühler, Vladyslav Moroshan, Lennart Purucker, Clara Cornu, Lilly Charlotte Wehrhahn, Alessandro Bonetto, Bernhard Schölkopf, Sauraj Gambhir, Noah Hollmann, and Frank Hutter. Tabpfn-2.5: Advancing the state of the art in tabular foundation models, 2026. URL https://arxiv.org/abs/2511.08667. [17] Kevin Hillstrom. Minethatdata e-mail analytics and data mining challenge. https://blog. minethatdata.com/2008/03/minethatdata-e-mail-analytics-and-data.html, March 2008. [18] Keisuke Hirano and Guido W. Imbens. The Propensity Score with Continuous Treatments, chapter 7, pages 73–84. John Wiley & Sons, Ltd, 2004. ISBN 9780470090459. doi: https: //doi.org/10.1002/0470090456.ch7. [19] Keisuke Hirano and Guido W. Imbens. The propensity score with continuous treatments. In Andrew Gelman and Xiao-Li Meng, editors, Applied Bayesian Modeling and Causal Inference from Incomplete-Data Perspectives, pages 73–84. Wiley, 2004. [20] Paul W. Holland. Statistics and causal inference. Journal of the American Statistical Association, 81(396):945–960, 1986. ISSN 01621459, 1537274X. [21] Noah Hollmann, Samuel Müller, Katharina Eggensperger, and Frank Hutter. TabPFN: A transformer that solves small tabular classification problems in a second. In The Eleventh International Conference on Learning Representations, 2023. [22] Nick Huntington-Klein. causaldata: Example data sets for causal inference textbooks, 2024. URL https://pypi.org/project/causaldata/. Python package. [23] International Warfarin Pharmacogenetics Consortium. Estimation of the warfarin dose with clinical and pharmacogenetic data. New England Journal of Medicine, 360(8):753–764, 2009. doi: 10.1056/NEJMoa0809329. [24] Nathan Kallus and Angela Zhou. Policy evaluation and optimization with continuous treatments. In Proceedings of the Twenty-First International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pages 1243–1251, 09–11 Apr 2018. [25] Ehud Karavani, Yishai Shimoni, and Chen Yanover. IBM causal inference benchmarking framework, January 2018. URL https://doi.org/10.5281/zenodo.1163587. [26] Roni W Kobrosly. Causal-curve: a python causal inference package to estimate causal doseresponse curves. Journal of Open Source Software, 5(52):2523, 2020. [27] Sören R. Künzel, Jasjeet S. Sekhon, Peter J. Bickel, and Bin Yu. Metalearners for estimating heterogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 116(10):4156–4165, 2019. doi: 10.1073/pnas.1804597116. [28] Lenta and Microsoft. BIGTARGET Hackathon Dataset, 2020. URL https://www.kaggle. com/datasets/mrmorj/bigtarget. [29] Fan Li, Peng Ding, and Fabrizia Mealli. Bayesian causal inference: A critical review. Philosophical Transactions of the Royal Society A, 381(2247):20220153, 2023. [30] Junwei Ma, Nour Shaheen, Alex Labach, Amine Mhedhbi, Frank Hutter, Anthony L Caterini, and Valentin Thomas. Generalization can emerge in tabular foundation models from a single table. arXiv:2511.09665, 2025. [31] Junwei Ma, Valentin Thomas, Rasa Hosseinzadeh, Alex Labach, Hamidreza Kamkari, Jesse C. Cresswell, Keyvan Golestan, Guangwei Yu, Anthony L. Caterini, and Maksims Volkovs. TabDPT: Scaling Tabular Foundation Models on Real Data. In Advances in Neural Information Processing Systems, 2025. 13
[32] Yuchen Ma, Dennis Frauen, Emil Javurek, and Stefan Feuerriegel. Foundation models for causal inference via prior-data fitted networks. In The Fourteenth International Conference on Learning Representations, 2026. [33] Javier Moral Hernández, Clara Higuera-Cabañes, and Álvaro Ibraín. An end-to-end pipeline for Causal ML with continuous treatments: An application to financial decision making. In 3rd Workshop on Causal Inference and Machine Learning in Practice, 2025. [34] Brady Neal. Introduction to Causal Inference from a Machine Learning Perspective. Course Lecture Notes, 2020. URL https://www.bradyneal.com/Introduction_to_Causal_ Inference-Dec17_2020-Neal.pdf. [35] Brady Neal, Chin-Wei Huang, and Sunand Raghupathi. RealCause: Realistic Causal Inference Benchmarking, 2021. URL https://arxiv.org/abs/2011.15007. [36] Lizhen Nie, Mao Ye, Qiang Liu, and Dan Nicolae. VCNet and Functional Targeted Regularization For Learning Causal Effects of Continuous Treatments. In International Conference on Learning Representations, 2021. [37] Arman Oganisian and Jason A. Roy. A practical introduction to Bayesian estimation of causal effects: Parametric and nonparametric approaches. Statistics in Medicine, 40(2):518–551, 2021. [38] Jingang Qu, David Holzmüller, Gaël Varoquaux, and Marine Le Morvan. TabICL: A Tabular Foundation Model for In-Context Learning on Large Data. In Proceedings of the 42nd International Conference on Machine Learning, volume 267 of Proceedings of Machine Learning Research, pages 50817–50847, 13–19 Jul 2025. [39] Jingang Qu, David Holzmüller, Gaël Varoquaux, and Marine Le Morvan. TabICLv2: A better, faster, scalable, and open tabular foundation model. arXiv:2602.11139, 2026. [40] Jake Robertson, Arik Reuter, Siyuan Guo, Noah Hollmann, Frank Hutter, and Bernhard Schölkopf. Do-PFN: In-Context Learning for Causal Effect Estimation. In Advances in Neural Information Processing Systems, 2025. [41] Donald B. Rubin. Bayesianly justifiable and relevant frequency calculations for the applied statistician. The Annals of Statistics, pages 1151–1172, 1984. [42] Patrick Schwab, Lorenz Linhardt, Stefan Bauer, Joachim M. Buhmann, and Walter Karlen. Learning counterfactual representations for estimating individual dose-response curves. Proceedings of the AAAI Conference on Artificial Intelligence, 34(04):5612–5619, 2020. ISSN 2159-5399. doi: 10.1609/aaai.v34i04.6014. [43] Uri Shalit, Fredrik D. Johansson, and David Sontag. Estimating individual treatment effect: generalization bounds and algorithms. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 3076–3085. PMLR, 06–11 Aug 2017. URL https:// proceedings.mlr.press/v70/shalit17a.html. [44] Maksim Shevchenko and contributors. scikit-uplift: Uplift modeling in scikit-learn style, 2020. URL https://github.com/maks-sh/scikit-uplift. [45] Yi Sui, Tongzi Wu, Jesse Cresswell, Ga Wu, George Stein, Xiaoshi Huang, Xiaochen Zhang, and Maksims Volkovs. Self-supervised representation learning from random data projectors. In International Conference on Learning Representations, 2024. [46] Stefan Tübbicke. Entropy balancing for continuous treatments. Journal of Econometric Methods, 11(1):71–89, 2022. doi: doi:10.1515/jem-2021-0002. [47] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, 2017. [48] Xin Wang, Shengfei Lyu, Xingyu Wu, Tianhao Wu, and Huanhuan Chen. Generalization bounds for estimating causal effects of continuous treatments. In Advances in Neural Information Processing Systems, volume 35, pages 8605–8617, 2022. 14
[49] M. Whirl-Carrillo, R. Huddart, L. Gong, K. Sangkuhl, C. F. Thorn, R. Whaley, and T. E. Klein. An evidence-based framework for evaluating pharmacogenomics knowledge for personalized medicine. Clinical Pharmacology & Therapeutics, 110(3):563–572, 2021. doi: 10.1002/cpt. 2350. [50] X5 Retail Group and ODS.ai. X5 RetailHero Uplift Modeling Dataset, 2019. URL https: //ods.ai/competitions/x5-retailhero-uplift-modeling.
15
A
Further Details About Benchmarks and Baselines
A.1
Synthetic and Semi-Synthetic Data Scenarios
Each subclass of Scenario at minimum implements the following three methods: • load_covariates: Generates the base covariates Dcov = {xn }N n=1 . For semi-synthetic data, these are loaded from real data; for fully synthetic data, these are generated synthetically. • treatment: Generates the treatment administered to each individual. • dose_response: The noise-free treatment-response function (x, t) 7→ µt (x) for this Scenario. Example plots of treatment-response curves generated by each scenario can be found in Appendix A.2 (test datasets) and Appendix C.2 (validation datasets). A.2
Benchmark Details
We used 6 different held-out datasets for testing CCPFN against other causal inference methods. For each set of covariates, a dataset provides a factual treatment and outcome, as well as counterfactual treatments along with their corresponding CEPO. • MVICU. A semi-synthetic medical scenario adapted from [42] and post-processed. It models the effect of different configurations of mechanical ventilation in the intensive care unit on patients. For treatment and outcome functions, we used the same ones as in News [48]. It contains 13 covariates and 4,963 rows. (PhysioNet Credentialed Health Data License 1.5.0) • Debt. A fully synthetic financial scenario concerning the impact of different levels of debt writedowns on total debt repaid, adapted from [33]. The generated dataset has 10 covariates and 10,000 rows. (MIT license) • News. A semi-synthetic scenario representing reader interaction with news articles from [42, 48]. It has 2,870 covariates. Using the original dataset downloaded from [42], we use the process outlined in [48] to process the data. Specifically, we randomly subsampled 10,000 rows. From this, we further subsampled 7,881 by dropping rows whose optimal CEPO (as determined by the oracle) was greater than 10. (MIT license) • NewsHet. A semi-synthetic scenario adapted from News, with treatment-response functions modified to introduce heterogeneity of optimal dosage levels. As implemented in [48], News suffers from the fact that all optima occur at the endpoints 0 or 1, making the optimal dosage policy prediction a degenerate task We introduce heteroscedastic horizontal shifts to rectify this. (MIT license) • TCGA. A semi-synthetic scenario representing the effect of medication dosage and risk of cancer recurrence from [42, 48], where we use the DGP outlined by [48] to generate ground-truth causal effects. We additionally removed rows whose optimal CEPO was greater than 50, leaving 4,428 rows. (MIT license) • Warfarin. A semi-synthetic medical scenario concerning warfarin dosing, adapted from [23, 49, 24]. The dataset specifically calculates the loss between the actual dosage received by an individual and the optimal warfarin dosage calculated by the IWPC pharmacogenetic algorithm [23] and uses this as the outcome. The generated dataset has 19 covariates and 4,490 rows. (CC BY-SA 4.0) Example individual treatment-response curves are shown for each benchmark dataset in Figure 7. A.3
Hyperparameter Optimization
EconML Baselines We tune results from the EconML Baselines (CausalForestDML and NonParamDML) using the FLAML (AutoML) library. Hyperparameter tuning is performed on the treatment model and the outcome model. We use k-fold cross-validation with k = 5, early stopping, and a time budget of 900 seconds. The following base estimators are used: "lgbm", "xgboost", "xgb_limitdepth", "rf", "kneighbour", "extra_tree". Neural Network Methods For the neural network baselines, we conduct a grid search over learning rates ∈ {0.005, 0.001, 0.0005, 0.0003, 0.0001} and batch sizes ∈ {128, 256, 512, 1024}. For each dataset, we sample 80% of the data and perform 5-fold cross-validation on this sample to identify the hyperparameter configuration yielding the lowest average training loss for each baseline. We use the training loss because in practice, one never has access to the synthetic validation metrics used. Once 16
Figure 7: Example individual treatment-response curves (ITRCs) for all six test scenarios. Solid curves are ITRCs for randomly-selected individuals; circles are the corresponding observed (T, Y ). Note that the observations do not lie precisely on the ITRCs due to the presence of exogenous noise; the curves represent expectations (CEPOs), not exact counterfactual outcomes.
the optimal hyperparameters are identified, we perform 5-fold cross-validation on the entire training dataset to elicit final performance. A.4
Evaluation Protocol
The fully optimized versions of all considered models (CCPFN, TFMs, neural networks, EconML, and statistical baselines) are evaluated on the benchmark datasets using 5-fold cross-validation. For the neural network, statistical, and EconML baselines, the training splits are used to explicitly fit the models to each benchmark scenario. Conversely, for CCPFN and TFMs, these training splits are provided solely as in-context examples, leaving the models’ underlying weights frozen. The final reported MISE and DPE metrics for each method are the mean error across the five validation folds as well as the standard deviation across these five folds.
B
Further Detail on Synthetic Priors
B.1
Our 3-MLP Prior
During training, the number of samples N is fixed to be 2048. The maximum number of covariates is 98. The number of layers LX , LT , LY are all sampled from independent truncated normal distributions with mean and variance α > 0, where α ∼ LogUniform(A, B). In our experiments, we 17
set A = 1, B = 10 for X, T , and Y . The number of hidden units HX , HT , HY are also sampled from the same class of distributions, with A = 10, B = 100 for X, T , and Y . The noise scale is also sampled from this class of distributions, with A = 1 × 10−4 and B = 0.5. The densities dX , dT , dY were sampled independently from U(0.1, 1). To ensure that the signal from T is not lost when dropping edges, we protect all outgoing edges from T in the input layer of MLPY . In the truncated normal distributions, we truncate to [3, ∞) for layer distributions and [4, ∞) for hidden size distributions. Integer values are then obtained by rounding.
Figure 8: Example individual treatment-response curves (ITRCs) from different DGPs produced by our 3-MLP prior. Solid curves are ITRCs for randomly-selected individuals; circles are the corresponding observed (T, Y ). Note that the observations do not lie precisely on the ITRCs due to the presence of exogenous noise; the curves represent expectations (CEPOs), not exact counterfactual outcomes.
B.2
Alternative Priors
We compare two alternative prior distributions over data-generating processes. Bernstein polynomial prior. We sample a table of N rows and partition its columns into covariates X and a noise column ϵ. Covariates are standardised, and treatments are drawn from a sigmoid-normal distribution, 2 t | x ∼ σ N (µt|x , σt|x ) , (12) where µt|x and σt|x are the outputs of a conditional MLP applied row-wise to x. Overlap is controlled by scaling σt|x by an overlap parameter α ∈ (0, 1]. The conditional expected potential outcome (CEPO) µt (x) is defined via a K-degree Bernstein polynomial, whose coefficients are formed as a convex combination of individual-specific coefficients c(x) (produced by a second conditional MLP) and a shared global coefficient vector c0 ∼ N (0, I): c = λ c(x) + (1 − λ) c0 ,
(13)
where λ ∈ [0, 1] governs the degree of treatment effect heterogeneity. The observed outcome is then y = µt (x) + ϵ, where ϵ is a scaled, centred noise term derived from the reserved noise column. Value-based prior. We again sample a table of N rows and select a subset of columns as covariates X. Rather than parameterising the treatment-response curve analytically, we directly read off potential outcomes at n randomly sampled (and sorted) treatment support points {t1 , . . . , tn } ⊂ [0, 1]. For each support point tk , we reserve one table column as the CEPO µtk (x) and a second column as an individual-level noise term ηk (x); noise columns are scaled so that their contribution is a fixed fraction of the corresponding signal variance. Observed outcomes at any treatment value t are obtained by linearly interpolating both µtk (x) and ηk (x) across the two nearest support points and summing the results. Treatments are sampled using the same sigmoid-normal mechanism as in the Bernstein prior. 18
C
Model Training Details
C.1
Model Architecture and Hyperparameters
We use a modified version of the TabDPT architecture [31] with a nonlinear T -encoder (see Figure 4). We initialize our weights to TabDPT’s trained weights on layers which support them. When the number of covariates K exceeds 100, we apply truncated SVD to the covariates before passing them to the model. During training, we used a batch size of 32, 8 gradient accumulation steps, and 128 model updates per epoch. Our final model was fine-tuned for 15 epochs; this number was selected by early stopping based on the validation dataset results. The σ used in the histogram loss (6) is 0.01. We use 20 transformer layers with 6 heads per layer, an embedding dimension of 384, and a feed-forward hidden dimension of 768. The final model architecture has 19,140,352 parameters. Of these, 297,600 are the nonlinear T encoder parameters, compared to 39,168 for the linear X-and-T -encoder and 768 for the Y -encoder. All training runs were performed on single NVIDIA A6000 GPUs. C.2
Validation Dataset Construction
The datasets we used for model validation and hyperparameter selection were a mix of synthetic and semi-synthetic. We used the fully synthetic dataset constructed in [48], as well as a simple linear dataset with a linear treatment-response function. We chose the best version of our model based on averaged IQR-normalised MISE performance across all validation datasets. For semi-synthetic data, we used covariates from several well-known causal inference datasets: ACIC2016 (Apache license) [11], ACIC2018 (Apache license) [25], Criteo (MIT license) [12], Hillstrom (MIT license) [17], Lalonde (MIT License) [10, 22], Lenta (MIT license) [28, 44], Twins (MIT license) [35], and X5 (MIT license) [50, 44]. We tasked an LLM agent with generating a plausible scenario for each set of covariates using the prompts included below. Example individual treatment-response curves for each of the validation datasets are shown in Figure 9. System Prompt for Generating Synthetic Validation Data #1 # Semi-synthetic data generation instructions (Note: run python scripts in the ‘conda tracee‘ environment.) ## Background You are working on a project in causal inference. The goal is to train a model to perform causal inference in the *continuous treatment* setting. ## Your Task Your task is to create semi-synthetic data consisting of real-world covariates ‘X‘ ( real data) and synthetic treatment and outcome variables by creating synthetic data-generating processes (DGPs). Adhere to the following instructions: 1. Ask the user which ‘csv‘ file to use as the base covariates ‘X‘. It is possible that there is no local csv, and the covariates will have to be downloaded in the script itself (e.g. using sklearn.datasets). You can download and view the covariates now, so that you have intuition for the context. 2. Ask the user for covariate context, i.e. what do the base covariates ‘X‘ represent in this dataset? 3. Ask the user for treatment and outcome context, i.e. what scenario the user has in mind for the treatment and outcomes. 4. Based on the information provided in steps 1 - 3, devise a *realistic* DGP to simulate treatment assignment and outcomes. Remember, the treatment variable should be continuous, *not* binary. This DGP should satisfy the following requirements: 1. There should be a high degree of confounding: at least 50% of the covariates should be causes of both the treatment and the outcome. 2. You should generate a *dose-response function* f(X, t) that maps an individual with covariates X and hypothetical treatment t to the *conditional expected
19
potential outcome (CEPO)*. This should be a suitably complex and realistic function which can be implemented in simple Python code. 3. You should generate a *treatment assignment function* T(X) that maps an individual with covariates X to the *observed* treatment T(X). This should be a suitably complex and realistic function which can be implemented in simple Python code. 4. In order to ensure that there is a high degree of confounding, the functions f and T should both depend on some subset of covariates comprising at least half of the total number of covariate features. 5. Once you have constructed this DGP, generate a *Python script* that outputs a csv file as follows: 1. Ask the user for the desired name of the Python script. 2. The Python script should include code for the dose-response function f(X, t) and the treatment assignment function T(X). 3. The Python script should output a single csv file with columns named x_0 through x_n (where n is the number of covariate features), t, y, t_test, cepo_test. The data should be filled as follows: 1. The values of columns x_0 through x_n should be the values of the original base covariates csv. 2. The value of t should be the value of T(X) for X the corresponding covariate value. 3. The value of y should be f(X, t) for X the corresponding covariate value and t = T(X), *plus Gaussian noise* which is iid for each row. 4. The value of t_test should be randomly sampled from [t_min, t_max]. 5. The value of cepo_test should be f(X, t_test). 6. All data should be numerical (e.g. string-based categorical variables should be encoded as integers). 4. Save the python script in tracee/inference/benchmarks/data_generation_scripts. When you are ready to proceed with this task, begin at step 1 above.
System Prompt for Generating Synthetic Validation Data #2
#Semi-synthetic data generation instructions ##Background You are working on a project in causal inference. The goal is to train a model to perform causal inference in the continuous treatment setting. ##Your Task Your task is to create semi-synthetic data consisting of real-world covariates X (real data) and synthetic treatment and outcome variables by creating synthetic datagenerating processes (DGPs). The dataset we are working with is the Lalonde dataset. You can decide the best way to access this dataset (realcause, jobs etc .) Only use the base covariates. Propose a potential CONTINUOUS-VALUED treatment (T) and CONTINUOUS-VALUED outcome (Y). You should then realise a DGP to generate this treatment and outcome. Some things to keep in mind are: We are working in the backdoor causal graph scenario, where X -> T, T -> Y, and X-> Y There should be a high degree of confounding - at least 50% of the covariates should affect outcome Y You should generate a dose-response function f(X, t) that maps an individual with covariates X and hypothetical treatment t to the conditional expected potential outcome (CEPO). This should be a suitably complex and realistic function which can be implemented in simple Python code. You should generate a treatment assignment function T(X) that maps an individual with covariates X to the observed treatment T(X). This should be a suitably complex and realistic function which can be implemented in simple Python code. This DGP should be generated in Python code. You should then save the data into a csv file as follows, where columns are X_0, ..., X_n, t, y, t_test, cepo_test. n refers to the number of covariates we originally had. Each row represents one individual in the original lenta dataset Rows X_0, X_n should retain the base covariate values from the original dataset
20
t = T(X) for corresponding X value (where X represents the set of covariates) y = f(X, t) for X the corresponding covariate value and t + Gaussian noise which is iid for each row. t_test should be randomly sampled from a suitable range based on the nature of the continuous treatment you propose cepo_test = f(X, t_test) for the corresponding t_test This should be a python script, with a main() function. You should also randomly sample only 4000 rows to include in the csv.
D
Ablation Studies
We perform ablations on important design choices related to our prior. In addition to those reported in Section 5, we experiment with a simplified prior design that uses a single MLP (Table 6), but find that the 3-MLP design acheives greater performance, especially on the DPE metric. Finally, instead of the histogram variant of the causal data-prior loss, we test a loss based on the continuous ranked probability score. In Table 7 we find that the CRPS loss gives slight improvement on TCGA, but otherwise is not as effective as the causal data-prior loss. Table 6: Ablation on the prior design. Ours is the 3-MLP prior discussed in Section 3; 1-MLP is a single MLP prior. Reported metrics are mean MISE and mean DPE (lower is better). Debt (×10−2 ) Method
MVICU (×10−3 )
Warfarin
TCGA
News
NewsHet (×10−2 )
MISE DPE MISE DPE MISE DPE MISE DPE MISE DPE MISE DPE
Ours (3-MLP) 2.22 0.011 1.45 1-MLP 5.32 0.122 1.67
– –
2.62 3.08
36.6 27.2
8.63 4.23
38.9 37.8
3.80 4.37
1.57 1.56
5.58 0.152 5.75 0.250
Table 7: Ablation on choice of loss function (CRPS v.s. CE). Reported metrics are mean MISE and mean DPE (lower is better). Debt (×10−2 ) Method
MVICU (×10−3 )
Warfarin
TCGA
News
NewsHet (×10−2 )
MISE DPE MISE DPE MISE DPE MISE DPE MISE DPE MISE DPE
Ours (CE) 2.22 0.011 1.45 CRPS 3.06 0.039 1.57
– –
36.6 38.9
2.62 2.73
21
8.63 8.03
38.9 37.2
1.57 1.58
3.80 3.80
5.58 0.152 5.90 0.450
Figure 9: Example individual treatment-response curves (ITRCs) for all eight validation scenarios. Solid curves are ITRCs for randomly-selected individuals; circles are the corresponding observed (T, Y ). Note that the observations do not lie precisely on the ITRCs due to the presence of exogenous noise; the curves represent expectations (CEPOs), not exact counterfactual outcomes.
22