Enhancing Treatment Effect Estimation with Generation-Driven Data Augmentation

Other comments

Clarifying $Y_{i,k}^{(gen)}$: We account for the potential stochastic nature of generative models (e.g. LLMs) by sampling $K$ outcomes $\{ Y_{i,k}^{(gen)} | k=1, ..., K \}$ from the distribution $P_{t, x_i}^{(gen)}$ for a given covariate $x_i$. We then aggregate this set into a single outcome, by computing the average $\frac{1}{K} \sum_{i=1}^{K} Y_{i,k}^{(gen)}$ , which defines $\bar{Y}_{i}^{(gen)}$.

We recognize that this was not clear in our original manuscript, which is why we have added an extended clarification at l.264.

Improvements to Figure 1. Thank you for pointing out the lack of clarity of Figure 1. We have made the changes to the notation used in the figure, as well as in the figure caption, to ensure that Figure 1 allows to comprehensively understand our work. Please let us know if there are any further suggestions that we can incorporate.

We hope the reviewer’s concerns are addressed and they will consider updating their score. We welcome further discussions.

References

Manousakas, Dionysis, and Sergül Aydöre. "On the usefulness of synthetic tabular data generation." arXiv preprint arXiv:2306.15636 (2023).

We appreciate the reviewer’s detailed and thoughtful feedback.

Clarifying the Contribution of this Work

Let us clarify the key novelties of our work in the following points.

Problem setting. Our work addresses a particularly challenging yet overlooked scenario in CATE estimation: the small-sample regime under high covariate shift between treatment groups. As demonstrated in Section 6, conventional approaches which rely on model specifications (e.g. inverse-propensity weighting, representation learning) typically struggle in this setting.

General augmentation framework. To address this problem, we introduce $`GATE`$, a general and model-agnostic framework to improve any CATE learner via data augmentation. $`GATE`$ is flexible, since it allows to leverage diverse generative models. Our theoretical analysis in Section 3.1 reveals the benefits of the $`GATE`$ framework, namely an increase of the effective sample size and the mitigation of covariate shift, which can offset potential inaccuracies of the generative model.

Difference with traditional data augmentation. Furthermore, the $`GATE`$ framework contrasts traditional data augmentation in that it does not require the generation of covariates, a complex task which could introduce additional covariate shift. Rather, we rely on generating the missing potential outcomes for already existing covariates. This approach directly reduces the covariate shift and as such, can provide performance benefits which can be strong enough to counterbalance the bias potentially introduced by an imperfect generative model. This distinguishes $`GATE`$ from traditional data augmentation, methods where inaccurate generative models can easily degrade performance (Manousakas & Aydöre, 2023).

External knowledge integration. Motivated by the inherent limitations of small observational datasets, we instantiate $`GATE`$ with LLMs to illustrate the advantages of generative models which incorporate external sources of knowledge. This key novelty in the context of CATE estimation represents a significant departure from existing imputation approaches, such as Aloui et al. (2023), which are ultimately contrained by the observational data (e.g. using local regression models). We further introduce a new selection mechanism based on the LLM's uncertainty to define the admissible set $\mathcal{X}_{t}$ used for augmentation.

Empirical performance. Our comprehensive set of experiments shows that this LLM-based instantiation of $`GATE`$ has the following benefits: 1) it consistenly improves the PEHE across $11$ CATE learners and $3$ datasets 2) it outperforms other instantiations of $`GATE`$ which use generative models trained solely on $D^{(obs)}$, with performance gains especially noticeable in the small-sample regime ($\rho=0.1$). In Section 6.4 we provide additional insights showing the importance of eliciting the prior knowledge of LLMs to achieve optimal performance gains. Further, in Section 6.5 we provide a detailed analysis verifying the theoretical insights about data augmentation for CATE estimation presented in Section 3.

We acknowledge that these key contributions were not stated clearly enough in our original manuscript. Hence, we have revised it to clarify its key novelty, with specific changes detailed in our global response.

Quantitative assessment verifying the covariate shift reduction

We would like to point out that the following experiments, present in our original manuscript, illustrate that the augmented data effectively mitigates covariate shift.

Quantifying the reduction in covariate shift. The experiment described in Section 6.5.2 demonstrates how the covariate shift between the control and treated groups decreases with increasing values of the parameter $\alpha$, as evidenced in Figure 5 Right. We quantified the distributional distance between the control and treated groups by computing the Sliced Wasserstein (SW) distance between the covariates in the control and treated groups. We provide more details about this metric in Appendix D.4.

GATE is particularly helpful in high covariate shift settings. Having shown that $`GATE`$ permits to effectively mitigate covariate shift, we demonstrated in Section 6.5.1 and Figure 5 Left that its performance gains were the most pronounced when the covariate shift between the treated and control groups in the observational dataset is important.

Therefore, these two experiments confirm the theoretical insights derived from our generalization bound (Theorem 3.1): augmenting the observational data with $`GATE`$ allows to reduce the covariate shift, hence leading to improved CATE estimation.

References

Van Breugel, Boris, et al. "Decaf: Generating fair synthetic data using causally-aware generative networks." Advances in Neural Information Processing Systems 34 (2021): 22221-22233.

Calmon, F. et al. (2017). Optimized pre-processing for discrimination prevention. Advances in neural information processing systems, 30.

Curth, Alicia, and Mihaela Van Der Schaar. "In search of insights, not magic bullets: Towards demystification of the model selection dilemma in heterogeneous treatment effect estimation." International Conference on Machine Learning. PMLR, 2023.

Dwork, C. et al. (2012). Fairness Through Awareness. In Proceedings of the 3rd Innovations in Theoretical Computer Science Conference (Cambridge, Massachusetts) (ITCS '12). ACM, New York, NY, USA, 214–226.

Feldman, M. et al. (2015). Certifying and removing disparate impact. In proceedings of the 21th ACM SIGKDD international conference on knowledge discovery and data mining (pp. 259-268).

Hardt, M. et al. (2016). Equality of opportunity in supervised learning. In Advances in neural information processing systems. 3315–3323.

Huang, Yuheng, et al. "Look before you leap: An exploratory study of uncertainty measurement for large language models." arXiv preprint arXiv:2307.10236 (2023).

Lan, Hui, and Vasilis Syrgkanis. "Causal q-aggregation for cate model selection." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.

Lin, Zhen, Shubhendu Trivedi, and Jimeng Sun. "Generating with confidence: Uncertainty quantification for black-box large language models." arXiv preprint arXiv:2305.19187 (2023).

Mahajan, Divyat, et al. "Empirical Analysis of Model Selection for Heterogeneous Causal Effect Estimation." arXiv preprint arXiv:2211.01939 (2022).

Manakul, Potsawee, Adian Liusie, and Mark JF Gales. "Selfcheckgpt: Zero-resource black-box hallucination detection for generative large language models." arXiv preprint arXiv:2303.08896 (2023).

Mehrabi, N. et al. (2021). A survey on bias and fairness in machine learning. ACM computing surveys (CSUR), 54(6), 1-35.

Saito, Yuta, and Shota Yasui. "Counterfactual cross-validation: Stable model selection procedure for causal inference models." International Conference on Machine Learning. PMLR, 2020.

Schuler, Alejandro, et al. "A comparison of methods for model selection when estimating individual treatment effects." arXiv preprint arXiv:1804.05146 (2018).

Wei, Jason, et al. "Chain-of-thought prompting elicits reasoning in large language models." Advances in neural information processing systems 35 (2022): 24824-24837.

Yao, Shunyu, et al. "Tree of thoughts: Deliberate problem solving with large language models." Advances in Neural Information Processing Systems 36 (2024).

Bias with data augmentation

We thank the reviewer for mentioning the potential bias of generative models. Let us explain how the practitioner using the LLM-instantiated $`GATE`$ has various tools at their disposal to address this point:

Auditing the generation of outcomes: Contrasting other generative models, the use of LLMs with $`GATE`$ permits to make the generation process more transparent. Beyond producing numerical values, LLMs can also elicit explicit causal reasoning chains underpinning outcome generation when prompted accordingly. This capability enables human-in-the-loop applications of $`GATE`$, where domain experts can evaluate the generated data by examining these reasoning traces against their domain knowledge. As such, we view the use prompting techniques for explicit reasoning, such as chain-of-through (Wei et al., 2022) or tree-of-thought (Yao et al., 2024), as a promising direction for future work.

Fairness analysis: Since $`GATE`$ is an augmentation method, practitioners can directly assess potential bias in the generated potential outcomes: (1) by doing a standard fairness analysis and computing different fairness metrics (Van Breugel et al., 2021; Dwork et al., 2012; Hardt et al., 2016; Mehrabi et al., 2021) on the augmented data. (2) If some bias is detected, the practitioner could use off-the-shelf debiasing methods (Calmon et al., 2017; Feldman et al., 2015).

Action Taken: We have added a discussion about the potential bias of the LLMs into our Appendix B.2.

Code release

We provide an anonoymized version of our code at the following link: https://anonymous.4open.science/r/GATE-8849

We hope the reviewer’s concerns are addressed and they will consider updating their score. We welcome further discussions.

Clarifying the Variance-Based Definition of the Admissible Set

Variance and imputation error. Let us stress that our proposed variance-based selector is specifically meant for the instantiation of $`GATE`$ with stochastic models trained on external datasets, such as LLMs, and is based on their unique characteristics. Indeed, in the context of LLMs, it has been shown that uncertainty measures such as variance can be used to discriminate between factually correct and incorrect responses (Huang et al., 2023; Manukal et al., 2023), as well as predict the quality of a response (Lin et al., 2023), thus validating the use of variance as the selection criterion. In case of deterministic models trained on $\mathcal{D}^{(obs)}$, rather than relying on variance (which then indeed is equal to zero), we propose to rely on other metrics, such as the propensity score (as we explain in l.293 in our manuscript).

Comparison with other selectors. We performed an additional experiment, detailed in Appendix G.3, to compare the variance-based selector with (1) a random selector and (2) a propensity-based selector, where we define the score function as $s(x,t) = P(T = t|X=x)$. We report the results obtained with the LLM instantiation of $`GATE`$ in Table 7, fixing $\alpha = 0.5$. These results show that the variance-based selector achieves optimal performance most consistently out of the considered selection criteria, with performance gains especially noticeable in the small-sample regime ($\rho = 0.1$).

Tuning the hyperparameter $\alpha$. The tuning of the hyperparameter $\alpha$ is linked to the broader problem of model selection in CATE estimation. Due to the fundamental problem of causal inference, true counterfactuals are unobservable in real-world experiments, preventing straightforward model selection through validation set performances. Therefore, model selection procedures for causal inference models remain an active area of research (Lan & Syrgkanis, 2024; Curth et al., 2023; Mahajan et al., 2022; Saito & Yasui, 2020; Schuler et al., 2018), and advances in this field will naturally benefit $`GATE`$.

With these challenges in mind, we propose three complementary strategies which allow to guide the selection of the value of $\alpha$ within the $`GATE`$ framework:

1. As we explain above, relying on the fixed-threshold definition of the admissible set (rather than the percentile-based definition of the threshold) can allow to guide the selection of $\alpha$ using domain knowledge.
2. We further propose to guide the selection of $\alpha$ by measuring the covariate shift between the sets $\tilde{\mathcal{D}}_0^{(obs)}$ and $\tilde{\mathcal{D}}_1^{(obs)}$, using for example the sliced Wasserstein distance (an example of such an analysis can be found in our Figure 5 Right). Then, we propose to choose the minimal value of $\alpha$ which allows to achieve significant reduction in a covariate shift (which might correspond to the 'elbow' in the graph). While such an 'elbow' might not always exist, this criterion provides additional guidelines in certain circumstances.
3. Finally, the choice of $\alpha$ can be further guided by standard methods for CATE model selection, particularly those based on comparing the downstream CATE models using a pseudo-outcome surrogate criteria evaluated on a held-out validation set (see Curth et al. (2023) for an overview). In particular, in view of strong covariate shift and small sample regime, we propose to rely on criteria which do not rely on estimating the propensity score, as these might lead to high variance in such cases.

Nevertheless, while finding the optimal value of $\alpha$ is non-trivial, our experiments showed the following: (1) Fixing $\alpha = 0.5$ consistently led to improved dowstream PEHE across the $3$ datasets and the $11$ CATE learners (cf. Table 1) (2) The strong effect of the reduction in covariate shift obtained with data augmentation (shown in Figure 5 Right) reduces the sensitivity with respect to $\alpha$. Indeed,Figure 5 Middle highlights that any value $\alpha > 0$ leads to performance gains compared to $\alpha = 0$.

We appreciate the reviewer’s detailed and thoughtful feedback.

Clarifying the Contribution of this Work

We respectfully disagree that the contribution of our work is "incremental", and that its core contributions have been "previously explored". Let us clarify the key novelties in the following points.

Problem setting. Our work addresses a particularly challenging yet overlooked scenario in CATE estimation: the small-sample regime under high covariate shift between treatment groups. As demonstrated in Section 6, conventional approaches which rely on model specifications (e.g. inverse-propensity weighting, representation learning) typically struggle in this setting.

General augmentation framework. To address this problem, we introduce $`GATE`$, a general and model-agnostic framework to improve any CATE learner via data augmentation. $`GATE`$ is flexible, since it allows to leverage diverse generative models. Our theoretical analysis in Section 3.1 reveals the benefits of the $`GATE`$ framework, namely an increase of the effective sample size and the mitigation of covariate shift, which can offset potential inaccuracies of the generative model.

External knowledge integration. Motivated by the inherent limitations of small observational datasets, we instantiate $`GATE`$ with LLMs to illustrate the advantages of generative models which incorporate external sources of knowledge. This key novelty in the context of CATE estimation represents a significant departure from existing imputation approaches, such as Aloui et al. (2023), which are ultimately contrained by the observational data (e.g. using local regression models). We further introduce a new selection mechanism based on the LLM's uncertainty to define the admissible set $\mathcal{X}_{t}$ used for augmentation.

Empirical performance. Our comprehensive set of experiments shows that this LLM-based instantiation of $`GATE`$ has the following benefits: 1) it consistenly improves the PEHE across $11$ CATE learners and $3$ datasets 2) it outperforms other instantiations of $`GATE`$ which use generative models trained solely on $D^{(obs)}$, with performance gains especially noticeable in the small-sample regime ($\rho=0.1$). In Section 6.4 we provide additional insights showing the importance of eliciting the prior knowledge of LLMs to achieve optimal performance gains. Further, in Section 6.5 we provide a detailed analysis verifying the theoretical insights about data augmentation for CATE estimation presented in Section 3.

We acknowledge that these key contributions were not stated clearly enough in our original manuscript. Hence, we have revised it to clarify its key novelty, with specific changes detailed in our global response.

Additional experimental contributions. In addition to clarifying the contributions, we have also followed the suggestion of the reviewer to provide a "more extensive study" of generative models. In light of this, we focused on our instantiation of $`GATE`$ with LLMs, and analyzed the sensitivity of $`GATE`$ to the prompt construction with two additional experiments:

• Number of in-context samples: in Section 6.4, Figure 4, we vary the number of in-context samples provided in the prompt, showing that $`GATE`$ instantiated with LLMs benefits from the use of a wide context to elicit the LLM's in-context learning abilities.
• Selection of in-context samples: in Appendix G.2, we show that random sampling of in-context samples (among both the treated and control group) outperforms a nearest-neighbor baseline which only considers the opposite treatment group. This demontrates better in-context learning capabilities of LLM when leveraging information from both treated and control groups under high covariate shift.

Nevertheless, while finding the optimal value of $\alpha$ is non-trivial, our experiments showed the following: (1) Fixing $\alpha = 0.5$ consistently led to improved dowstream PEHE across the $3$ datasets and the $11$ CATE learners (cf. Table 1) (2) The strong effect of the reduction in covariate shift obtained with data augmentation (shown in Figure 5 Right) reduces the sensitivity with respect to $\alpha$ in high covariate shift settings. Indeed,Figure 5 Middle highlights that any value $\alpha > 0$ leads to performance gains compared to $\alpha = 0$.

Action Taken: We have added an extended discussion about hyperparameter tuning to Appendix B.2.

Reference to [1] and [2]. As the reviewer noted, our original manuscript cited [2], with references in our related work sections (Section 5, Appendix A) and an experimental comparison (in Appendix F.2 of the original manuscript). We chose to cite [2] instead of [1] since (a) the publication date of [2] is more recent than the publication date of [1], (b) we did not want to cause confusion by citing the same work twice.

We hope the reviewer’s concerns are addressed and they will consider updating their score. We welcome further discussions.

References

Aloui, Ahmed, et al. "Counterfactual Data Augmentation with Contrastive Learning." arXiv preprint arXiv:2311.03630 (2023).

Curth, Alicia, and Mihaela Van Der Schaar. "In search of insights, not magic bullets: Towards demystification of the model selection dilemma in heterogeneous treatment effect estimation." International Conference on Machine Learning. PMLR, 2023.

Huang, Yuheng, et al. "Look before you leap: An exploratory study of uncertainty measurement for large language models." arXiv preprint arXiv:2307.10236 (2023).

Johansson, Fredrik D., et al. "Generalization bounds and representation learning for estimation of potential outcomes and causal effects." Journal of Machine Learning Research 23.166 (2022): 1-50.

Lan, Hui, and Vasilis Syrgkanis. "Causal q-aggregation for cate model selection." International Conference on Artificial Intelligence and Statistics. PMLR, 2024.

Lin, Zhen, Shubhendu Trivedi, and Jimeng Sun. "Generating with confidence: Uncertainty quantification for black-box large language models." arXiv preprint arXiv:2305.19187 (2023).

Mahajan, Divyat, et al. "Empirical Analysis of Model Selection for Heterogeneous Causal Effect Estimation." arXiv preprint arXiv:2211.01939 (2022).

Manakul, Potsawee, Adian Liusie, and Mark JF Gales. "Selfcheckgpt: Zero-resource black-box hallucination detection for generative large language models." arXiv preprint arXiv:2303.08896 (2023).

Saito, Yuta, and Shota Yasui. "Counterfactual cross-validation: Stable model selection procedure for causal inference models." International Conference on Machine Learning. PMLR, 2020.

Schuler, Alejandro, et al. "A comparison of methods for model selection when estimating individual treatment effects." arXiv preprint arXiv:1804.05146 (2018).

Proportion $\rho =0.5$

Proportion $\rho =1.0$

The results demonstrate that using the variance-based selection can lead to performance improvements, particularly in the small sample regime ($\rho = 0.1$). Further, the variance-based selection provides optimal performance most consistently out of the three methods considered.

Action Taken: Following your comment, we have now added the results of the experiments presented above to Appendix G.3, with a reference to this section in l. 303 in the revised manuscript. Further, we have included a discussion of the variance-based selection in Appendix B.3.

Other questions

Size of the augmented dataset. In the instantiation of $`GATE`$ that we have used in the experiments, our main focus was to control the number of generated potential outcomes, and as a result we have decided to use a percentile-based definition of the scoring function $s(x, t)$ (where choosing $\alpha = 0.5$ guarantees that 50% of missing potential outcomes are generated, thus allowing to fix the proportion of generated outcomes across datasets).

However, in real-world applications a more optimal strategy might be to let $\alpha$ be a fixed variance threshold instead, the value of which can be guided by domain-knowledge or exploratory analysis of the data. This would more explicitly guardrail against the inclusion in the augmented dataset of particularly 'poor' generated outcomes. Then, we would define $\mathcal{X}_{t} = \\{X_i \mid i \in [n], s(X_i, T_i) < \alpha_t\\}$. We note that this in case, the proportion of generated outcomes depends on the properties of the generative model. In particular, if the model is particularly bad, no potential outcomes are generated and our method recovers the baseline performance.

Action Taken: We have added a footnote about the fixed-value threshold for the scoring function in l.322 with extended discussion in Appendix B.3.

Tuning the hyperparameter $\alpha$. The tuning of the hyperparameter $\alpha$ is linked to the broader problem of model selection in CATE estimation. Due to the fundamental problem of causal inference, true counterfactuals are unobservable in real-world experiments, preventing straightforward model selection through validation set performances. Therefore, model selection procedures for causal inference models remain an active area of research (Lan & Syrgkanis, 2024; Curth et al., 2023; Mahajan et al., 2022; Saito & Yasui, 2020; Schuler et al., 2018), and advances in this field will naturally benefit $`GATE`$.

With these challenges in mind, we propose three complementary strategies which allow to guide the selection of the value of $\alpha$ within the $`GATE`$ framework:

1. As we explain above, relying on the fixed-threshold definition of the admissible set (rather than the percentile-based definition of the threshold) can allow to guide the selection of $\alpha$ using domain knowledge.
2. We further propose to guide the selection of $\alpha$ by measuring the covariate shift between the sets $\tilde{\mathcal{D}}_0^{(obs)}$ and $\tilde{\mathcal{D}}_1^{(obs)}$, using for example the sliced Wasserstein distance (an example of such an analysis can be found in our Figure 5, right). Then, we propose to choose the minimal value of $\alpha$ which allows to achieve significant reduction in a covariate shift (which might correspond to the 'elbow' in the graph). While such an 'elbow' might not always exist, this criterion provides additional guidelines in certain circumstances.
3. Finally, the choice of $\alpha$ can be further guided by standard methods for CATE model selection, particularly those based on comparing the downstream CATE models using a pseudo-outcome surrogate criteria evaluated on a held-out validation set (see Curth et al. (2023) for an overview). In particular, in view of strong covariate shift and small sample regime, we propose to rely on criteria which do not rely on estimating the propensity score, as these might lead to high variance in such cases.

Motivation for using a variance-based selection criterion. To overcome this problem, we have proposed to use a proxy measure to quantify the fidelity of $P_{t, x}^{(gen)}$ and efficiently identify regions which should be excluded from the admissible set $\mathcal{X}_t$. In particular, we rely on the uncertainty in the generated potential outcomes (measured with variance). While, as you point out, such a measure might not always be the optimal choice for all generative models, in the context of LLMs it has been shown that uncertainty measures such as variance can be used to discriminate between factually correct and incorrect responses (Huang et al., 2023; Manakul et al., 2023), as well as predict the quality of a response (Lin et al., 2023), further validating our choice.

Correlation between variance and imputation errors. Following your question, we have computed the Spearman correlation coefficient between the variance of the generated potential outcomes and the mean squared error on the predicted counterfactual. The results are presented in the table below.

We note that the correlation between the variance and the mean squared error on the generated counterfactual outcomes are significantly larger than zero for the Hillstrom and STAR datasets, both of which are datasets with real-world outcomes obtained by subsampling a RCT dataset. In contrast, the missing potential outcomes in the Lalonde CPS1D are obtained using a black-box generative model (see Appendix D.2 for an extended discussion of the construction of the datasets).

Cautionary remark: lack of direct link between the imputation error and PEHE. In aiding the interpretation of the above results, we would like to clarify why relying on imputation error in the generated counterfactuals (for example computing an imputation MSE) to evaluate the selection mechanism might not be an optimal strategy. Indeed, the imputation MSE is not necesarily equivalent to the distributional distance between $P(Y(t) | X=x)$ and $P(Y^{(gen)}(t) | X=x)$, averaged over the different values of $X=x$, which is ultimately the quantity impacting downstream performance. In other words, two given generative models can achieve similar imputation MSE but lead to drastically different downstream PEHE.

We empirically validate this insight with an additional experiment where we compare the performance of two generative models:

• Model A: adds Gaussian noise with standard deviation $\sigma = \sqrt{5000}$ to the true counterfactual outcomes.
• Model B: adds the constant $\sigma$ to the true counterfactual outcomes.

By construction, with a reasonable number of samples, the imputation MSE of Model A will be close to $\sigma^2$, which is the imputation MSE of Model B.

We use the STAR dataset and for each model, we augment the observational dataset with its corresponding generated outcomes. We use a DR-Learner and repeat the experiment across $3$ different seeds.

The following table summarizes the results:

Despite having similar MSE, the two models lead to very different downstream $\sqrt{\epsilon_{PEHE}}$, showing that imputation error should be interpreted carefully before being used for example to evaluate a selection mechanism.

Comparison to other selection mechanisms. To evaluate the performance of the variance selection criterion more robustly, we have compared it against random selection and propensity score selction (where in the latter case we define $s(x,t) = P(T = t|X=x)$). We have included the details of this experiment in Appendix G.3, as well as presented them in the table below.

Proportion $\rho =0.1$

We appreciate the reviewer’s detailed and thoughtful feedback.

Originality

Following your suggestions, we have revised our manuscript to emphasize the paper's main contributions, including the use of external sources of data to enhance CATE estimation, which is possible through the LLM-based instantiation of $`GATE`$. We outlined in details the steps taken to achieve that in the general comment, posted above. All the introduced changes can be found in the revised version of the paper, which we have also already uploaded to OpenReview (changes are marked in teal). We thank you for the suggestion to introduce these changes, as we believe that they clarify the novelty and contributions of $`GATE`$.

Theoretical results

On top of the LLM-refocus, in our paper we have added an extended discussion comparing our theoretical results to Aloui et al. (2023), with additional references to this work (and ideas presented within) incorporated into the Introduction (l.064), Section 3 (l.207), Section 4.1 (l.239), and Section 4.2 (l.294).

However, we still believe that the result we presented in Theorem 3.1 is a fundamental part of our paper, theoretically validating our approach, and further expanding the previous results presented by Aloui et al. (2023), as we outline below. Further, we note that our Theorem 3.1, while similar, was derived using entirely different approach, relying on different metrics (IPM vs Total Variation and L2 distance) and different results from previous works (this work relies on Johansson et al. (2022), while Aloui et al. rely on Ben-David et al. (2010)).

We believe that the result we presented (Theorem 3.1) is indeed a generalisation and extension of the result presented in Aloui et al. (2023), with the following key differences:

1. The formulation of the generalisation bound presented in this work accommodates an arbitrary loss function $L$.
2. Theorem 3.1 presented in this work explicitly accounts for the potential stochasticity of generative models (such as LLMs), by employing an IPM term to measure the distributional distance between the ground-truth outcome and the generator's distribution.
3. Our bound characterizes the finite-sample regime, showing how the benefits of data augmentation correlate with sample size.

Comparison of Insights from the Theoretical Results

Further, we believe that the insights drawn from the theoretical results also differ between the two works. In particular:

1. We focus on the small sample regime, emphasising that efficient data augmentation can be particularly beneficial in the case when the observational data is small (we further verify this insight experimentally in Section 6.3).
2. We recognise that data augmentation becomes particularly important in problems with a strong covariate shift, which arises for example when the treatment is guided by expert judgment (we further verify this insight experimentally in Section 6.5).
3. We emphasise that the selection can be guided not by properties of the dataset, but rather by the properties of the generative model, which is particularly relevant when using generative models not trained on $\mathcal{D}^{(obs)}$.

Additional experimental results

To further empirically analyse the LLM-based instantiation of $`GATE`$, we have also empirically investigated the sensitivity of the LLM-instantiated $`GATE`$ to the prompt construction by conducting two additional experiments:

• Number of in-context samples: in Section 6.4, Figure 4, we vary the number of in-context samples provided in the prompt, showing that $`GATE`$ instantiated with LLMs benefits from the use of a wide context to elicit the LLM's in-context learning abilities.
• Selection of in-context samples: in Appendix G.2, we show that random sampling of in-context samples (among both the treated and control group) outperforms a nearest-neighbor baseline which only considers the opposite treatment group. This demontrates better in-context learning capabilities of LLM when leveraging information from both treated and control groups under high covariate shift.

Selection of the Missing Potential Outcomes

The difficulty of evaluating the bias of the generative model. As you have pointed out, our Theorem 3.1 indicates that to maximise the performance gains obtained by using $`GATE`$, we should aim to exclude from the admissible set those regions of the covariate space, where $P_{t, x}^{(gen)}$ significantly deviates from $P(Y(t)|X=x)$. However, this requires estimating the distance between these two distributions, which is highly non-trivial, as it would require us to obtain multiple samples from $P(Y(t)|X=x)$ for each value of $x$. This is particularly infeasible under strong covariate shift and in the small-sample regime: both cases which call for the use of $`GATE`$.

2. Theoretical Insight: Excluding from the admissible set regions of the covariate space where the generative model is particularly "incorrect" (the IPM term measuring the distance between the true distribution $Y(t) | X$ and the generated distribution $Y^{(gen)}(t) | X$ is large) can improve performance.

Design Choice: As we explain in Section 4.2, assessing the statistical distance between these two distributions is challenging, particularly in the small sample regime and for models trained on external data sources. Thus, we propose to identify such "incorrect" regions using a proxy measure: the uncertainty in the generated outcomes. Indeed, in the context of LLMs, it has been shown that uncertainty measures such as variance can be used to discriminate between factually correct and incorrect responses (Huang et al., 2023; Manakul et al., 2023), as well as predict the quality of a response (Lin et al., 2023).

Empirical Verification Through Experiment 6.5.2 (Figure 5, right): We verify that as we increase the allowed level of uncertainty of $P_{t, x}^{(gen)}$, the bias introduced by data augmentation increases, while the covariate shift decreases.

Datasets used in the experiments

How are the datasets we use in our experiments created? Let us clarify that in our experiments (Section 6), none of our benchmarks dataset is fully synthetic. Indeed, they are obtained as follows:

• Lalonde CPS1: Constructed using a generative model trained on the real-world Lalonde CPS1 dataset. The covariates are the same as in the original dataset, while the treatments and outcomes are obtained with the generative model. This generative model is extensively tested in the RealCause benchmark (Neal et al. 2020) to ensure that it is representative of real-world data.
• STAR Project and Hillstrom: Constructed by subsampling real-world randomised controlled trials, following the method described in Gentzel et al. (2021). We obtain the counterfactual outcomes by fitting a 'ground truth CATE model' on the full (significantly larger), randomised version of the datasets, following standard practices (Künzel et al. 2018).

We included more detailed description for how these datasets are constructed in Appendix D.2 in our manuscript.

Why are these techniques necessary? The fundamental challenge of evaluating CATE models lies in the fact that in real-world datasets, the ground truth value of the treatment effect is never observed. This makes the objective evaluation of CATE model performance on a held-out test set impossible when using real-world observational data. To counteract this problem, many different evaluation methods for constructing synthetic or semi-synthetic datasets have been proposed, however all of them have their respective flaws (Curth et al., 2021b).

What alternative datasets are used in the literature? CATE models are typically evaluated using fully synthetic datasets, where covariates, treatments, and outcomes are all simulated (e.g., Curth & van der Schaar, 2021a; Wager & Athey, 2018), or semi-synthetic datasets like IHDP (Hill et al., 2016), which use real-world covariates but synthetically generated outcomes, using pre-defined parametric functions. Many causal inference papers rely exclusively on such datasets (e.g., Shi et al., 2018; Wager & Athey, 2018; Hassenpour & Greiner, 2018; Jesson et al., 2021), even though they may fail to capture the complexity of real-world problems.

However, semi-synthetic datasets are particularly unsuitable for evaluating the $`GATE`$ framework, as their outcomes lack grounding in reality, preventing a proper assessment of the LLM's ability to incorporate external knowledge. For this reason, we focused on more reality-grounded datasets and evaluation practices. Nevertheless, for completeness, we also tested LLM-instantiated $`GATE`$ on the IHDP dataset in Appendix F.2 to evaluate its few-shot learning capabilities.

We appreciate the reviewer’s detailed and thoughtful feedback.

Results for more recent CATE learners

We agree with the reviewer that evaluating the benefits of $`GATE`$ across various CATE learners is important. However, let us clarify that in addition to the results presented in Table 1 (with 7 CATE learners), our original manuscript included results for TARNet (Shalit et al., 2017), DragonNet (Shi et al., 2019), CFR-MMD (Shalit et al., 2017) and BART (Athey & Imbens, 2016), presented in Appendix F.1. of our original manuscript. We realized that there was a wrong reference to this Appendix in the main text and apologize for this error.

For clarity, we present these results below. We note the following: (1) $`GATE`$ instantiated with LLMs leads to significant performance improvements across most datasets and models, (2) it allows to decrease the performance gap between the different CATE models, making it a great data-preprocessing step which could potentially minimise the effect of model selection.

Supplementary results for Table 1:

Action Taken: Following your comment, we have moved the results on the additional learners into the main text, adding them to Table 1.

Details of how to utilise the LLM for the $`GATE`$ framework

Details on LLM instantiation. Following your comment, we have dedicated Section 4.1 in the revised version of the manuscript to detailing what strategies we use to fully leverage the LLM's capabilities within the $`GATE`$ framework. This includes a discussion about (1) the metadata-informed prompts, (2) in-context samples provided, as well as (3) methods for aggregating the potential outcomes sampled from the LLM. We have also reeemphasised that the details and templates of the prompts used can be found in the appendix, to ensure the reproducibiity of our results.

Empirical validation. We investigate the effect of these prompting strategies on the downstream performance of the CATE models in Section 6.4. The results indicate that informing the LLM about the context of the task (Figure 2), as well as providing it with in-context samples from $\mathcal{D}^{(obs)}$ (Figure 4) helps to significantly improve the downstream performance.

Action Taken: We have revised Section 4.1 to detail how to utilise the LLM within the $`GATE`$ framework.

Connection between theoretical analysis and design choices for the $`GATE`$ framework

We would like to highlight the following connections between the theoretical analysis (Section 3) and the design choices made for the $`GATE`$ framework (Section 4):

1. Theoretical Insight: Tuning the distribution $Q$ via the admissible set $\mathcal{X}_t$ allows to navigate the trade-off between the bias introduced by the generator, and the reduction of variance and covariate shift achieved via data augmentation. In other words, performing data augmentation only in a selected subset of the covariate space $\mathcal{X}_t \subset \mathcal{X}$ can allow to obtain performance benefits even when the generative model is imperfect.

Design Choice: We generate the missing potential outcomes only in selected a subset of the covariate space. The exact proportion of potential outcomes generated is modulated by the parameter $\alpha$, and we instantiate $`GATE`$ with $\alpha = 0.5$ for our experiments (unless otherwise stated).

Empirical Verification Through Experiment 6.5.2 (Figure 5, middle): We verify that modulating $\mathcal{X}_t$ allows to navigate this trade-off.

References

Athey, Susan, and Guido Imbens. "Recursive partitioning for heterogeneous causal effects." Proceedings of the National Academy of Sciences 113.27 (2016): 7353-7360.

Curth, Alicia, and Mihaela Van der Schaar. "Nonparametric estimation of heterogeneous treatment effects: From theory to learning algorithms." International Conference on Artificial Intelligence and Statistics. PMLR, 2021a.

Curth, Alicia, et al. "Really doing great at estimating CATE? a critical look at ML benchmarking practices in treatment effect estimation." Thirty-fifth conference on neural information processing systems datasets and benchmarks track (round 2). 2021b.

Dorie, Vincent, et al. "Automated versus do-it-yourself methods for causal inference: Lessons learned from a data analysis competition." (2019): 43-68.

Gentzel, Amanda M., Purva Pruthi, and David Jensen. "How and why to use experimental data to evaluate methods for observational causal inference." International Conference on Machine Learning. PMLR, 2021.

Hassanpour, Negar, and Russell Greiner. "CounterFactual Regression with Importance Sampling Weights." IJCAI. 2019.

Hill, Jennifer L. "Bayesian nonparametric modeling for causal inference." Journal of Computational and Graphical Statistics 20.1 (2011): 217-240.

Jesson, Andrew, et al. "Causal-bald: Deep bayesian active learning of outcomes to infer treatment-effects from observational data." Advances in Neural Information Processing Systems 34 (2021): 30465-30478.

Johansson, Fredrik, Uri Shalit, and David Sontag. "Learning representations for counterfactual inference." International conference on machine learning. PMLR, 2016.

Neal, Brady, Chin-Wei Huang, and Sunand Raghupathi. "Realcause: Realistic causal inference benchmarking." arXiv preprint arXiv:2011.15007 (2020).

Shalit, Uri, Fredrik D. Johansson, and David Sontag. "Estimating individual treatment effect: generalization bounds and algorithms." International conference on machine learning. PMLR, 2017.

Shi, Claudia, David Blei, and Victor Veitch. "Adapting neural networks for the estimation of treatment effects." Advances in neural information processing systems 32 (2019).

Wager, Stefan, and Susan Athey. "Estimation and inference of heterogeneous treatment effects using random forests." Journal of the American Statistical Association 113.523 (2018): 1228-1242.

Significance of improvement when using data augmentation

We verify that LLM-instantiated $`GATE`$ offers significant performance improvements compared to the no-augmentation baseline by conducting two-sample t-tests on the $\sqrt{\epsilon_{PEHE}}$ obtained with and without data augmentation (for the complete datasets, i.e $\rho=1.0$). We report the p-values in the table below, which we also included in Appendix G.1 of the revised manuscript.

These results demonstrate that performance gains obtained with $`GATE`$ are statistically significant at the $0.05$ level across the majority of the CATE learners and datasets.

Action Taken: To further emphasize the significance of our results, we have included the above table in Appendix G.1 of the revised manuscript.

Other questions

Training of generative models. Regarding the baseline generative models trained on $\mathcal{D}^{(obs)}$, we have included the details of how these were trained and set up in Appendix D.4. In particular, for the mean imputation model, we set $Y^{(gen)}(t) = \frac{1}{n_t}\sum_{i=1}^n Y_i \mathbb{1} (T_i = t)$, where $(X_i, T_i, Y_i) \in D^{(obs)}_t$ (we compute this mean over the training set only).

Scoring function. In the instantiation of $`GATE`$ we use a percentile-based definition of the scoring function $s(x, t)$ (where choosing $\alpha = 0.5$ guarantees that 50% of missing potential outcomes are generated, thus allowing to fix the proportion of generated outcomes across datasets. The $`GATE`$ framework also allows the practitioner to use fixed variance thresholds instead, e.g. defining ($\alpha_{0}$, $\alpha_{1}$) for the control and treated group respectively, and setting $\mathcal{X}_{t} = \\{X_i \mid i \in [n], s(X_i, T_i) < \alpha_t\\}$.

The values of $\alpha_{0}$ and $\alpha_{1}$ could be guided by domain-knowledge or exploratory analysis of the data. Using a fixed threshold would more explicitly guardrail against the inclusion in the augmented dataset of particularly 'poor' generated outcomes, as the proportion of generated outcomes would depend on the properties of the generative model. In particular, if the model is particularly bad, no potential outcomes are generated and our method recovers the baseline performance.

Actions taken: We have included a more detailed discussion about the threshold for the scoring function in Appendix B.3.

We hope the reviewer’s concerns are addressed and they will consider updating their score. We welcome further discussions.