CausalPFN: Amortized Causal Effect Estimation via In-Context Learning

We thank the reviewer for their thoughtful review and constructive feedback. Below, we have provided detailed responses to each of the concerns raised. If our clarifications address these points satisfactorily, we would be very grateful if you consider raising the score.

Baselines are somewhat outdated, and thus, the benefit in performance is unclear. Although they used the CATENets library, they did not compare SOTA DNN-based methods in it, such as SNet and FlexTENet.

We have used two widely-used libraries for heterogeneous treatment effect estimation: EconML (4.2k stars on Github), and CATENets (another widely used library), in addition to more classical R-based baselines like GRF and BART. Moreover, we tuned the hyperparameters carefully for most of the baselines. Based on the reviewer’s suggestion, we also provide additional baselines from the CATENets library, namely SNet, FlexTENet, as well as XNet.

Here, we compare CausalPFN to these additional baselines on the IHDP, ACIC, and Lalonde datasets for both ATE estimation (relative error) and CATE estimation (PEHE). We perform hyperparameter tuning similar to that used for the other baselines in the CATENets library, and still show competitive performance often exceeding that of these new baselines.

1) Mean PEHE ± Standard Error (↓ better)

2) Mean ATE Relative Error ± Standard Error (↓ better)

Causal inference usually involves statistical performance issues, and computational performance is not particularly important in practical terms. The advantage at this point may not be appealing to readers. Therefore, speed alone is not a compelling reason for adoption. Authors might want to add some explanations of use-cases where speed is important, or of the value added by high-speed inference.

First, we note that speed is not the main advantage of CausalPFN. As shown in Figure 1, as well as Tables 1 and 2, CausalPFN provides superior performance (for CATE estimation) on average, even without taking the computational performance into account.

Second, while it is true that many traditional causal inference applications prioritize statistical performance, a large number of real-world industry use-cases require fast/adaptive causal estimators. For instance, real-time bidding systems must estimate the incremental effect of ads and make bidding decisions with minimum latency [A]. Another example is e-commerce personalization systems that rely on fast uplift models to tailor recommendations in the course of a short user session, where latency directly impacts conversion rates [B].

Therefore, even though speed is not a primary advantage of our method, we find it to be beneficial in many practical scenarios. We will add further clarifications and add some discussion to further clarify this in our revision.

[A] Yuan et al. "Real-time bidding for online advertising: measurement and analysis." Proceedings of the seventh international workshop on data mining for online advertising. 2013.

[B] Sawant et al. "Contextual multi-armed bandits for causal marketing." arXiv:1810.01859.

[Q1] Since quite high accuracy is reported for IHDP CATE and Lalonde CPS ATE, is there any concern that the pre-training data might contain these datasets? How is it possible to confirm that such dataset leaks have not occurred?

This is quite an important concern for all large-scale training paradigms, including CausalPFN. We made sure that none of the benchmarking datasets was part of the pre-training data based on the following:

1. Most of the pre-training data used in our model are purely synthetic, following a similar approach of using randomly generated MLPs as in TabPFN v1 [42]. Specifically, no information regarding the benchmarks is used during this synthetic generation.

2. For the real-world part of the pre-training data, we only use real-world covariates and generate the treatments through synthetic functions and add synthetic noise to the CEPO columns to generate potential outcomes. In that sense, there is only a chance for covariate leakage. However, we have also published the list of real-world datasets (OpenML datasets) in footnote 2 on page 28 of the appendix, which shows that none of the benchmarks are part of those datasets.

[Q2] Is it possible to train a model that directly estimates CATE instead of CEPO? That is, a transformer that takes (X_train, T_train, Y_train, X_test) as input and outputs τ_test? In causal inference, it is important to focus on the causal effect τ itself, which is often relatively smaller than the variation of Y on X (since we focus only on interventionable factors as ). The proposed method selects Y1 and Y0 randomly from X, which does not reproduce such a situation.

Yes, our framework allows us to train CausalPFN to directly estimate CATE values instead of separate estimation of CEPO values $\mathbb{E}[Y_1 | X]$ and $\mathbb{E}[Y_0|X]$. We chose to estimate CEPOs instead of CATE values mainly for the following reason: when calibrating the estimator, we often do not have access to ground-truth CATE values but factual observations, which can correspond to $Y_0$ or $Y_1$ based on the treatment value. Therefore, with the current choice, we can use the observational data to calibrate the model, as we detail in lines 271-307. Ultimately, estimating CEPOs can easily lead to many well-known causal estimands, such as CATE, ATE, local average treatment effects, and can potentially generalize to multiple-treatment settings.

We thank the reviewer for their positive feedback and encouraging comments. We are glad they found the paper well-written, the experiments extensive, and the combination of amortization and PFNs promising. We also appreciate their view that this work could become a valuable tool for practitioners.

[1,2,3] are very closely related papers. Some even doing causal inference via PFNs similar to the current work. While all of these papers have been very recent (most of them actually only available after the NeurIPS deadline), it might still be good to discuss these approaches at least or potentially even compare e.g. to [3].

Thank you for pointing out these related works. Except for [1], which is already discussed in the paper as reference [14], the other papers first became public on arxiv approximately one month after the conference submission deadline. Among the mentioned papers, [3] does not provide any code for using their model and mainly benchmarks on synthetic datasets that we do not have access to. Similarly, we could not find any code for the method in [1]. Specifically, as mentioned in Section 3.6 (page 9) of [1], the method requires having access to the underlying family of data-generating processes for each benchmark (e.g., for IHDP, they use 13,000 training datasets sampled from the underlying outcome surface model). In contrast, CausalPFN provides a zero-shot estimator. Below, we compare to Do-PFN [2], which is the most similar method to ours.

We use the IHDP, ACIC, and Lalonde datasets for both ATE estimation (relative error) and CATE estimation (PEHE). Note that Do-PFN is also a zero-shot method so we do not run any hyperparameter tuning or pre-processing on the datasets.

1) Mean PEHE ± Standard Error (↓ better)

2) Mean ATE Relative Error ± Standard Error (↓ better)

As it is clear from both tables, CausalPFN achieves superior performance in all the benchmarks. We hypothesize that this is due to the non-identifiability of the prior data that they use, which we will elaborate in the next part.

Why do you only gather identifiable DAGs? Would you not want the distribution over all equivalent DAGs no matter if they are identifiable or not? This would probably give you less clean estimates but would that not better? I.e. right now the argument seems that you only care about a DAG if it is identifiable while in the real world that might not be the case.

We have the following reasons to only include identifiable data-generating processes (DGP) in our prior generation:

1. Non-identifiability of the DGP means that for the same observational distribution, we can get an unbounded set of causal effects. Therefore, if we allow for non-identifiable settings, in the worst case, we can have DGPs that all agree on the observational data (the only available information during inference) while having a range of $(-\infty, \infty)$ for the causal effects, which lacks any useful insights. Specifically, this worst-case scenario is more likely when generating large-scale prior data. This is perhaps one of the reasons we see worse performance for Do-PFN compared to our CausalPFN, as demonstrated in the previous part, since Do-PFN uses non-identifiable prior data, potentially having unobserved confounding.

2. More concretely, we have provided an updated version of Proposition 1 here, that states that identifiability is necessary for convergence of the posterior-predictive distribution of causal effects (the quantity that CausalPFN aims to learn) in the limit:

Proposition 1 (Informal). Under mild regularity assumptions, for almost all $\psi^\star \sim \pi$ and any set of i.i.d. samples $\mathcal{D}\_\text{obs}\sim\mathrm{P}^{\psi^\star}\_\text{obs}$, we have that as $|\mathcal{D}\_\text{obs}|\to\infty$,

$ \mathbb{E}^{\pi^{\mu\_t}} \bigl[\mu \mid **X**,\mathcal{D}\_\text{obs}\\bigr] \overset{a.s.}{\longrightarrow} \mu_t(**X**; \psi^\star), \quad \forall t \in \mathcal{T}, \quad **X**\sim\mathrm{P}^{\psi^\star}, $

if and only if the support of the prior $\pi$ only contains $\psi$ with identifiable CEPOs almost everywhere.

What this theory directly implies is that the identifiability (or in our case ignorability) requirement of the DGPs is not only sufficient for consistent estimation of the causal effects, but it is also a necessity. We have a formal proof for this proposition, which we will include in the final version of the paper. We would also be happy to provide the proof during the discussion period if requested.

Why do you only provide inference code and not training code?

We will release the training code, along with the synthetic prior data used for the training, with the published version of the paper.

We thank the reviewer for their thoughtful review, insightful questions, and their positive assessment of the paper, as well as their view that CausalPFN is both convincing and easy to use.

More Confident Model With More Data

CausalPFN becomes more confident as it observes more data. As an example, we use the synthetic polynomial datasets described in L1262–L1275 of the appendix and provide the length of the 95% credible interval of CATE estimates from CausalPFN, along with the confidence interval derived from a DR-Learner:

Table: The Effect of Sample Size on the Length of the 95% Credible Interval

Although more conservative than the doubly-robust learner, the length of CausalPFN’s credible interval goes down as we observe more samples. Moreover, to ensure well-calibration, we also calculate the integrated coverage error (ICE):

CausalPFN’s ICE for Different Sample Sizes

The Limit of Covariate and Sample Size

During the causal phase of the training, we consider sample sizes up to 2048 and covariates up to 99. However, during inference, CausalPFN can take $50{,}000$ (Table 2) or more (Table 5).

To assess the effect of sample/covariate size on CausalPFN’s performance, we provide additional experiments. Similar to above, we focus on the synthetic polynomial datasets. We fix the test size at 100 samples in all experiments, and for each setting of sample size and covariate size, we report the average±standard error over 50 tables generated from the polynomial prior with a fixed random seed.

1) Effect of Sample Size: Here, we consider the same data-generating process while increasing the number of samples. We fix the number of covariates to 10. The table below shows the PEHE values for CausalPFN’s CATE compared to other baselines. It is clear that CausalPFN has a better ability to utilize large sample sizes, as the PEHE drops faster.

2) Effect of Covariate Size: Here, we fix the number of samples to 1000 and increase the number of covariates. Similar to above, the table below compares the performance of CausalPFN to other methods in terms of PEHE. Although CausalPFN consistently maintains a performance gap, the gap becomes smaller as the number of covariates increases. This is likely due to the limited number of covariates during training (up to 100) and can potentially be addressed by increasing the number of covariates during training.

Ensuring Training Datasets Did Not Contain Test Data

1. Most of the pre-training data used in our model are purely synthetic, following a similar approach to TabPFN [42]. We will publish the training code, along with the pre-training data, for the final version.

2. For the real-world part of the pre-training data, we only use real-world covariates and generate the treatments through synthetic functions and add synthetic noise to the CEPO columns to generate potential outcomes. We included the list of real-world datasets (OpenML datasets) in footnote 2 on page 28 of the appendix, which shows that none of the benchmarks are part of those datasets.

Extension to Time Series Tasks

Potentially, there have already been a few attempts to use in-context learners on time-series data [B, C].

[B] Hoo, Shi Bin, et al. "From Tables to Time: How TabPFN-v2 Outperforms Specialized Time Series Forecasting Models." arXiv:2501.02945

[C] Taga, Ege Onur, et al. "TimePFN: Effective multivariate time series forecasting with synthetic data." AAAI 2025.

L58: The Ignorability Assumption

CausalPFN is only applicable for the ignorability assumptions where identifiability holds. In most cases with no ignorability (i.e., with unobserved confounding), the causal effects can be generally unbounded. This is why many standard causal effect estimators are based on this assumption, and rely on a domain expert to assess whether ignorability holds.

L158: ATE Calculation by Averaging Across Units in $\mathcal{D}\_\text{obs}$

We run the model to estimate CATE for each observed unit and calculate the average of all to estimate the ATE.

L192: Is Positivity Guaranteed?

Positivity is guaranteed in the population as the treatment logits go through a Sigmoid function to generate propensities $P(T=1|X) \in (0, 1)$. Still, there is a small chance that the data will only contain one treatment due to finite samples. We explicitly discard those edge cases during training.

Continuous Treatments

In principle, the CEPO-PPDs are still well-defined and the causal data-prior loss should still yield valid CEPO-PPDs. The challenge here, as correctly mentioned by the reviewer, is to implement a diverse and smooth data generation pipeline.

L207: Histogram Loss

As provided in Section B.2 L1175-L1179 of the appendix, the data is normalized first, then 1024 equidistant bins are allocated from -10 to +10; every point outside of this [-10, +10] is treated as an outlier and clamped.

Once the prediction is done in [-10, +10], we use the stored scale and shift parameter using for normalization to undo it and bring it back to the original scale.

Metric Definitions

For CATE evaluation, for a test dataset of size $N$ and ground-truth CATEs $\tau$ alongside their estimations $\hat{\tau}$, the evaluation metric, PEHE, is defined as:

$\text{PEHE}(\hat{\tau}) = \sqrt{\frac{1}{N} \sum_{n=1}^{N} (\hat{\tau}(**x**^{(n)}) - {\tau}(**x**^{(n)}))^2}$

For ATE evaluation, given a ground-truth ATE $\lambda$ and its estimation $\hat{\lambda}$, the relative error is defined as: $\frac{|\lambda - \hat{\lambda}|}{|\lambda|}$.

We will include these equations in the revised version.

Why are relative errors directly averageable and PEHEs not?

PEHE depends on the scale of CATE values. Hence, averaging across different datasets would be biased towards datasets with larger output values, while the scale of ATE is normalized in calculating the relative error. For more consistency, we will update Table 1 with average rank instead of average relative errors for ATE estimation, similar to CATE.

Why not use relative errors for CATE?

PEHE is a standard metric used for CATE estimation in the causal inference literature. Moreover, the CATE error is defined as an average metric over $N$ samples, while the ATE error is defined only for one value per dataset.

L243: ITE vs CATE

We consider CATE, but the conditioning variables consist of all the observed covariates. In that sense, it is close to ITEs (though we still estimate an expected value). We will clarify this in the paper.

Eq 10: What is $\lambda$?

$\lambda$ coincides with ATE in equation (2) (same notation). We will recall its definition for more clarity.

Table2: Score Normalization

We normalize each column by dividing all values by the maximum score in that column. This normalization enables us to compute an average performance metric across datasets

Table 3: Fine-tuning vs. Performance

In Table 3, the first row corresponds to CausalPFN, while the second and third rows report results for a different underlying architecture (TabPFNv2): one without fine-tuning (row 3) and one with fine-tuning on our causal prior (row 2). In all cases, fine-tuning improves performance when comparing the second and third rows.

Comparison to "Towards causal foundation model”

We have already briefly discussed this method (see [110]). Unfortunately, we could not find any implementation for this method. Moreover, as mentioned in Section 5.3 in [110], the zero-shot version of their method, CInA (ZS), is only applicable for benchmarks that naturally come with multiple datasets and there is no clear instruction on how to use the zero-shot version for standard casual effect estimation tasks akin to the ones we have used.

We thank the reviewer for their effort and considering our paper as novel, well-motivated, clearly explained, and our experiments as extensive. In what follows, we have provided detailed responses to each individual comment:

Details regarding method scalability are not discussed. How large are the datasets used in terms of sample size and number of covariates? How many covariates are considered both during training and evaluation?

In the appendix (footnote 2 on page 28), we provided an anonymized link to the list of all the real-world datasets (OpenML tables), which includes details such as sample sizes and the number of features. For the causal phase of the training (synthetic datasets), the following table summarizes the characteristics of the tables used, which we will include in our revision.

Table: Description of Training Data

For evaluation, we used standard benchmarks, and we detail the number of samples and covariates below:

Table: Description of Benchmarks

For the marketing datasets, as stated in the paper, we perform 5-fold honest splitting. We evaluate both the original sizes (Table 5) and a stratified 50k sub-sampling (Table 2).

Table: Description of Benchmarks

To handle more than 99 covariates, we apply PCA to the covariates and retain only the first 99 principal components to obtain a condensed representation of the covariates that can fit into the input of the model.

How does the method’s performance in terms of accuracy scale with the number of covariates compared to other methods?

Based on your suggestion, we conducted an additional ablation experiment to assess the effect of the number of covariates on the performance of CausalPFN compared to other methods. For this ablation experiment, we focus on the synthetic polynomial datasets described in lines (1262–1275) of the appendix. We fix the test size at 100 samples and number of training samples as 1000, and increase the number of covariates. For each covariate size, we report the average ± standard error of PEHE of CATE estimates over 50 tables generated from the polynomial prior with a fixed random seed to control for randomness:

Table: Effect of Covariate Size

The table suggests that CausalPFN consistently maintains a performance gap over different numbers of covariates. However, the gap becomes smaller as the number of covariates increases. This is likely due to the limited number of covariates during training (up to 100). We believe that increasing the number of covariates during training can address this issue. Still, we observe that CausalPFN can generalize well to a larger number of covariates.

Could you provide a version of Figure 1 that isolates inference time only, excluding training and hyperparameter search for the baseline methods? This would help in fairly assessing the computational efficiency of CausalPFN.

We included the training and hyperparameter tuning time for other baselines in Figure 1 as this is what would normally happen in practice: when faced with a causal task, domain experts typically search over a set of possible design choices, including different hyperparameters. On the other hand, CausalPFN is only pre-trained once, and for each new dataset, it only requires running a few forward passes of the transformer model. We want to emphasize that Figure 1 is not intended to claim that the inference speed of CausalPFN is inherently faster than that of the baselines; rather, it illustrates what a practitioner would typically experience in terms of end-to-end runtime when choosing an approach.

Having said that, we agree with the reviewer that providing inference-only time can provide more insights into the efficiency of CausalPFN. However, including the inference-only time requires re-training all the baselines, which can take more than a week of computation and is infeasible for the rebuttal period. Instead, we provide results that we already had available that contain the training + inference time for the baselines, excluding the hyperparameter optimization time.

Table: Median Time for CATE Estimation (Seconds per 1000 Samples)