Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
52,884 characters · 16 sections · 53 citation commands
In Search of Insights, Not Magic Bullets: Towards Demystification of the Model Selection Dilemma in Heterogeneous Treatment Effect Estimation
\twocolumn[ \icmltitle{In Search of Insights, Not Magic Bullets: Towards Demystification of the Model Selection Dilemma in Heterogeneous Treatment Effect Estimation}
\icmlsetsymbol{equal}{*}
\icmlaffiliation{yyy}{Department of Applied Mathematics and Theoretical Physics, University of Cambridge, UK} \icmlaffiliation{comp}{The Alan Turing Institute}
\icmlcorrespondingauthor{Alicia Curth}{[email removed]}
\icmlkeywords{Heterogeneous Treatment Effect, CATE, ITE, Model Selection}
\vskip 0.3in ]
\printAffiliationsAndNotice
Applications in which the causal effects of treatments (or actions, interventions & policies) are of interest are ubiquitous in empirical science, and personalized effect estimates could ultimately be used to improve decision making in many domains, including healthcare, economics and marketing. Machine learning (ML) has shown great promise in providing such personalized effect estimates bica2021real, and the ML literature on the topic has matured over the last five years: a plethora of new methods for estimating conditional average treatment effects (CATE) have been proposed recently, including method-agnostic so-called meta-learner strategies that could be implemented using any ML prediction method kunzel2019metalearners, nie2021quasi, kennedy2020optimal, curth2021nonparametric as well as adaptations of specific ML methods shalit2017estimating, alaa2018limits, wager2018estimation to the treatment effect estimation context.
Personalized treatment effect estimates, crucially, are often of interest in safety-critical applications, particularly in medicine and policy making. Thus, prior to use in practice we would like to ensure that we have selected a ML estimator from this vast toolbox that is trustworthy and outputs the best possible estimates. While this sounds like a straightforward requirement, it remains a big hurdle in practice: because of the fundamental problem of causal inference holland1986statistics, ground truth treatment effects are usually not available to perform standard model validation, and alternative solutions need to be considered to overcome this model selection dilemma. As we discuss below, despite recent proposals of new model selection strategies and recent empirical studies comparing different strategies, we believe that there is still a general lack of understanding of the (dis)advantages of different strategies and how they are entangled with underlying data-generating processes (DGPs) -- which we aim to provide in this work.
Related work. Despite its practical relevance, the problem of model selection for heterogeneous treatment effect estimation has received only very limited attention so far (this stands also in stark contrast to the plethora of new estimators proposed in recent years, see Appendix (ref) for an overview). An intuitive and often applied solution is to rely on a simple prediction-type validation and evaluate a model's performance in predicting the observed (factual) outcome associated with the factual treatment instead of evaluating the quality of the effect estimate. More targeted alternatives have recently been developed: rolling2014model propose to construct approximate effect validation targets by matching the nearest treated and control units and comparing their outcomes, nie2021quasi highlight that their R-learner objective could also be used for model selection, saito2020counterfactual similarly propose the use of a criterion that corresponds to kennedy2020optimal's meta-learner objective and alaa2019validating propose to rely on influence functions to de-biase plugin estimates. We are aware of two independent benchmarking studies that compare (subsets of) such criteria: schuler2018comparison find that the R-learner objective performs best overall, while mahajan2022empirical find that no criterion dominates all others over all datasets considered (and in particular do not find the R-learner objective to perform remarkably), which highlights to us that there is much room for understanding of the relative strengths but also relative weaknesses of different selection criteria.
Contributions. In this paper, we focus on building systematic understanding of the (dis)advantages of different model selection criteria. Note that therefore our aim is not to propose new methodology, but rather to establish understanding and insight into the tools already available in the literature. We believe that this is one of the most crucial and necessary next steps for this community in order to enable actual adoption of personalized treatment effect estimators in practice, and may inspire further methodological research to fill gaps highlighted by this understanding. In doing so, we make three contributions:
We consider the by now standard CATE estimation setup within the potential outcomes framework rubin2005causal. That is, we assume access to a dataset consisting of $n$ i.i.d. tuples $(X, A, Y)$, where $Y$ is an outcome of interest, $X$ consists of pre-treatment covariates and $A \in \{0, 1\}$ is a binary treatment (action, intervention or policy), which is assigned according to some (often unknown) propensity $\pi(x)=\mathbb{P}(A=1|X=x)$. We assume that conceptually each individual is a priori associated with two potential outcomes (POs) $Y(a), a \in \{0, 1\}$, capturing outcome under either treatment $a$, however, we observe only the outcome associated to the treatment $A$ actually received, i.e. $Y=Y(A)$, We can thus naturally define an individualized treatment effect through the (unobserved) PO contrast $Y(1)\!-\!Y(0)$. We focus on estimating the conditional average treatment effect (CATE) $\tau(x)$, i.e. the expected PO difference for an individual with covariates $X=x$:
where $\mu_a(x)=\mathbb{E}[Y(a)|X=x]$. To ensure that effects are identifiable and nonparametrically estimable from observational data, we rely on the standard ignorability assumptions rosenbaum1983central:
Then, $\mathbb{E}[Y(a)|X\!=\!x]=\mathbb{E}[Y|X\!=\!x, A\!=\!a]$ so that observed statistical associations have a causal interpretation.
A plethora of strategies for estimating CATE have been proposed in the recent literature. One strand of this work has recently relied on the meta-learner framework of kunzel2019metalearners, where a meta-learner provides a `recipe' for estimating CATE using any arbitrary\footnote{as opposed to adaptation of specific ML methods proposed in another stream of work, as discussed further in Appendix (ref).} ML method $\mathcal{M}$. Due to their ease of implementation with different underlying ML methods, existing theoretical understanding and correspondence to the model selection strategies discussed in the following section, we focus on the problem of choosing between such meta-learners in this paper.
Following curth2021nonparametric we distinguish between (i) indirect estimation strategies, which estimate CATE indirectly by outputting estimates $\hat{\mu}_a(x)$ of the PO regressions and then setting $\hat{\tau}=\hat{\mu}_1(x)-\hat{\mu}_0(x)$, and (ii) direct estimation strategies which output an estimate $\hat{\tau}(x)$ directly without outputting PO estimates as byproducts. kunzel2019metalearners discuss two indirect learners: a T-learner strategy, where the training data is split by treatment group and $\mathcal{M}$ is trained independently (twice) on each sample to output two regressors $\{\hat{\mu}_0(x), \hat{\mu}_1(x)\}$, and a S-learner strategy, where the treatment indicator $A$ is simply appended to $X$ so that $\mathcal{M}$ can be trained a \textbf{s}ingle time using covariates $(X, A)$ and outputting a single estimated function $\hat{\mu}(x, a)$ that can be used to impute both POs. Because the latter formulation can lead to implicit regularization of CATE schuler2018comparison (as any heterogeneous effect has to be represented by \textit{learned interaction terms} of $X$ and $A$), we also include an extended version (ES-learner) which is trained on the covariates $(X, X*A, A)$ explicitly.
Alternatively, there exist direct estimators that output an estimate of the CATE only by relying on estimates of (some of) the nuisance parameters $\eta = (\mu_0(x), \mu_1(x), \pi(x))$ obtained in a pre-processing step using ML method $\mathcal{M}$. Most such strategies, in particular, X-learner kunzel2019metalearners, the DR-learner kennedy2020optimal and RA- and PW-learner curth2021nonparametric, rely on a pseudo-outcome approach where, using plug-in nuisance estimates $\hat{\eta}$, one constructs a pseudo outcome for which it holds that $\mathbb{E}[Y_{\eta}|X=x]=\tau(x)$ for ground truth nuisance parameters ${\eta}$, and then regresses $Y_{\hat{\eta}}$ on $X$ using $\mathcal{M}$ to give an estimate $\hat{\tau}(x)$. nie2021quasi's R-learner is similar in spirit but relies on a modified loss function instead. These multi-stage estimation procedures have recently gained popularity in the literature because they have good theoretical properties curth2021nonparametric, are more robust kennedy2020optimal and have been observed to perform better across a variety of scenarios than vanilla S- and T-learner individually nie2021quasi.
We study the problem of selecting an estimator from a set of CATE estimators $\mathcal{T}\!=\!\{\hat{\tau}_1(\cdot), \ldots, \hat{\tau}_K(\cdot)\}$, containing different meta-learner+ML method combinations, that minimizes the precision of estimating heterogeneous effects (PEHE) hill2011bayesian, the root-mean-squared error of estimating the underlying effect function over a test-set of size $n$:
This metric is an oracle metric because it cannot be evaluated in practice -- making the CATE model selection problem highly nontrivial. Below, we therefore provide an in-depth discussion of the inherent challenges of this problem, and then establish a classification of model selection criteria that have been proposed to overcome these challenges.
{Challenge 1: Lack of supervised signal for the individual treatment effect. Due to the fundamental problem of causal inference holland1986statistics, i.e. the fact that we can only either observe $Y(0)$ or $Y(1)$ for any one individual, the true supervised target label $Y(1)-Y(0)$ for estimation of $\mathbb{E}[Y(1)-Y(0)|X=x]=\tau(x)$ is not available for model selection through a standard held-out validation approach. This lack of supervised label is also the issue that motivated the construction of the direct meta-learners for estimation of effects using pseudo-outcomes kennedy2020optimal, curth2021nonparametric, but it results in another challenge for model selection: even though direct estimation of effects is possible, it is not possible to compare multiple direct estimators on basis of their output $\hat{\tau}_k(x)$ using observed data only because there is no natural outcome to validate $\hat{\tau}_k(x)$ against directly.\parfillskip=0pt}
Challenge 2: Confounding leads to covariate shift between treatment groups. One straightforward option to validate estimators using factual (observed) outcomes only is to simply evaluate them based on their outcome prediction ability; i.e. to use $Y_i - \hat{\mu}_{A_i}(X_i)$ for validation. This, however, is only an option for evaluating indirect learners because it requires an output $\hat{\mu}_a(x)$. Even when one is willing to restrict attention only to indirect learners to make use of a factual evaluation strategy, a remaining inherent challenge is that evaluating $\hat{\mu}_a(x)$ only on individuals with $A=a$ observed inherently suffers from covariate shift whenever $\pi(x)$ is not constant because treatment is not assigned completely randomly. This too is a challenge also when estimating effects shalit2017estimating.
Challenge 3: Selection for good PO estimation and CATE estimation may not be the same. Finally, even when selecting only among indirect learners and in absence of covariate shift, in finite samples the estimator with the best performance on estimating (potential) outcomes might not do best at estimating CATE. It is clear that when models are correctly specified and unlimited data is available, perfectly estimating the POs will immediately lead to perfect CATE estimates. However, when data is limited and/or the model is misspecified, this might lead to a trade-off between estimating the POs well and estimating their difference\footnote{To see this, note that when comparing an indirect estimator $\hat{\mu}^1_a(x)=\{\mu_1(x) + \tilde{\epsilon}_1(x), \mu_1(x) - \tilde{\epsilon}_1(x)\}$ to an estimator $\hat{\mu}^2_a(x)=\{\mu_1(x) + \tilde{\epsilon}_2(x), \mu_1(x) + \tilde{\epsilon}_2(x)\}$ with estimation errors $\tilde{\epsilon}_1(x) < \tilde{\epsilon}_2(x)$ for all $x$, the MSE of estimating the POs will be lower for estimator $\hat{\mu}^1_a(x)$ -- yet its estimation error on CATE will be $2 \tilde{\epsilon}_1(x)$ for every $x$ while estimator $\hat{\mu}^2_a(x)$ with worse estimation on the POs will have CATE estimation error $0$. }.
To overcome the challenges discussed above, numerous alternative model selection criteria have been used or proposed in related work. Below, we establish a classification of existing model selection criteria into three categories based on their most salient characteristics: we consider factual (prediction) criteria, plug-in surrogate criteria and pseudo-outcome surrogate criteria. We provide a conceptual overview of the three types of strategies in Fig. (ref). Overall, we consider a similar set of model selection criteria as the (union of) the benchmarking studies presented in schuler2018comparison, mahajan2022empirical\footnote{schuler2018comparison, predating publication of most model selection papers, miss some of the plug-in and pseudo-outcome criteria. mahajan2022empirical only do not consider factual (prediction) criteria, which we consider of major interest as discussed in the following section. We only drop policy value criteria, i.e. those optimizing the derived treatment policy $\mathbf{1}\{\hat{\tau}_k(x)>0\}$, from consideration, both because we focus on PEHE and because these criteria were shown to underperform even when evaluated in terms of policy value schuler2018comparison.}.
1. Factual (prediction) criteria. First, as discussed in the previous section, a possible way of evaluating models $\hat{\tau}_k(x)$ that also output a pair of regressors $\{\hat{\mu}^k_0(x), \hat{\mu}^k_1(x)\}$ is to rely on a simple prediction loss considering only the observed potential outcome
In order to correct for possible covariate shift, this can also be transformed into an importance weighted criterion $\mathcal{E}^{fact, w}_{Y}$ using a propensity score estimate in $w(X_i, A_i) = A_i(\hat{\pi}(X_i))^{-1}+(1-A_i)(1-\hat{\pi}(X_i))^{-1}$.
2. Plug-in surrogate criteria. To actually evaluate estimates $\hat{\tau}_k(x)$ directly one is thus forced to construct surrogates for CATE. One way of doing so is by fitting a new CATE estimator $\tilde{\tau}(x)$ on held-out data and using this to compare against the estimates:
For implementation of this criterion, any CATE estimator could be used -- thus trying to select the best plug-in estimator may potentially lead us in a circle and back to the problem we are originally trying to solve. Related work alaa2019validating, mahajan2022empirical considered only indirect estimators (S and T-learners) as a plug-in surrogate criterion -- possibly because it is possible to choose between those factually -- but we note that technically any estimator, including direct ones, could be used as $\tilde{\tau}(x)$.
3. Pseudo-outcome surrogate criteria. Finally, one could make use of the same pseudo-outcome based strategy that underlies the direct learners: given auxiliary nuisance estimates $\tilde{\eta} = (\tilde{\mu}_0(x), \tilde{\mu}_1(x), \tilde{\pi}(x))$ obtained from the validation data using ML method $\mathcal{M}$, one can construct pseudo-outcomes $Y_{\tilde{\eta}}$ for which it holds that for ground truth nuisance parameter $\eta$, $\mathbb{E}[Y_{\eta}|X=x]=\tau(x)$ and -- instead of using them as regression outcomes as in the learners themselves -- one can use them as validation targets in
which is reasonable because the conditional mean $\mathbb{E}[Y_{\tilde{\eta}}|X=x]$ minimizes the MSE $\mathcal{E}^{pseu}_{Y_{\tilde{\eta}}}(\cdot)^2$ in expectation. Use of the doubly robust pseudo-outcome,
which is what the direct meta-learner known as the DR-learner kennedy2020optimal is based on, gives rise to the selection criterion proposed in saito2020counterfactual. We note here that it would also be possible to use the other meta-learner pseudo-outcomes discussed in curth2021nonparametric for this purpose, e.g. the singly-robust propensity-weighted $\textstyle{{Y}_{PW, \tilde{\eta}} = \left(\frac{A}{\hat{\pi}(X)}- \frac{(1-A)}{1-\hat{\pi}(X)}\right) Y}$. We also consider rolling2014model's matching based model selection strategy to fall into this category. In this case $ \textstyle{Y_{\text{match}, \tilde{\eta}}=(2A_i-1)(Y_i - \tilde{NN}_{1-A_i}(X_i))}$ where $\tilde{NN}_a(X)$ is the nearest neighbor of $X$ in treatment group $a$; this essentially corresponds to the pseudo-outcome associated with the RA-learner of curth2021nonparametric, implemented using 1-NN regression to estimate the nuisance parameters $\hat{\mu}_a(x)$. We also put alaa2019validating's influence function based criterion, which we discuss in Appendix (ref), into this category. Finally, the R-learner objective of nie2021quasi, which requires an estimate of the treatment-unconditional mean $\mu(x)=\mathbb{E}[Y|X=x]$, relies on a similar idea\footnote{$\mathcal{E}^{pseu}_{R}$ can be rewritten in pseudo-outcome form as $\textstyle{{Y}_{R, \tilde{\eta}} = \frac{Y_i - \tilde{\mu}(X_i)}{A_i - \tilde{\pi}(X_i)}}$ if combined with weights $(A_i - \tilde{\pi}(X_i))^2$ to be used in the sum inside the RMSE knaus2021machine.} and can also been used for the selection task nie2021quasi, resulting in the criterion
Having established a high-level classification of model selection strategies, we can now build intuition and hypothesize about their expected (dis)advantages based on their inherent characteristics (Sec. (ref)). These expectations lead to numerous interesting research questions, which we believe can only be disentangled by designing carefully controlled experiments -- a feature that we would argue related work has neglected so far. Therefore, we then move to discuss design principles for construction of empirical studies of the CATE model selection problem (Sec. (ref)), which we will then apply in our experiment section.
Combining on our high-level overview of different strategies for CATE model selection with the previously discussed challenges, we argue that every class of criteria comes with their own inherent (dis)advantages.
Factual (prediction) criteria. We believe that factual criteria could be very appealing for use in practice because -- at least the unweighted $\mathcal{E}^{fact}_{Y}$ -- does not require estimates of any nuisance parameters and only relies on observed data; this means that there is no additional overhead and its results could be considered trustworthy in the sense that there is no dependence on possibly misspecified or biased nuisance estimates. However, such criteria mainly evaluate the performance in terms of estimation of the POs, which, as discussed above, may wrongly prioritize good fit on the POs over good CATE fit -- while the latter cannot be measured at all. This last point, crucially, also means that the factual criterion cannot evaluate all types of methods, excluding all direct estimators in particular, and therefore has to select among a smaller set $\mathcal{T}_{indirect}\subset \mathcal{T}$ of all possible estimators -- potentially missing out on the estimators with the best performance by construction.
{Plug-in surrogate criteria. Plug-in surrogate criteria have the clear advantage over factual criteria that they can evaluate all types of estimators and are targeted at the outcome of interest (i.e. the CATE). Yet, because a plug-in estimate $\tilde{\tau}(x)$ is needed, this introduces additional potential for estimation error or variance. Further, $\tilde{\tau}(x)$ could be any CATE estimator, thus choosing a good plugin $\tilde{\tau}(x)$ leads us back to the dilemma we were initially trying to overcome. Finally, we believe that such surrogate criteria may also suffer from a phenomenon we will refer to as congeniality bias: they may advantage CATE estimators $\hat{\tau}_k(x)$ that are structurally similar to their plug-in estimator $\tilde{\tau}(x)$. Even though $\hat{\tau}_k(x)$ and $\tilde{\tau}(x)$ should be fit on different data folds, we expect that a plug-in criterion $\tilde{\tau}(x)$ may still prefer estimators imbued with similar inductive biases. That is, we expect that e.g. a criterion using plug-in surrogate $\tilde{\tau}(x)$ implemented using linear regression may favor CATE estimators $\hat{\tau}_k(x)$ similarly relying on linear regressions (over estimators implemented using other methods $\mathcal{M}$), and one relying on an S-learner surrogate $\tilde{\tau}(x)$ may have a preference for selecting S-learners $\hat{\tau}_k(x)$. Here, we borrow the term `congeniality bias’ from the psychology literature, where it is used to indicate that individuals may have a systematic preference for information consistent with current beliefs hart2009feeling.\parfillskip=0pt}
Pseudo-outcome surrogate criteria. We expect these criteria to share the advantages, and some of the disadvantages of the plug-in criteria (namely the need to estimate additional parameters, and resulting possibility for increased error or variance). Because they do not use a final $\tau(x)$ estimate but only a pseudo-outcome $Y_{\tilde{\eta}}$, we expect that they might be less likely to suffer from congeniality bias, but could still favor estimators with similar inductive biases, e.g. direct estimators trained on the same pseudo-outcome.
This paper was motivated by many of the interesting research questions outlined below that naturally follow from the discussion above, none of which we believe have been addressed in related work-- we believe that this is a result of the fact that the focus so far has been on establishing global best performance of some criterion. Instead, we are interested in understanding scenarios in which there could be performance differences of competing criteria, in the hope that this will help practitioners in choosing the right criterion in their specific application. In particular, we are interested in exploring three questions in depth in this paper:
As outlined above, we believe that there are many very interesting questions to explore in this CATE model selection dilemma. Existing empirical studies that we are aware of -- both those proposing new model selection criteria alaa2019validating, saito2020counterfactual and those benchmarking existing ones schuler2018comparison, mahajan2022empirical -- have mainly focused on the question `what strategy works best globally?', and have taken a black-box approach in doing so: by considering opaque datasets for benchmarking, by considering a large inextricable set of estimators to select from and by mainly reporting on averages across a number of different DGPs.
It has recently been highlighted that the benchmarking practices in the ML CATE estimation literature more generally have many shortcomings curth2021really, especially the reliance on single semi-synthetic datasets\footnote{Note that, due to the lack of $Y(1)\!-\!Y(0)$ in real data it is generally necessary to simulate outcomes in experiments to allow for known ground truth; most existing benchmark datasets are semi-synthetic in that they use real covariates but simulate outcomes.} that encode very specific problem characteristics in their DGP, without discussing the effect of these choices. We therefore believe that it is crucial to carefully design controlled experimental environments that allow to disambiguate the effects that different components of a DGP may have on selection criterion performance. Below we discuss desiderata for designing an empirical study that allows for the insight into model selection performance we seek, which we use to design an empirical study in the following section.
1. Use DGPs with interesting `experimental knobs'. To gain systematic understanding of when different selection criteria work well, we need to be able to systematically vary the underlying experimental characteristics -- as suggested by dorie2019automated we therefore design simulations that enable turning of important `experimental knobs'. In pursuit of interesting insights into relative performance of selection criteria, we will therefore choose to investigate axes that have been shown to matter for estimator performance itself as well as others that we would expect to matter.
{2. Examine the performance of candidate estimators in $\mathcal{T}$.} To understand when different selectors perform well, we believe that it is important to first establish how the underlying candidate estimators in $\mathcal{T}$ perform. All related work that we are aware of skip this step, yet we consider it crucial because the performance of different selection strategies may be deeply entangled with the performance of the underlying estimators: if, for example, there is congeniality between selectors and estimators, and/or a specific type of estimator is advantaged on a specific dataset, then the corresponding selector may perform well by construction.
3. Analyze how and when performance of selection criteria varies. Finally, once DGPs are constructed and estimators examined, we aim to analyze the performance of the selection criteria in detail. That is, we wish to explicitly understand how the relative performance of selectors varies as an experimental knob is turned. In many related works this is not possible as results are reported as averages across many different settings alaa2019validating, mahajan2022empirical, obfuscating possible interesting insights into systematic performance differences.
Setup. In this section, we conduct an empirical study comparing CATE selection criteria following the three steps outlined above. Throughout, we rely on two ML-methods $\mathcal{M}$ to instantiate all meta-learners and selection criteria: Extreme Gradient Boosted Trees chen2016xgboost (GB) and linear regressions with ridge penalty (LR). We chose these two because they encode very different inductive biases, and allow us to give insights into performance differences between very flexible versus rigid models. Whenever propensity score estimates are needed, we estimate these using logistic regressions. As meta-learners to choose between, we consider (indirect) S-, T- and ES-learners and the (direct) DR- and R-learner. As selection criteria we consider $\mathcal{E}^{fact}_{Y}$ (factual), $\mathcal{E}^{plug, ES}_{\tilde{\tau}}$, $\mathcal{E}^{plug, T}_{\tilde{\tau}}$, $\mathcal{E}^{plug, DR}_{\tilde{\tau}}$ & $\mathcal{E}^{plug, R}_{\tilde{\tau}}$ (surrogate plug-in) and $\mathcal{E}^{pseudo, DR}_{Y_{\tilde{\eta}}}$ & $\mathcal{E}^{pseudo, R}_{Y_{\tilde{\eta}}}$ (surrogate pseudo-outcome) in the main text, further results using criteria that generally performed worse can be found in Appendix (ref). Implementation details\footnote{Code to replicate all experiments is available at \url{https://github.com/AliciaCurth/CATESelection}} can be found in Appendix (ref).
We build a DGP loosely inspired by the setup used in curth2021inductive: we similarly make use of the covariates $X$ of the ACIC2016 dataset dorie2019automated, and simulate our own outcomes and treatment assignments for greater transparency and control. We begin by binarizing all continuous covariates at randomly sampled cutoff points, obtaining processed covariates $X^*$, and then use them in a linear model for $\mu_0(\cdot)$, including up to third-order interaction terms of $X^*$, while $\tau(\cdot)$ is simply linear in $X^*$. This setup includes 3 main experimental knobs:
{We vary $\rho$ on the x-axis of our plots, which we split into four setups based on characteristics 2 and 3: Setup A is unconfounded and estimators & selectors get $X$ as input, Setup B is unconfounded and estimators & selectors get $X^*$ as input, Setup C is confounded and estimators & selectors get $X$ as input and Setup D is confounded and estimators & selectors get $X^*$ as input. A more formal description of the DGP can be found in Appendix (ref). Further, throughout we split training data of size $n$ into $n_{train}= \frac{2n}{3}$ for training of all estimators and $n_{val}=\frac{n}{3}$ to be used by the selection criteria, and use an independent test-set of size $n_{test}=500$ to evaluate a criterion by the test-set PEHE of its selected best model. We use $n=1000+500$ as a default, but also consider the effect of having more ($n=2000+1000$) or less ($n=500+250$) data available in Appendix (ref). Every experiment is repeated for 20 random seeds, across which we report means and standard errors (SEs). Finally, we present additional results using the standard ACIC2016 and IHDP benchmarks in Appendix (ref), where we also outline why we believe they allow for less interesting analyses -- highlighting further that careful design of DGPs was important to gain interesting insights.}
{Here, we briefly examine the performance of the underlying learners themselves in Fig. (ref). We include an oracle that is trained directly on the (usually unknown) $\tau(X_i)$ to provide a lower bound on error due to misspecification (in particular, this highlights that in setup A and C, LRs cannot capture the CATE well while GBs could). We see that there are indeed interesting performance differences across learners and settings, meaning that no single learner, estimation strategy or ML method consistently dominates all others. As expected, R- and DR-learner show good performance especially when CATE is relatively simple. T- and ES-learners perform worst for $\rho=0$, but relatively better for $\rho$ large; the opposite is true for S-learners. Note that the performance of the LR S-learner is particularly poor for $\rho>0$ because it can only learn a constant treatment effect (i.e. it is severely misspecified for $\rho>0$). In Appendix (ref), we additionally consider imbalanced treatment group sizes and find that this worsens mainly the performance of the R-learner. \parfillskip=0pt}
We are now ready to examine the performance of the different model selection criteria: In Fig. (ref) we present results of different criteria choosing between all 10 learner-method combinations. For legibility, we report performance in terms of $\Delta \text{PEHE}_{fact}$, the difference between the test-PEHE of any given selector and the factual choice (lower is better, negative means better than factual).
{Q1: General performance trends. In Fig. (ref), we observe that across all selectors, most performance differences and gains relative to factual performance are observed when i) treatment is randomized and the treatment effect is complex ($\rho$ large in A&B), ii) there is confounding and the treatment effect is simple ($\rho$ small in C&D) and iii) LRs are misspecified for the POs but not CATE (B& D). We observe that the performance of the plug-in criteria often follows the performance of their underlying methods (i.e. comparing the trends in Figs. (ref) and (ref)). Further, the plug-in criteria based on ES- and T-learner generally perform the worst, especially when implemented using LRs. In setups C & D, the plug-in surrogate based on the R-learner works well when CATE complexity is low, but not when it is high -- mimicking the pattern of the underlying estimator observed in the previous section. The pseudo-outcome criteria appear to perform best overall, with the R-pseudo-outcome often performing most similarly to the oracle selector. Note also that while the plug-in criteria based on indirect learners deteriorate substantially when misspecified (LR in Setups A & C), the pseudo-outcome criteria still perform well even when implemented using a misspecified model (LR).\parfillskip=0pt}
{Q2. Congeniality bias. Next we examine whether there is evidence for congeniality bias -- i.e. whether plug-in or pseudo-outcome surrogate criteria appear to inherently favor estimators with similar inductive biases as the strategy used to provide a validation target. We propose to measure this by calculating the proportion of times an (i) estimation strategy or (ii) underlying ML method is selected by a validation criterion whenever it does not identify the best estimator (intuitively, we make this distinction because whenever an estimator is the oracle choice, selecting it should not be considered biased).\footnote{{More formally, we investigate how often a selector $\mathcal{E}_j$ chooses the best estimator $\hat{\tau}^*$ that minimizes the true (unobservable) PEHE $\mathcal{E}^{oracle}_\tau$ and, if not this top choice $\hat{\tau}^*$, which other type of estimator is chosen. That is, in Fig. (ref), we measure using the leftmost bar labeled `best' how often the correct best estimator is chosen --i.e. $\hat{P}(\arg\min_{\hat{\tau}_k} \mathcal{E}_j(\hat{\tau}_k)\!=\!\hat{\tau}^*)=1\!-\!\alpha_{\mathcal{E}_j}$, where $\alpha_{\mathcal{E}_j}$ is the error-rate of the selector -- , and other bars measure the proportion of times any specific type of estimator $\hat{\tau}_l\neq \hat{\tau}^*$ is chosen whenever $\mathcal{E}_j$ does not make the right choice. If some type of estimator is chosen disproportionately often -- i.e. if $\hat{P}(\arg\min_{\hat{\tau}_k} \mathcal{E}_j(\hat{\tau}_k)=\hat{\tau}_l|\hat{\tau}_l\neq \hat{\tau}^*)=1-\alpha_{\mathcal{E}_j}\!\!>>\!\frac{\alpha_{\mathcal{E}_j}}{|\mathcal{T}|}$ -- we consider this evidence that the selector $\mathcal{E}_j$ may be biased towards choosing estimators of the type $\hat{\tau}_l$.\parfillskip=0pt} } In Appendix (ref), we present a similar plot without making this distinction.}
In Fig. (ref) (top) we investigate congeniality between selection criteria and estimator strategy (i.e. R-, DR-, S-, ES- or T-learner, implemented using either ML method), pooled across all settings of Fig. (ref). We observe that there is clear evidence for congeniality bias between some of the plug-in criteria and their corresponding learning strategy; this is most pronounced for the criteria relying on indirect learners, the plug-in S-learner and ES-learner in particular. The plug-in criterion based on the R-learner also clearly suffers from this, while the DR-plug-in criterion exhibits less of this behavior. The pseudo-outcome criteria overall display less pronounced preference for their own strategy, with the LR-implementations of pseudo-outcome R- and DR-criteria giving least evidence for such congeniality bias overall.
{In Fig. (ref) (bottom) we then investigate congeniality bias between selection criteria and estimator method (i.e. LR or GB, used with any estimation strategy). Also here we observe clear evidence for congeniality biases in almost all criteria: LR- (GB-)based criteria appear to prefer learners implemented using LR (GB). The only exception appears to be the LR-pseudo-outcome selectors, who actually select GB learners more often; this may partially explain their good relative performance compared to the other LR-based selectors.\parfillskip=0pt}
{Q3. What do we lose through factual evaluation? Finally, we consider the question of why we see in Fig. (ref) that so many selectors using surrogates for the treatment effect can outperform factual selection in some scenarios. One possible explanation would be the exclusion of direct learners from the candidate pool available to $\mathcal{E}^{fact}_{Y}$-- however, as can be seen in Fig. (ref), there is often some indirect learner that matches performance of a direct learner. Two alternative explanations we wish to test would be i) the presence of covariate shift due to confounding and ii) the incorrect focus of factual selection on performance in terms of estimating the POs. We can test i) by including an importance weighted factual selector, and ii) by restricting the candidate estimator pool available to all selectors to $\mathcal{T}_{indirect}$ (i.e. excluding R- and DR-learner from the estimator pool).
In Fig. (ref)(left) we observe that i) does not seem to be the case as weighted and unweighted factual selection perform identically (one possible explanation for this is that none of the considered indirect estimators perform an internal covariate shift correction themselves). Considering ii) we do however observe that indeed both oracle and pseudo-outcome selectors appear to select different (better) indirect estimators than $\mathcal{E}^{fact}_{Y}$, also in the absence of covariate shift (more saliently in Setup B). In Fig. (ref)(right), we show that these, in turn, indeed perform worse in terms of estimating the POs themselves, a trade-off that we expected in Sec. (ref). }
We studied the CATE model selection problem and focused on building understanding of the (dis)advantages of different model selection strategies -- using factuals, plug-in surrogates or pseudo-outcome surrogates -- that have been used or proposed in recent work. Instead of attempting to declare a global `winner', we empirically investigated success- and failure modes of different strategies -- and in doing so found that there are scenarios where factual selection can be appropriate but also scenarios where pseudo-outcome surrogate approaches are likely to perform better (only plug-in surrogate approaches seemed likely to underperform throughout). We hope that some of the insights presented here will give a starting point for practitioners able to identify how the likely characteristics of their own application translate to the scenarios we considered -- for this purpose we include an additional digest of our findings in form of a Q&A with an imaginary reader in Appendix (ref). We also highlighted that there is a complex interplay between selection strategies, candidate estimators and the DGP used for testing -- congeniality bias is likely to arise when the inductive biases of estimators and selection strategies align. By doing so, we also hope to have demonstrated to the community the need to conduct more simulation studies relying on carefully constructed DGPs to allow to disentangle different forces at play in this problem, enabling more nuanced analyses.
{Limitations. Finally, note that we do not claim our results to be complete: to allow for interesting and nuanced insights, we needed to restrict our attention to specific questions and candidate estimators. We believe that there are a plethora of interesting questions to explore in this area, of which we only made an initial selection to serve as a starting point for discussion. It would, for example, be an interesting next step to consider how different criteria fare at selection between other classes of estimators, e.g. the method-specific neural-network-based estimators extending the work of shalit2017estimating, or, at a more microscopic level, at hyperparameter-tuning for any specific method. While our experimental results are limited to answering some of the questions we found most intriguing, we hope that the desiderata for experimental design that we discuss and implementations that we provide will allow future research to easily expand to further questions and associated DGPs. }
{
We would like to thank Toon Vanderschueren, the members of the vanderschaar-lab and anonymous reviewers for insightful comments and discussions on earlier drafts of this paper. AC gratefully acknowledges funding from AstraZeneca.}