AutoCATE: End-to-End, Automated Treatment Effect Estimation

Thank you for your thorough review!

Related work

While prior work automates parts of CATE estimation, no approach–to our knowledge–provides an end-to-end automated framework. Existing work only automates part of our approach (e.g., learning pseudo-outcomes with AutoML). In contrast, AutoCATE extends automation across the entire pipeline, integrating model selection, preprocessing, and ensembling. In doing so, we address novel challenges like balancing data for training and validation or using multi-objective optimization. Our empirical study is also considerably more extensive, covering diverse datasets and revealing new insights into model selection.

We will revise our paper to better highlight these distinctions.

Metalearners

AutoCATE is a general-purpose framework for CATE estimation, not limited to metalearners. While we use them for their flexibility, our key contribution is automating selection, tuning, and validation—critical for all CATE methods. Other approaches (e.g. TARNet) can easily be integrated in this framework. In future versions, we aim to support custom models through a user-friendly API.

Similar to general AutoML frameworks (e.g. FLAML), which do not incorporate each supervised ML algorithm, AutoCATE focuses on structured automation within an extensible search space. This remains a fundamental challenge, regardless of specific estimators.

Clarity: AllMeta and AllBase

These terms define search spaces in our automated ML pipeline for CATE validation and estimation. AllMeta includes all metalearners in our search (S, T, DR, X, R, RA, Lo, Z, U, F), covering all known techniques in the literature. AllBase consists of the nine base learners in our framework (see Figure 1). BestMeta and BestBase are selected subsets based on performance—metalearners are chosen from experiments and the baselearners are known to work well for tabular data.

We will clarify these in the revised paper—thank you for pointing this out! If anything else is unclear, we would be happy to clarify further.

AutoCATE stages

In Stage 1, AutoCATE constructs pseudo-outcomes as proxies for ground truth CATE. These are derived from metalearners, but not full CATE learners themselves: e.g., we use the DR-Learner’s pseudo-outcome directly, but without fitting the final model to predict it. This step is also performed on validation data, not training data.

Stage 2 trains a CATE estimator on training data and evaluates it using pseudo-outcomes from Stage 1. As evaluation depends on a risk measure, it logically precedes estimation, avoiding a chicken-and-egg problem. Both stages involve choices not dictated by data (e.g., which risk measure to use), which we study systematically.

Repeating stages until convergence is not straightforward, as estimation and evaluation are independent stages with distinct data splits. While evaluation informs estimator selection, estimators do not directly influence evaluation. Nevertheless, exploring these interactions (e.g., by sharing information in their optimization) could be an interesting avenue for future work.

Limited data

We agree that limited data can be a challenge for CATE validation. We emphasize the importance of having sufficient data for validation to mitigate such issues (as highlighted in Figure 3). We also study cross-validation (see Figure 7 and Table 5), which can help to stabilize the estimates for smaller datasets. Empirically, we observe good results for the IHDP data with only 672 instances. Some baselearners (e.g. linear regression) can perform reasonably well with very little data.

That said, we acknowledge that for very small datasets, other approaches may be more appropriate. We will include a more detailed discussion on scenarios where our framework may not be optimal in the conclusion. We refer to the response to reviewer NSJK, where we enlist more limitations.

Computational complexity

Efficiency is an important consideration in our framework. AutoCATE runs fairly quickly on small to moderate-size datasets, often completing in minutes locally (see Table 9). AutoCATE also allows for trade-offs by limiting search space, reducing trials, or choosing faster learners (see Figure 9). Further improvements could be made based on more efficient search and pruning algorithms. Nevertheless, we agree that computational complexity is a potential limitation and will discuss it more in depth in the revision.

Multiple risk measures

We provide an initial exploration of combining risk measures. While no combination consistently outperforms the best single measure, we see significant potential for further development (see Figure 13e). Though not yet optimal, it provides a solid foundation for future improvements in both performance and efficiency. Although this approach can increase computational complexity, reusing components like propensity scores helps mitigate this impact.

Thank you once again! Please let us know if there are any remaining concerns.

Thank you for your insightful review–this is highly appreciated!

Real-world data and violations

We agree that these are crucial considerations. While we use real-world data to validate AutoCATE as much as possible (e.g, Twins and the uplift data in Appendix D.4), we acknowledge that more real-world validation would be useful. Unfortunately, we are unfortunately limited by the fact that the CATE is unknown in real data.

We agree that considering violations of identifiability assumptions is important. We have added a synthetic experiment varying selection bias (γ) and found that while performance declines with higher bias, more optimization trials can partially mitigate the effect (we refer to the response to Reviewer NSJK). Even with strong overlap violations (γ > 10), AutoCATE remains competitive. We will update the paper to clarify the data-generating process and these findings. Finally, we also point out AutoCATE’s good performance on IHDP, which also contains overlap violations [1].

More generally, we recognize the need for improving robustness to such violations and will include this limitation in the revised version.

Domain knowledge and observational data diagnostics

Thank you for the insightful suggestion. We agree that incorporating domain knowledge and data diagnostics would be valuable, and we would love to support such features in future releases. AutoCATE allows for trimming extreme pseudo-outcomes (e.g., from very small propensity scores), though this was not included in our paper. We also agree that checking overlap/covariate balance and filtering propensity scores would be useful additions.

While sensitivity analyses for hidden confounders are appealing, they may fall outside the scope of our automated approach. Instead, we aim for AutoCATE to be complementary to other packages (e.g. DoWhy), which address other aspects of causal inference. Our main goal is automating CATE estimation, a gap we see in current tools. Nevertheless, we see many remaining challenges for truly supporting practical adoption of these methods.

Computational complexity and heuristics

We have designed AutoCATE with efficiency in mind, using parallelization via Optuna, efficient ML implementations via scikit-learn, and configurable constraints. Computational times for different datasets are summarized in Table 9.

Nevertheless, there are several opportunities to improve AutoCATE’s speed. We currently use a naive random search algorithm without heuristics. As our search is implemented with Optuna, more advanced search algorithms and pruning strategies can easily be used instead. Finally, different metalearners have widely varying time complexities (Figure 9). Future research could try to use these discrepancies to further optimize the search.

In the revised version, we will stress the complexity of AutoCATE more explicitly as a limitation.

Combining risk measures

Thank you for these insightful suggestions!

We explore various static strategies for combining risk measures (e.g., averaging, ranking) and compare them empirically in Table 3. While we analyse correlations between risk measures (Figure 13a), we do not yet use this information, though it seems promising. Our current work focuses on the feasibility of combining risk measures using simple approaches. While the suggested dynamic approaches are really exciting directions, their implementation is challenging–key issues include the lack of ground truth and the variance of risk measures. While tackling these challenges is beyond the scope of this work, we see them as important future research directions and hope our framework can serve as a foundation for further advancements.

Stacking

Why might stacking underperform? Without a ground truth CATE, validating and tuning stacking weights is challenging. Stacking may overfit pseudo-outcomes with high variance or outliers. The squared error loss is sensitive to outliers, and alternatives like Huber loss could improve robustness. Stacking also requires more training data, and with limited data, its complexity may hurt generalization.

To improve stacking, we could try to filter pseudo-outcomes based on reliability, possibly using risk measure agreement. Multi-objective stacking could create more generalizable models. Using loss functions like Huber loss could mitigate noisy outcomes. Further research into tailored stacking and ensembling for CATE estimation is needed. Limiting stacking to top models could also improve stability.

Thank you again for your detailed and thoughtful review. We agree that there is much work ahead, and while some suggestions are beyond the scope of this paper, we are happy to consider any specific approaches you feel should be included for acceptance within the remaining time.

[1] Curth, A., Svensson, D., Weatherall, J., & van der Schaar, M. (2021). Really doing great at estimating CATE? A critical look at ML benchmarking practices in treatment effect estimation. NeurIPS.

Thank you for your time and effort in reviewing our work!

Obtaining the best meta/base-learner and training nuisance models with AutoML

Please allow us to clarify our approach. AutoCATE follows an AutoML-based procedure to tune ML pipelines at multiple stages: first, to optimize risk measures for model selection, and second, to construct optimal CATE estimators by searching over meta- and base-learners. We obtain the best CATE estimator–i.e., a metalearner constructed with (multiple) baselearner(s)--in the second stage, based on the risk measure selected in the first stage. As the reviewer suggests, this indeed ensures that we pick the CATE estimator that best predicts the (observational) validation data.

Our design ensures a fully automated, end-to-end approach that integrates evaluation, estimation, and ensembling within a single framework (see Section 4.5). Importantly, all components of a metalearner (i.e., each baselearner) are tuned automatically and simultaneously, ensuring the best possible downstream performance. We believe that this is one of the key strengths of our work compared to previous approaches!

We hope that these additions clarify our approach. We will update the paper to make this aspect more clear. Nevertheless, if there are any remaining questions, please do let us know.

Tuning the S- and T-Learner

Regarding the tuning of S-/T-Learners, this is done based on a random search with the number of trails equal to AutoCATE, allowing for a fair comparison. As such, we do not use a manual grid search for either AutoCATE or the benchmarks. Mahajan et al. (2023) use AutoML separately for evaluation and for learning the nuisance models underlying CATE estimators. As neither S- and T-Learners use nuisance models, we believe that our approach is similar to the one in Mahajan et al. (2023). Additionally, they do not consider the impact of only using a limited number of optimization trials, which is an important consideration in practice. Nevertheless, if there are distinctions we have missed, we would be happy to look into them and adjust the experiments.

We will revise our manuscript to better clarify the reasoning behind and training procedure for the benchmarks.

Combining risk measures

Thank you for these thoughtful questions. We appreciate the opportunity to clarify our approach and conclusions.

How are risk measures combined?

We explore multiple strategies for combining risk measures (see also Appendix B.5), including:

• Averaging normalized risk measures (as in Table 1b). This approach corresponds to the reviewer’s suggestion.
• Averaging rankings of risk measures to improve robustness to outliers.
• Euclidean distance to the origin (best possible performance).
• Selecting all Pareto-optimal points.

What are the main conclusions?

Our experiments serve as an initial exploration into combining risk measures. While our hypothesis was that relying on multiple risk measures could enhance robustness and performance, we find that no strategy consistently outperforms using a single T-risk. However, results indicate that this approach is promising, and further research—such as leveraging risk measure correlations (Figure 14a) or advanced ensembling—could improve performance. Therefore, we believe there is evidence that this novel approach provides a fruitful and promising direction for future research on CATE estimator validation. By highlighting the potential of this approach and providing a foundation for future research with our software package, we hope to encourage the community to further explore these ideas.

We will update our paper to clearly state the main conclusions.

Scale of the empirical study

Thank you for this great suggestion! We agree that summarizing the scale of our empirical study upfront helps highlight our contributions, and we will incorporate this more expliclity in the revised manuscript.

In the main body, we present experiments across a total of 247 distinct datasets, across four benchmark families, spanning binary and continuous outcomes, as well as different sizes and dimensionalities. Additionally, we include additional experiments on synthetic data for this rebuttal. In the appendix, we also explore AutoCATE on two uplift datasets.

AutoCATE’s full search space consists of 2,187 possible pipelines (3 feature selection × 3 scaling × 27 meta/base-learner configurations × 9 base learners), excluding hyperparameters (Appendix B.3). It incorporates 8 different risk measures. A full overview of its configuration options can be found in Appendix B.6.

We appreciate the reviewer’s suggestion and will ensure this information is presented more explicitly in the paper.

Once again, thank you for your detailed review! While we hope that these responses address your concerns, please let us know if there are any remaining points. We would be happy to engage further if needed.

Thank you for your thoughtful review!

Robustness to selection bias

We agree that synthetic data allows for precise control over selection bias and covariate shift, enabling a more systematic evaluation. We have added a synthetic experiment where we vary the degree of selection bias (controlled by the parameter γ\gammaγ). Our results indicate that while AutoCATE’s performance degrades as selection bias increases, increasing the number of search trials helps mitigate this effect. Even under strong overlap violations ($\gamma > 10$), AutoCATE can result in good performance. We also compare AutoCATE to benchmark models across different bias levels. The results confirm that AutoCATE consistently performs competitive to each baseline in settings with moderate bias and remains relatively robust under extreme bias. We will update the paper to include a detailed explanation of the data-generating process (DGP) and expand on these findings. Thank you for this valuable suggestion—we believe this addition strengthens our evaluation and provides clearer insights into AutoCATE’s generalization capabilities.

Synthetic data: setup

Synthetic data: $\sqrt{\text{PEHE}}$ (SE) for AutoCATE with different number of evaluation and estimation trials

Synthetic data: $\sqrt{\text{PEHE}}$ (SE) for AutoCATE and benchmarks with 50 evaluation and estimation trials

Customizability

AutoCATE is designed to be highly customizable yet easy to use. Users can specify the search space (preprocessors, metalearners, baselearners) and set design parameters (evaluation protocols, trial numbers, ensemble methods). Domain experts can also fine-tune risk measures, evaluation metrics, and validation procedures. The interface follows scikit-learn conventions for intuitive use, with configuration details available in Appendix B.6. Future versions will include an API for custom algorithms, enabling users to integrate their own CATE estimation and evaluation methods.

Limitations

While AutoCATE is designed to be broadly applicable, we acknowledge that no method is universally optimal for all CATE estimation scenarios. As such, there are certain settings where AutoCATE may be less suitable:

• Very small datasets (n<50), where model selection based on pseudo-outcomes may be unreliable. For IHDP, we achieve good performance with only n=672 instances in the training set.
• Large datasets with constrained compute, where an AutoML-based approach may be too computationally expensive.
• Scenarios requiring strong domain knowledge integration, where a fully customized pipeline may be preferable.
• Data that requires extensive preprocessing, such as raw image or text data.
• Settings with fairness or regulatory constraints, where automated model selection may need additional safeguards.
• Cases violating our causal assumptions, such as strong violations of overlap, hidden confounders.
• Cases considering a different setting, such as instrumental variables or time series/panel data settings.

We agree that clarifying these limitations will strengthen the paper and will update the manuscript accordingly. Thank you for this valuable suggestion. If the reviewer thinks of other limitations that we have missed, we would be happy to also include those.

Theoretical contributions

While our paper does indeed not present new formal proofs or theoretical guarantees, it aligns with ICML's emphasis on application-driven research, as highlighted in the call for papers. ICML encourages innovative techniques and problems motivated by real-world needs, with an emphasis on reproducible experiments and sound analysis, rather than mandatory theoretical components. Many influential ICML papers prioritize empirical insights and conceptual innovations, and we believe our work contributes to this tradition by addressing a significant practical challenge in causal inference.

Thank you for your time and effort in reviewing our work! Please let us know if you have any remaining concerns.