EconBase
← Back to paper

Causal Q-Aggregation for CATE Model Selection

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.

50,745 characters · 9 sections · 42 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

\twocolumn[

\aistatstitle{Causal Q-Aggregation for CATE Model Selection}

\aistatsauthor{ Hui Lan \And Vasilis Syrgkanis* }

\aistatsaddress{ Stanford University\\ [email removed] \And Stanford University\\ [email removed]} ]

abstractAccurate estimation of conditional average treatment effects (CATE) is at the core of personalized decision making. While there is a plethora of models for CATE estimation, model selection is a nontrivial task, due to the fundamental problem of causal inference. Recent empirical work provides evidence in favor of proxy loss metrics with double robust properties and in favor of model ensembling. However, theoretical understanding is lacking. Direct application of prior theoretical work leads to suboptimal oracle model selection rates due to the non-convexity of the model selection problem. We provide regret rates for the major existing CATE ensembling approaches and propose a new CATE model ensembling approach based on Q-aggregation using the doubly robust loss. Our main result shows that causal Q-aggregation achieves statistically optimal oracle model selection regret rates of $\frac{\log(M)}{n}$ (with $M$ models and $n$ samples), with the addition of higher-order estimation error terms related to products of errors in the nuisance functions. Crucially, our regret rate does not require that any of the candidate CATE models be close to the truth. We validate our new method on many semi-synthetic datasets and also provide extensions of our work to CATE model selection with instrumental variables and unobserved confounding.

Introduction

Identifying optimal decisions requires understanding the causal effect of an action on an outcome of interest. With the emergence of rich and large datasets in many application domains such as digital experimentation, precision medicine, and digital marketing, identifying optimal personalized decisions has emerged as a mainstream topic in the literature. Identifying optimal personalized decisions requires understanding how the causal effect changes with observable characteristics of the treated unit. For this reason, many recent works have studied the estimation of conditional average treatment effects (CATE):

align[align omitted — 59 chars of source]

where $X$ are observable features and $Y(d)$ is the potential outcome of interest under treatment $d\in \{0,1\}$.

This has led to a surge of many different methods for CATE estimation using machine learning techniques, such as deep learning shalit2017estimating, shi2019adapting, lasso lasso, random forests wager2018estimation,oprescu2019orthogonal, Bayesian regression trees hahn2020bayesian, as well as model-agnostic frameworks such as meta-learners kunzel2019metalearners and double machine learning kennedy2020towards, foster2023orthogonal,nie2021quasi. Machine learning CATE estimation has been considered both under the assumption of unconfoundedness, i.e., that outcomes $Y(1),Y(0)$ are independent of the assigned treatment $D$, conditional on the observed features $X$, as well as when there is unobserved confounding and access to instrumental variables IV,Tripadvisor.

However, the effectiveness of each approach depends on the learnability of the various causal mechanisms at play and the structure of the CATE function. This has led to the pursuit of automated and data-driven model selection approaches for CATE estimation schuler2018comparison,mahajan2022empirical,curth2023search,alaa2019validating. Due to the fundamental problem of causal inference, i.e., that we do not observe both counterfactual outcomes $Y(1),Y(0)$ for each unit, there does not exist a perfect analogue of out-of-sample model selection based on mean squared error, as is typically invoked in regression problems. For this reason, many works have proposed proxy loss metrics nie2021quasi,foster2023orthogonal,kennedy2020towards,alaa2019validating, analogous to the mean square error, which, under assumptions, can be used for data-driven model evaluation. In particular, mahajan2022empirical, highlights that proxy metrics that incorporate double robustness properties (that we either have a good estimate of the outcome regression function $\mathbb{E}[Y\mid D,X]$ or the propensity function $\mathbb{E}[D\mid X]$), tend to perform well in many practical scenarios.

Moreover, ensembling of CATE models tends to typically increase performance, as compared to simply choosing the single best model based on the loss. The empirical success of model aggregation in causal inference has also been demonstrated multiple times. In the 2016 Atlantic Causal Inference Conference Competition, 3 out of the 5 best performing models to estimate the average treatment effect were ensemble models dorie2019automated. Moreover, athey2019ensemble also showed that ensemble methods can outperform single models at predicting counterfactuals.

\paragraph{Contributions and prior work.} Motivated by these empirical findings and by the success of ensemble methods in the regression and supervised machine learning literature, we consider the problem of CATE model selection: given a set of $M$ candidate CATE models $\cF=\{f_1,\ldots,f_M\}$ and $n$ samples, can we construct an ensemble $\hat{f}=\sum_j \theta_j f_j$ that competes with the best model:

align[align omitted — 140 chars of source]

No rigorous theoretical treatment of this CATE model selection problem from observational data has been provided in the literature. Prior theoretical work considered only randomized trial data han2022ensemble, or considered only CATE estimation over rich function spaces nie2021quasi,foster2023orthogonal, which is technically very different from the goal of competing with the best of $M$ given models, or consider only guarantees for model evaluation alaa2019validating, but not selection. For the case of regression problems, the optimal oracle selection error $\epsilon$ that is achievable is of the order of $\log(M)/n$. This rate allows selecting from an exponentially sized set of candidate models $M$ and decays fast with the sample size, as $n^{-1}$ instead of $n^{-1/2}$.

Our main result is such an optimal oracle selection error for CATE model selection. We propose a novel CATE model ensembling method based on Q-aggregation Rigollet2014 with a doubly robust loss and prove that it attains the oracle selection error of $\log(M)/n$ plus second-order terms that depend on the regression and the propensity estimation errors. Moreover, even the $\log(M)$ term can be removed if one has a relatively good prior over which models will be the better performers, which can be quite natural in the context of CATE model selection.

In the regression setting, optimal selection error rates cannot be achieved using ERM approaches, i.e., by minimizing squared loss over some ensemble hypothesis space Rigollet2014. Thus, prior approaches to CATE ensembles nie2021quasi,han2022ensemble that either select the best model based on a proxy loss, or optimize over the space of all possible convex combinations of models, will necessarily achieve suboptimal selection error rates (either $\sqrt{\log(M)/n}$ or $M/n$ in the worst case). As a set of side results, we provide oracle selection rates for both of these methods for the CATE problem (which are novel in the literature and require a generalization of the main statistical learning theorems with nuisance functions presented in foster2023orthogonal). These results also highlight when one should expect ERM approaches to achieve close to optimal error.

Moreover, a theoretical result on optimal oracle model selection necessarily needs to use techniques beyond what is outlined in the prior theoretical work of nie2021quasi,foster2023orthogonal,kennedy2020towards, which is our main technical contribution. Notably, the general theorems of foster2023orthogonal require a first-order optimality condition, which is not satisfied due to the non-convexity of the function space ${\mathcal F}$ with which we are trying to compete and cannot be applied to prove our main Q-aggregation theorem. Due to space constraints, we defer a detailed comparison to previous work in the Appendix (ref).

Finally, we provide an extensive empirical evaluation of our causal Q-aggregation ensemble method and of the ERM based ensemble approaches, on a broad spectrum of synthetic data as well as semi-synthetic data based on five real-world datasets from a multitude of application domains, such as political science welfare, economics poverty,401k_1,401k_2, education STAR and digital advertising criteo. We showcase how ensemble methods offer more robust solutions than any single meta-learning approach proposed in the literature and how Q-aggregation can many times lead to improvements in performance than ERM based approaches.

Preliminaries

We consider a setting where we are given a set of $M$ candidate CATE models $\{f_1, \ldots, f_M\}$ and a data set of $n$ samples. Each candidate model $f_i$ maps covariates $X$ to a CATE prediction and can be thought of as a CATE model that was fitted on a separate set of samples (e.g. fitted on training set and correspond to different hyperparameters, methods, and model spaces), and the $n$ samples that we are using are the validation set. Our goal is to select a model or an ensemble of models, i.e. $f_{\theta}=\sum_{j=1}^M \theta_j f_j$, that achieves mean-squared-error with respect to the true CATE function $\tau_0(x)=\mathbb{E}[Y(1)-Y(0)\mid X=x]$ that is comparable to the best model, i.e.:

align[align omitted — 140 chars of source]

As we will see in Section (ref), the true CATE model $\tau_0$, can be expressed as the minimizer of a risk $\min_{f} \mathbb{E}[\ell(Z;f(X), g_0)]$, which, in addition to the CATE function $f$ also depends on other nuisance functions $g_0$, which are the solution to auxiliary statistical estimation problems and which also need to be estimated from data. Moreover, the loss function varies depending on the argument that is used to identify the CATE (e.g. conditional exogeneity, instrumental variables, etc.).

For this reason, we will consider the more general problem of model selection in the presence of nuisance functions and present our main theorems at this level of generality. Subsequently in Section (ref) and Appendix (ref), we will instantiate the general theorems for particular CATE selection problems.

\paragraph{Model Selection with Nuisance Functions} Given a set of candidate functions $\ensuremath{{\cal F}}=\{f_1, \ldots, f_M\}$, we care about minimizing an expected loss function $\mathbb{E}[\ell(Z;f(X), g_0)]$, where $g_0$ is the ground truth nuisance function. In particular, our goal is to identify an ensemble of the candidate estimators, i.e.

align[align omitted — 151 chars of source]

that controls the oracle selection error:

align[align omitted — 139 chars of source]

Our target is to prove an oracle inequality of the order: $O\left(\frac{\log(M)}{n}\right)$, with access to $n$ samples of the variables $(X,Z)$, which is the optimal rate in the absence of nuisance functions. Since, we also do not know $g_0$, we need to construct an estimate $\hat{g}$ from the data. Our goal is to derive an error that would have a second order dependence on the estimation error of $g_0$

align[align omitted — 123 chars of source]

for some appropriately defined error function that would decay at least as fast or faster than $1/n$ under assumptions on the estimation error of $\hat{g}$. Unless otherwise stated, the estimate $\hat{g}$ will be assumed to be trained on a separate sample of size $n$ (e.g. the training sample used to fit the models $f_i$).

\paragraph{Notation} We use short-hand notation:

align[align omitted — 128 chars of source]

where $Z$ denotes all the random variables and $X$ is some subset of them. Moreover, we denote with $\|f-f'\|^2 = P(f'(X) - f(X))^2$ and with $\|X\|_{L^p}=\mathbb{E}[X^p]^{1/p}$. For any convex combination $\theta$ over a set of candidate functions $\{f_1, \ldots, f_M\}$, we define the notation:

align[align omitted — 109 chars of source]

We also denote with $e_j\in \Theta$ the vector that has $1$ at coordinate $j$ and zero otherwise (that is, $f_{e_j}=f_j$) and denote with $\ensuremath{\mathrm{co}}(\ensuremath{{\cal F}})$ the convex hull of space $\ensuremath{{\cal F}}$.

ERM-Based Selection

We start our investigation by exploring what regret rates are achievable by ensembling or stacking approaches based on ERM, i.e. minimizing the empirical risk over the weight parameters $\theta\in \Theta$, for some constraint space $\Theta$. To simplify exposition, we restrict attention in this section to square losses:

align[align omitted — 80 chars of source]

where $g$ is a vector of nuisance functions and $\hat{Y}(g)$ is some functional of these functions that satisfies:

align[align omitted — 84 chars of source]

for some target function $\tau_0$.

Two major ensemble methods are choosing the single best model, ($\Theta_b=\{e_1, \ldots, e_M\}$), and choosing the best convex combination,

align[align omitted — 283 chars of source]

It is known that such ERM approaches lead to suboptimal regret rates without further assumptions, even in the absence of nuisance functions (i.e., when we know $g_0$) Rigollet2014. As a side result of this paper, we provide oracle results for these estimators in the presence of nuisance functions and refine prior results in the regression literature, providing some insights as to when we should expect these methods to achieve near optimal performance in practice.

theorem[Gurantees for Best-ERM] Consider the case of the square loss and let: \begin{align} \ensuremath{Bias}(X; \hat{g}) = \mathbb{E}\left[\hat{Y}(\hat{g}) - \hat{Y}(g_0)\mid X\right] \end{align} Assume that $\hat{Y}(g)$ is absolutely uniformly bounded by $U$ for all $g\in G$ and all functions $f_j$ are absolutely uniformly bounded by $U$. Let \begin{align} \bar{f}_* :=& \operatorname*{arg\,min}_{f\in \ensuremath{\mathrm{co}}(\ensuremath{{\cal F}})} \|f - \tau_0\|, & f_* :=& \operatorname*{arg\,min}_{f\in \ensuremath{{\cal F}}} \|f - \tau_0\| \end{align}{ and let: \begin{align} \epsilon = & \|f_*-\tau_0\|^2 - \|\bar{f}_*-\tau_0\|^2 \end{align} } Then with probability $1-\delta$: \begin{align} \ensuremath{{\cal S}}(\hat{\theta}_b) \leq & O\left(\epsilon + \frac{\log((M\vee n)/\delta)}{n} + \mathbb{E}[\ensuremath{Bias}(X;\hat{g})^2]\right) \end{align}

The result is enabled by a generalization of the main theorems in foster2023orthogonal. Incorporating the bound $ \|f_*-\tau_0\|^2 - \|\bar{f}_*-\tau_0\|^2$ is a novel addition to the literature, even in the absence of nuisance functions. We see that the best selector achieves the target rate only when $\epsilon$ vanishes faster than $\log(M)/n$, which would happen only if the best model in $\ensuremath{{\cal F}}$ is very close to the true model $\tau_0$ (c.f. $\min_{f\in \ensuremath{{\cal F}}}\|f-\tau_0\|^2\leq \epsilon$), which is quite unlikely when $\tau_0$ is complex, or when the best model $\bar{f}_*$ in the convex hull ends up putting most of the support on only one model, i.e., there is a dominant model (c.f. $\min_{f\in\ensuremath{{\cal F}}} \|f-\bar{f}_*\|\leq \epsilon$). Typical prior bounds depend on the larger quantity $\epsilon=\|f_*-\tau_0\|^2$, which is not close to zero if the optimizer in the convex hull is primarily supported on one model.

Theorem (ref) is also closely related to the results in van2003unified. As we show in Appendix (ref), our Theorem (ref) improves upon this prior work in: i) only incurring a dependence on the nuisance error through the term $\mathbb{E}[\ensuremath{\textsc{Bias}}(X;\hat{g})^2]$, which possesses doubly robust properties in our main application of CATE model selection, unlike the bound in van2003unified, which would also depend on a term of the form $\mathbb{E}[(\hat{Y}(\hat{g})-\hat{Y}(g_0))^2]$, which does not possess doubly robust properties , ii) only depending on the smaller quantity $\epsilon$, unlike the bound in van2003unified which would depend on $\|f_*-\tau_0\|^2$; as we discussed $\epsilon$ can be much smaller than $\|f_*-\tau_0\|^2$, when there is a dominant model, but all models are far from $\tau_0$.

theorem[Gurantees for Convex-ERM] Consider the case of the square loss and assume that $\hat{Y}(g)$ is absolutely uniformly bounded by $U$ for all $g\in G$ and all functions $f_j$ are absolutely uniformly bounded by $U$. Then with probability $1-\delta$: \begin{align} \ensuremath{{\cal S}}(\hat{\theta}_c) \leq & O\left(\frac{M\log(n/\delta)}{n} + \mathbb{E}[\ensuremath{Bias}(X; \hat{g})^2]\right)\\ \ensuremath{{\cal S}}(\hat{\theta}_c) \leq & O\left(\sqrt{\frac{\log(M/\delta)}{n}} + \mathbb{E}[|\ensuremath{Bias}(X; \hat{g})|]\right) \end{align}

As we show in the appendix, the benchmark in the oracle selection guarantee in both the above theorems can be strengthened to be the best model in the convex hull of $\ensuremath{{\cal F}}$ and not just the best in $\ensuremath{{\cal F}}$. Moreover, we note that the slow rate result of $\sqrt{\log(M/\delta)/n}$ is closely related to the work of han2022ensemble, who provided the exact same rate for convex-ERM ensembling using the doubly robust loss that we will introduce in Section (ref), albeit only for the case of randomized trials, where we have $\mathbb{E}[|\ensuremath{\textsc{Bias}}(X;\hat{g})|]=0$, due to knowing the propensity function $\operatorname{Pr}(D=1\mid X)$. Thus this result is a direct generalization of their result to observational data. None of these aforementioned theorems achieves the optimal rate when we want to compete with the best in $\ensuremath{{\cal F}}$. The inability of ERM based approaches to achieve the optimal oracle selection rates motivates the main result of this paper in the next section: the use of ensemble approaches based on Q-aggregation.

Q-Aggregation with Nuisances

We assume that the function $\ell$ is $\sigma$-strongly convex in $f$ in expectation, i.e., with $\ell^{(2)}$ being the derivative of $\ell$ with respect to $f(X)$:

multline[multline omitted — 160 chars of source]

Consider a modified loss, which penalizes the weight of each model, based on individual model performance:

align[align omitted — 118 chars of source]

For any prior weights $\{\pi_1, \ldots, \pi_M\}$, the Q-aggregation ensemble is the solution to the convex problem:

align[align omitted — 183 chars of source]
theorem[Main Theorem] Assume that for some function $\ensuremath{\textsc{Error}}(\hat{g})$: \begin{multline} P(\tilde{\ell}_{\theta,g_0}(Z) - \tilde{\ell}_{\theta',g_0}(Z)) - P(\tilde{\ell}_{\theta,\hat{g}}(Z) - \tilde{\ell}_{\theta',\hat{g}}(Z)) \\ \leq \ensuremath{Error}(\hat{g})\, \|f_{\theta} - f_{\theta'}\|^{\gamma} \end{multline} for some $\gamma\in [0, 1]$ and $\ell(f, g_0)$ is $\sigma$-strongly convex in expectation with respect to $f$. Moreover, assume that the loss $\ell(Z; f(X), \hat{g})$ is $C_b(\hat{g})$-Lipszhitz with respect to $f(X)$, and both the loss and the candidate functions are uniformly bounded in $[-b,b]$ almost surely. For any $\mu \leq \min\{\nu, 1-\nu\} \frac{\sigma}{14}$ and $\beta \geq \{ \frac{8C_b^2(1-\nu)^2}{\mu}, 8\sqrt{3}bC_b(1-\nu),\frac{2C_b\nu(\nu C_b + \mu b)}{\mu}\}$, w.p. $1-\delta$: \begin{align} R(\hat{\theta}, g_0) \leq & \min_{j=1}^M\left( R(e_j, g_0) +\frac{\beta}{n}\log(1/\pi_j)\right) \\ & + O\left( \frac{\log(1/\delta) \max\{\frac{1}{\mu} C_b(\hat{g})^2, b\, C_b(\hat{g}) \}}{n}\right) \\ &+ \left(\frac{1}{\mu}\right)^{\frac{\gamma}{2-\gamma}} \ensuremath{Error}(\hat{g})^{\frac{2}{2-\gamma}} \end{align}

The theorem gives a more refined version of oracle selection that is adaptive to prior beliefs on which model is best and can even save the $\log(M)$ term, if the prior beliefs end up being correct. In its simplest form, when prior weights are uniform, then $\beta$ is irrelevant in the optimization problem and can be taken to satisfy the condition of the theorem, as part of the analysis, i.e. the method does not have any hyperparameter $\beta$ to choose. Moreover, we can always take $\nu=1/2$, in which case the ensemble has the oracle selection error:

align[align omitted — 209 chars of source]
remark[Computational Considerations] The computation of the weights $\theta$ that optimize the Q-aggregation objective corresponds to a convex optimization problem, which can be solved efficiently by modern convex optimization solvers. Drawing from the results in greedy, we also show in Appendix (ref) that a greedy approximate solution to the Q-aggregation problem provides the same statistical guarantees as in Theorem (ref). The greedy approximation requires only finding the best single model to the Q objective, as well as the best convex combination between that best single model and any other model (see Algorithm (ref)). This will always return a convex combination of at most two baseline models and can be implemented by a simple line search over the scalar that represents the convex combination of the two models, as well as the procedure that finds the minimum of the objective over the $M$ models (which is linear in $M$).

Square Losses with Nuisance Functions

We specialize the main theorem to the case of square losses, where the target label depends on unknown and estimated nuisance functions as defined in Equation (ref), with target parameter $\tau_0$ as defined in Equation (ref). For such loss functions, we can instantiate the main theorem to obtain the following corollary.

corollary[Square Losses with Nuisances] Let: \begin{align} \ensuremath{Error}(\hat{g}) = 2\sqrt{\mathbb{E}\left[\ensuremath{Bias}(X; \hat{g})^2\right]} \end{align} Assume that $\hat{Y}(\hat{g})$ and all functions $f_j$ are absolutely uniformly bounded by $U$. Then the Q-aggregation ensemble $\hat{\theta}$ with $\nu=1/2$, $\beta\geq\max\{112\,U^2, 56 U^3\}$ and loss $(\hat{Y}(\hat{g})-f(X))^2$, satisfies w.p. $1-\delta$: \begin{multline} \|f_{\hat{\theta}} - \tau_0\|^2 \leq \min_{j=1}^M \|f_j - \tau_0\|^2 + \frac{\beta}{n} \log(1/\pi_j) \\ + O\left( \frac{\log(1/\delta)}{n} + \ensuremath{Error}(\hat{g})^2\right) \end{multline}

If we impose more restrictions on the functions in $\ensuremath{{\cal F}}$, then we can show a corollary, with weaker requirements on the nuisance functions:

corollaryAssume that for all $\theta,\theta'\in \Theta$: \begin{align} \|f_{\theta}(X)-f_{\theta'}(X)\|_{L^\infty} \leq C \|f_{\theta}(X)-f_{\theta'}(X)\|_{L^2}^{\gamma} \end{align} for some constants $C$ and $0<\gamma\leq 1$. Moreover, let: \begin{align} \ensuremath{Error}(\hat{g}) = 2\, C\, \mathbb{E}\left[\left|\ensuremath{Bias}(X; \hat{g})\right|\right] \end{align} Then under the remainder conditions and definitions of Corollary (ref), w.p. $1-\delta$ \begin{multline} \|f_{\hat{\theta}} - \tau_0\|^2 \leq \min_{j=1}^M \|f_j - \tau_0\|^2 + \frac{\beta}{n} \log(1/\pi_j) \\ + O\left( \frac{\log(1/\delta)}{n} + \ensuremath{Error}(\hat{g})^{\frac{2}{2-\gamma}}\right) \end{multline}

If the function class $\ensuremath{{\cal F}}$ contains smooth enough functions, then the above condition tends to hold. For instance, mendelson2010regularization shows that if the functions $f_j\in F$ lie in an RKHS with a polynomial eigendecay, then the above condition holds for some constant $\gamma$ that depends on the rate of eigendecay. Moreover, as we show in Appendix (ref), even when $\ensuremath{{\cal F}}$ does not contain smooth functions, we can still show the above property for $\gamma=1$ and $C \sim M$, as long as functions in $\ensuremath{{\cal F}}$ are not co-linear. Here, the number of models $M$ and not the logarithm of $M$ appears. However, assuming that the nuisance error term decays faster than $n^{-1/2}$, this linear in $M$ term appears only in lower-order terms.

\paragraph{Guarantees without Sample Splitting}

When learning with nuisance functions, it is often desirable to train the nuisance functions on a separate sample, as it reduces overfitting bias. However, sample splitting decreases efficiency and statistical power, and, thus, it might not be beneficial for small datasets. The following theorem provides theoretical guarantees when the nuisance estimate and the ensemble weights are constructed using the same sample.

theorem[Guarantees without Sample Splitting] Under the same conditions as in Corollary (ref), with the exception that the nuisance function $\hat{g}$ is trained on the same sample as the parameter $\hat{\theta}$. The Q-aggregation ensemble $\hat{\theta}$ satisfies w.p. $1-\delta$: \begin{multline} \|f_{\hat{\theta}} - \tau_0\|^2 \leq \min_{j} \|f_j - \tau_0\|^2 + \frac{\beta}{n} \log(1/\pi_j) \\ + O\left( \frac{\log(1/\delta)}{n} + r_n^2 + \ensuremath{Error}(\hat{g})^2\right) \end{multline} where $\ensuremath{\textsc{Error}}(\hat{g})^2$ is of the order: \begin{align} \mathbb{E}\left[\ensuremath{Bias}(X;\hat{g})^2\right] + \mathbb{E}[(\hat{Y}(\hat{g}) - \hat{Y}(g_0))^4] \end{align} and $r_n$ denotes the critical radius of the function space $\operatorname{star}(Y_G)$ where $Y_G := \left\{ (\hat{Y}(g) - \hat{Y}(g_0))^2 \mid g \in G\right\}$.

Connections to Neyman Orthogonality

In this section, we draw a connection between the bias in labels and the literature of Neyman orthogonality. When we are interested in estimating an average treatment effect, then we can typically phrase this as estimating a parameter of the form: $\tau_0 = \mathbb{E}[\hat{Y}(g_0)]$, for some appropriately defined function $\hat{Y}(g)$. In such settings, we say that the moment function $M(g):=\mathbb{E}[\hat{Y}(g)]$ is Neyman orthogonal chernozhukov2018double with respect to the nuisance functions $g$ if for all $g\in G$, the following directional derivative is zero:

align[align omitted — 86 chars of source]

Similarly, we can define a conditional analogue, when the quantity of interest is a conditional expectation: $ \tau_0(X) = \mathbb{E}[\hat{Y}(g_0)\mid X]$. We say that the conditional moment function $M_X(g)=\mathbb{E}[\hat{Y}(g_0)\mid X]$ is conditionally Neyman orthogonal if for all $g\in G$, a.s. over $X$:

align[align omitted — 88 chars of source]

When the label function $\hat{Y}(g)$ satisfies the latter property then by a second order Taylor expansion we have:

align[align omitted — 138 chars of source]

Thus we can instantiate all the theorems in the aforementioned sections, by using the above bias term. Typically this second order remainder term will depend quadratically in the error of the nuisances, i.e.

align[align omitted — 150 chars of source]

or will contain only products of errors of the different nuisance functions that compose $g$ (see next section). In both cases, the $\ensuremath{\textsc{Error}}(\hat{g})^2$ term, will most times be of lower order than $\log(M)/n$ or $\min_j \|f_j - \tau_0\|_{L^2}^2$. The next section and Appendix (ref), draw on this connection and provide two Neyman orthogonal label functions for the case of model selection for CATE under unobserved confounding and with access to an instrument.

Doubly Robust Q-Aggregation

We now focus on our main question: model selection for CATE, i.e., $\tau_0(X)=\mathbb{E}[Y(1)-Y(0)\mid X]$. We focus on the setting where we observe all confounding variables and we defer to Appendix (ref) the case of unobserved confounding with access to an instrumental variable (e.g. A/B testing with non-compliance).

In the no unobserved confounding case, we observe a set of variables $W$ (that is, a superset of $X$), such that the conditional exogeneity property holds: $Y(1),Y(0) \perp \!\!\! \perp D\mid W$. Under conditional exogeneity, the CATE is identified as:

align[align omitted — 71 chars of source]

where $h_0(D,W) := \mathbb{E}[Y\mid D, W]$. As is well known oprescu2019orthogonal,semenova2021debiased,foster2023orthogonal,kennedy2020towards, this identifying formula can be robustified by incorporating propensity estimation: $\tau_0(X) = \mathbb{E}[\hat{Y}(g_0)\mid X]$

align[align omitted — 69 chars of source]

with $g=(h,a)$, $h$ an estimate of $h_0$ and $a$ an estimate of the signed inverse propensity function $a_0$ (a.k.a. the Riesz representer in chernozhukov2022riesznet):

align[align omitted — 101 chars of source]

This gives rise to the doubly robust square loss:

align[align omitted — 71 chars of source]

For any estimate $\hat{g}=(\hat{h}, \hat{a})$, define $\hat{q}(D, W)$ as:

align[align omitted — 74 chars of source]

An important fact of the doubly robust target, exploited also in prior works, is that it satisfies the mixed bias property (see Appendix (ref)):

align[align omitted — 164 chars of source]

Thus we can apply Corollary (ref) to obtain:

corollary[Main Corollary] Let $\tau_0(X)=\mathbb{E}[Y(1)-Y(0)\mid X]$ and suppose that conditional exogeneity holds conditional on some observed $W$ that is a superset of $X$. Let: \begin{align} \ensuremath{Error}(\hat{g}) = 2\sqrt{\mathbb{E}\left[\mathbb{E}\left[ \hat{q}(D, W)\mid X\right]^2\right]} \end{align} and assume that $\hat{Y}(\hat{g})$ all functions $f\in F$ are uniformly and absolutely bounded by $U$. Then the doubly robust Q-aggregation ensemble $\hat{\theta}$, based on the DR-Loss, with $\nu=1/2$ and a uniform prior $\pi$, satisfies that w.p. $1-\delta$: \begin{multline} \|f_{\hat{\theta}} - \tau_0\|^2 \leq \min_{j} \|f_j - \tau_0\|^2 \\ + O\left( \frac{\log(M/\delta)}{n} + \ensuremath{Error}(\hat{g})^2\right) \end{multline} For any prior $\pi$, if $\beta\geq\max\{112\,U^2, 56 U^3\}$, w.p. $1-\delta$: \begin{multline} \|f_{\hat{\theta}} - \tau_0\|^2 \leq \min_{j} \|f_j - \tau_0\|^2 + \frac{\beta}{n} \log(1/\pi_j) \\ + O\left( \frac{\log(1/\delta)}{n} + \ensuremath{Error}(\hat{g})^2\right) \end{multline}

The error term can also be further upper bounded by an application of a Cauchy-Schwarz inequality by:

align[align omitted — 226 chars of source]

Thus we can either control the $L^4$ moment of the prediction error of the two nuisances, or guarantee an $L^\infty$ rate for one and an $L^2$ rate for the other. Note that if the above product of error terms is of lower order than $n^{-1/2}$, then the effect of these errors on the quality of the causal ensemble is of second-order importance.

If we further make an assumption that:

align[align omitted — 125 chars of source]

Then by Corollary (ref), we can take in Corollary (ref):

align[align omitted — 159 chars of source]

As we saw in Section (ref), the latter assumption holds if either the functions in $\ensuremath{{\cal F}}$ are smooth or they contain independent components in their predictions.

remark[Priors on CATE models] In the context of CATE estimation, incorporating priors is quite natural. For instance, we have a strong prior that models that come out of meta-learning approaches that only use outcome modeling should perform worse. Thus, we can put lower prior probability on S- and T- learner models and larger prior probability on X-, DR- and R- learner models. Moreover, we can incorporate priors that more complicated models are less probable than simpler models. For instance, we can put higher prior on linear cate models or shallow random forest models.

Doubly Robust Q-Aggregation for CATE with Instruments

We consider a widely encountered setting of estimating a CATE from a stratified randomized trial with non-compliance. In this setting we know the cohort assignment policy $\pi_0(X) = \operatorname{Pr}(Z=1\mid X)$, and we are interested in estimating the conditional local average treatment effect $$\tau_0(X) = \mathbb{E}[Y(1)-Y(0)\mid D(1)>D(0), X]$$ where $D(z)$ is the potential outcome of the chosen treatment under different cohort assignments (aka recommended treatments). This effect is identified as:

align[align omitted — 276 chars of source]

where $\tilde{Z} = Z - \pi_0(X)$.

Let $\hat{\alpha}$ be an estimate of:

align[align omitted — 121 chars of source]

and $\hat{\beta}$ an estimate of:

align[align omitted — 125 chars of source]

and let $\hat{\tau}=\hat{\alpha}/\hat{\beta}$. Then we can construct the random variable

align[align omitted — 106 chars of source]

and we can consider Q-aggregation with the loss:

align[align omitted — 56 chars of source]

In this case, we can show the following mixed bias property (see Appendix (ref)):

align[align omitted — 153 chars of source]

Thus applying Corollary (ref) we have:

corollary[Main Corollary for CATE with Instruments] Let $\tau_0(X)=\mathbb{E}[Y(1)-Y(0)\mid D(1) > D(0), X]$ and suppose that the instrument is a stratified randomized trial with non-compliance and with a known propensity model $\pi_0(X)=\operatorname{Pr}(Z=1\mid X)$. Let: \begin{align} \ensuremath{Error}(\hat{g})^2 = 4 \mathbb{E}\left[(\tau_0(X) - \hat{\tau}(X))^2 \frac{(\beta_0(X) - \hat{\beta}(X))^2}{\hat{\beta}(X)^2}\right] \end{align} and assume that $\hat{Y}(\hat{g})$ is uniformly and absolutely bounded by $U$ in a sufficiently small neighborhood of $g_0$ and that all functions $f\in F$ are uniformly and absolutely bounded by $U$. Then the doubly robust IV Q-aggregation ensemble $\hat{\theta}$ with $\nu=1/2$ and a uniform prior $\pi$, satisfies that w.p. $1-\delta$: \begin{multline} \|f_{\hat{\theta}} - \tau_0\|^2 \leq \min_{j} \|f_j - \tau_0\|^2\\ + O\left( \frac{\log(M) +\log(1/\delta)}{n} + \ensuremath{Error}(\hat{g})^2\right) \end{multline} Moreover, for any prior $\pi$, for $\beta\geq\max\{112\,U^2, 56 U^3\}$, w.p. $1-\delta$: \begin{multline} \|f_{\hat{\theta}} - \tau_0\|^2 \leq \min_{j} \|f_j - \tau_0\|^2 + \frac{\beta}{n} \log(1/\pi_j)\\ + O\left( \frac{\log(1/\delta)}{n} + \ensuremath{Error}(\hat{g})^2\right) \end{multline}

Moreover, note that since $\hat{\tau} = \hat{\alpha}/\hat{\beta}$, we also have that $$\tau_0 -\hat{\tau} = \frac{\alpha_0 \hat{\beta} - \hat{\alpha} \beta_0}{\hat{\beta}\beta_0 }= \frac{(\alpha_0 - \hat{\alpha})\hat{\beta} + \hat{\alpha} (\hat{\beta}-\beta_0)}{ \hat{\beta}\beta_0}$$ Thus overall we have:

align[align omitted — 309 chars of source]

Assuming that $\beta_0(X)\geq \underline{\beta}$ and that we impose the same restriction on our estimate $\hat{\beta}(X)$ (which is a "minimal conditional instrument strength") and assuming that $\hat{\alpha}(X)^2\leq \bar{\alpha}^2$, then we have:

align[align omitted — 223 chars of source]

By Cauchy-Schwarz inequality the latter is upper bounded by:

align[align omitted — 244 chars of source]

Thus if we learn the effect of the instrument on the treatment (typically referred to as the compliance model) with an $L^{\infty}$ guarantee (e.g. a high-dimensional sparse linear logistic regression), then it suffices to learn the effect of the instrument on the outcome with an $L^2$ error. Moreover, as long as the error on the compliance model $\hat{\beta}$ is $o(n^{-1/4})$ and the product of the errors of $\hat{\alpha}$ and $\hat{\beta}$ is $o(n^{-1/2})$, then the nuisance impact is of lower order in the final guarantee for the Q-aggregation ensemble.

Again, if we further make an assumption that:

align[align omitted — 118 chars of source]

Then we can take:

align[align omitted — 256 chars of source]

and it suffices to control only the RMSE errors of $\hat{\alpha}$ and $\hat{\beta}$ to be $o(n^{-1/4})$.

remark[Beyond Stratified Trials with Non-Compliance] When the instrument policy $\pi_0(X)$ is un-known, i.e. when we have a conditional instrument that corresponds to some natural experiment, then the quantity $\hat{Y}(\hat{g})$ is sensitive to errors in the estimate $\hat{\pi}$ of $\pi_0$. To reduce this sensitivity, we can incorporate residualization of $Y$ and $D$ (as also proposed in Tripadvisor): \begin{align} \hat{Y}(g) = \hat{\tau}(X) + \frac{(\tilde{Y} - \hat{\tau}(X) \tilde{D}) \tilde{Z}}{\hat{\beta}(X)} \end{align} where $\tilde{Y}=Y-\hat{h}(X)$, $\tilde{D}=D-\hat{r}(X)$ and $\tilde{Z}=Z-\hat{\pi}(X)$ and $\hat{h}$ is an estimate of $h_0(X)=\mathbb{E}[Y\mid X]$, $\hat{r}$ is an estimate of $r_0(X)=\mathbb{E}[D\mid X]$ and $\hat{\pi}$ is an estimate of $\pi_0(X)=\mathbb{E}[Z\mid X]$. With such a proxy label the error term in the Corollaries presented in this section, will contain extra terms that can be upper bounded by terms of the order $\|\hat{\pi}-\pi_0\|_{L^4}^2\, \|\hat{r}-r_0\|_{L^4}^2$, $\|\hat{\pi}-\pi_0\|_{L^4}^2\, \|\hat{h}-h_0\|_{L^4}^2$, $\|\hat{\pi}-\pi_0\|_{L^4}^2\, \|\hat{\beta}-\beta_0\|_{L^4}^2$, $\|\hat{r}-r_0\|_{L^4}^2\, \|\hat{\beta}-\beta_0\|_{L^4}^2$ and $\|\hat{h}-h_0\|_{L^4}^2\, \|\hat{\beta}-\beta_0\|_{L^4}^2$.

Experiments

We assessed the performance of our doubly robust Q-aggregation method on fully simulated and semi-synthetic data sets. We constructed several candidate CATE models based on meta-learning approaches using XGboost for regression and classification sub-problems, such as for estimating the nuisance functions: $\hat{\mu}(X,D)$ (an estimate of $\mathbb{E} [Y \mid X, D]$), $\hat{\mu}_d(X)$ (an estimate of $\mathbb{E} [Y \mid X, D=d]$), and $\hat{\pi}(X)$ (an estimate for the propensity). For constructing the CATE models we considered 8 meta-learning strategies, namely S-, T-, IPW-, X-, DR-, R-, DRX-, and DAX-learners. Details of each learner are given in Appendix (ref).

We present results for Q-aggregation, convex stacking, and best-ERM (selecting the model with the best doubly robust loss). Note that convex stacking is a special case of Q-aggregation with $\nu$ set to 0. For Q-aggregation, we chose $\nu=0.1$ and did not incorporate any priors. Each data set is divided into 3 portions: $60\%$, $20\%$, and $20\%$. Candidate CATE models and nuisance functions are trained on the $60\%$ of the data. Subsequently, the CATE ensemble weights are trained on $20\%$ of the data. Finally, CATE RMSE is evaluated on the remaining $20\%$ of the data. For each model, we report the RMSE regret when compared to the oracle model selection - the model that achieves the lowest RMSE on the test set (i.e. $\|f_{\hat{\theta}} - \tau_0\|^2 - \min_{j} \|f_j-\tau_0\|^2$).

table*[table* omitted — 894 chars of source]
table*[table* omitted — 2,081 chars of source]

\paragraph{Simulated Data} We employed six distinct data generation processes and performed experiments (DGP) on 100 instances of each DGP. In all cases, a scalar covariate $X$ is drawn from a uniform distribution. Each DGP is characterized by 3 functions: propensity $\pi(X)$, conditional baseline response surface $m(X)$, and treatment effect $\tau(X)$. Treatment $D$ is sampled from a binomial distribution with probability $\pi(X)$, and outcome Y is simulated by: $$Y = m(X) + D \tau(X) + \epsilon$$ where $\epsilon \sim \mathrm{N}(0,\sigma^2)$. The details of each DGP can be found in the supplementary materials.

\paragraph{Semi-synthetic Data} Semi-synthetic data was generated from five datasets used in previous studies: welfare dataset welfare, poverty datasetpoverty, Project STAR dataset STAR, 401k eligibility dataset 401k_1,401k_2, and the Criteo Uplift Modeling criteo dataset. Since our analysis assumes a binary treatment, only one treatment is considered for datasets with multiple treatments.

We used two simulation modes to generate the outcome from covariates. In the simple mode, the outcome is generated based on a simple linear CATE and baseline response function plus Gaussian noise (with $\sigma = 0.2$). In the fitted mode, the output is generated from a random forest regressor that is fitted on the real dataset, with randomness introduced by the addition of the model residual uniformly at random.

\paragraph{Results} Overall, there is no single meta-learning model that performs consistently well across all DGPs or semi-synthetic datasets. We present in Table (ref) the mean RMSE regret for different DGPs and different semisynthetic datasets, where the RMSE regret is normalized by the mean RMSE regret of all models for each DGP/dataset. The results in Table (ref) show that best-model selection, convex stacking, and Q aggregation achieve a significantly lower average RMSE regret compared to choosing any single model.

On semi-synthetic datasets, Q-aggregation achieves comparable RMSE regret with convex stacking (see Figure (ref)) and consistently outperforms both best-ERM and convex stacking in the case of a simple CATE and baseline response function. Table (ref) summarizes the RMSE regret for model selection, convex stacking, and Q-aggregation for semi-synthetic datasets with simple CATE functions, where each cell in the table reports [mean $\pm$ st.dev.] median (95%) of RMSE Regret over 100 experiments for each semi-synthetic dataset simulated using the simple mode. The last two rows present the percentage decrease in the mean and median RMSE regret for Q-aggregation with respect to convex stacking and model selection.

figure*[figure* omitted — 177 chars of source]

For simple CATE functions, Table (ref) shows that there is an decrease in the mean and median RMSE regret for Q-aggregation, when compared with convex stacking and model selection for every dataset. For datesets generated with the fitted CATE functions, the underlying model follows a S-learner setup. Thus, we expect the S-learners to be dominant models, which is evident Table (ref). As a result, we observe that the there is often an increase in the RMSE regret for ensemble models over the best single model, as is hinted by Theorem (ref), since in this case we have a dominant model.