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,089 characters · 17 sections · 50 citation commands
\twocolumn[
\aistatstitle{Robust and Agnostic Learning of Conditional Distributional Treatment Effects}
\aistatsauthor{ Nathan Kallus \And Miruna Oprescu}
\aistatsaddress{Cornell University and Cornell Tech } ]
Measuring treatment-effect heterogeneity along observed covariates is an important tool for interpreting the results of A/B tests on online platforms, program evaluations in social science whether experimental or observational, and clinical trials in medicine. These analyses can help diagnose how the treatment works, understand for whom it does and does not work, and assess fairness heckman1997making, crump2008nonparametric, kent2010assessing, kallus2022treatment. On the other hand, measuring distributional treatment effects (DTEs) such as quantile treatment effects (QTEs) is an important tool for understanding the impact of interventions beyond the mean, especially when outcomes are naturally very skewed, like income or platform usage bitler2006mean, firpo, belloni, ldml.
Motivated by the need to assess both heterogeneity and distributional impact, in this paper we study flexible, agnostic, and robust machine learning tools to estimate conditional DTE (CDTE) functions. Recent advances in causal machine learning have offered new methods for assessing effect heterogeneity by learning conditional average treatment effects (CATEs) slearner,causaltree, causalforest, kunzel2019metalearners,kennedy2020optimal,nie2021quasi. These works have highlighted the importance of learning CATEs directly, rather than learning conditional-average outcomes by treatment arm and taking their difference (also known as the plug-in approach). One issue with the plug-in approach is that it can wash out the effect signal (not robust): \eg, many variables strongly predict baselines but only a few modulate the effect. Another issue is that it fails to give best-in-class predictions (not agnostic): \eg, taking the difference of the best linear predictions of outcome by arm does not yield the best linear prediction of treatment effect.
We tackle the same challenges for CDTEs. We consider CDTEs for a very rich class of distributional metrics that includes quantiles, super-quantiles, and other coherent risk measures. Given any distributional metric in our class, we construct a pseudo-outcome that combines an initial guess for the CDTE along with a debiasing term. Our algorithm is then to regress this pseudo-outcome on covariates, using any given blackbox learner. In the case of the CATE, our method recovers the DR-Learner kennedy2020optimal. We show that this procedure is robust in the sense that the blackbox regression mimics having used the pseudo-outcome with the true CDTE as the “initial guess”, thus removing any bias or noise from fitting the baselines. We further show that our method is model-agnostic in the sense that, if the blackbox is not a universal approximator, we still get the best approximation to the CDTE function offered by the blackbox. Lastly, we show that our estimating procedure allows for valid statistical inference on the best linear projection of CDTE, thus enabling interpretable analyses of distributional effect heterogeneity. We demonstrate in a comprehensive simulation study that we obtain uniformly better performance than the plug-in approach for several types of CDTEs. Finally, we apply our method on a real world study of 401(k) eligibility and its impact on financial wealth.
We consider either an experimental or observational dataset with two treatments, denoted by $0$ and $1$. Each unit in the dataset is a draw from a population of baseline covariates $X\in\Xcal$, treatment indicator $A\in\{0,1\}$, and observed outcome $Y\in\RR$. The dataset consists of $n$ such independent draws, $Z_i=(X_i,A_i,Y_i)\sim Z=(X,A,Y),\,i=1,\dots,n$. We define the propensity score as $e^*(X)=\Prb{A=1\mid X}$. We assume throughout that $e^*(X)\in(0,1)$ almost surely, known as overlap.
Each unit is additionally associated with two unobserved potential outcomes, $Y(0),\,Y(1)\in\Rl$, representing the potential outcome we would observe if (possibly counter to fact) each treatment were applied. We assume we observe the potential outcome corresponding to the treatment indicator, $Y=Y(A)$, which also encapsulates an assumption of no interference between unit treatments. We assume unconfoundedness (ignorability) throughout: $Y(a)\indep A\mid X$. For experimental data this is ensured by design via random assignment of $A$ (often with covariate-agnostic assignment, $A\indep X$). For observational data, this is an assumption that all potential sources of confounding are captured in $X$. For our purposes, the only difference between the two cases is whether the propensity score $e^*(X)$ is known.
We are interested in the differences between the conditional distributions of $Y(1)$ and $Y(0)$, given $X$. In the next section, we describe specific metrics for these differences.
Notation. Given a distribution $F$, we define $\E_F[f(Z)]=\int f(z)dF(z)$. We let $\widehat\E_n$ denote the empirical expectation $\widehat\E_nf(Z)=\frac{1}{n}\sum_{i=1}^n f(Z_i)$. For a parameter $f$, we reserve $f^*$ to represent its true value and $\widehat{f}$ a value learned from the data. We let $\|f\|:=\EE_F[f(z)^2]^{1/2}$ be the $L_2$ norm of $f$. We use $D$ to denote directional derivatives: $D_hF(h)\lvert_{h=h'}=\frac{\partial}{\partial a}F(h'+a)\lvert_{a=0}$, whenever this exists. If $h$ is a vector of functions, $D_hF(h)=(D_{h_1}F(h), ..., D_{h_k}F(h))$ and $D_{h_i}F(h)\lvert_{h=h'}=\frac{\partial}{\partial a}F((h'_1, ..., h'_i+a,...,h'_k))\lvert_{a=0}$. We let $\text{conv}(\mathcal S)$ denote the convex hull of $\mathcal S$. For two numbers $a, b$, we take $a\lesssim b$ to mean $a\leq Cb$ for some universal constant $C$ and $a\asymp b$ to mean $cb\leq a\leq Cb$ for some constants $c$ and $C$. Finally, we let $\overline{1, n}$ denote the set of integers $\{1,...,n\}$.
CDTEs are functions mapping $x$ to a difference in some statistic of the conditional distributions of $Y(1)$ and $Y(0)$, given $X=x$. Examples of such statistics are the mean, yielding the CATE, and the $\tau$-quantile, yielding the CQTE. In this paper, we handle a very wide range of CDTEs given by statistics defined by moment equations chamberlain1992efficiency, ai2003efficient.
For brevity, we will define the following functions (scalar-valued and $\R m$-valued, respectively)
We also assume these functions exist in that at least one solution to (ref) exists for $F_{Y(a)\mid X}$ (see (ref) regarding multiplicity). In the examples below, we give specific names for $\kappa$ or $h$ (\eg, $q(F;\tau)$ for the $\tau$-quantile of $F$), in which case we use analogous abbreviations (\eg, $q_a(X;\tau)$).
Note that under unconfoundedness and overlap, $F_{Y(a)\mid X}$ is the same as $F_{Y\mid X,A=a}$, the conditional distribution of $Y$ given $X$ and $A=a$. Therefore, $\kappa^*_a$, $h_a^*$, and the CDTE are all identifiable from the data. That is, despite being defined in terms of potential outcomes, the $\operatorname{CDTE}$ depends only on the distribution of observed data $(X,A,Y)$ and not on that of the unobserved data $(X, Y(0), Y(1))$, under our assumptions. The question we address in this paper is how to learn $\operatorname{CDTE}$s from data.
We next review some important examples of CDTEs that fit into our framework. First note that CATE fits into this setup by setting $\rho(y,\kappa)=y-\kappa$ with $m=0$ in (ref) (no $h$'s). The power of our framework is that it captures many CDTEs of interest beyond averages.
Our first example is the conditional quantile treatment effect (CQTE). Given $\tau\in(0,1)$, the quantile at level $\tau$ of the distribution $F$ is defined as $q(F;\tau)=\inf\{y:F(y)\geq \tau\}$ (identifying $F$ with its cumulative distribution function). Whenever $F$ has a positive derivative at $q(F;\tau)$, it is given by (ref) with $m=0$ (no $h$'s) and
The corresponding CDTE is called the CQTE at level $\tau$.
Non-conditional QTEs are an important tool for quantifying the effects of treatments throughout the outcome distribution firpo, belloni, ldml. This is especially important when we suspect that the outcome distribution might be skewed or heavy-tailed (\eg, income).
CQTEs offer an opportunity to assess such effects at the individual level. In particular, the CQTE can capture the potential increase in an individual's chance to have very poor outcomes due to treatment, even if the average effects are good or neutral. That is, if $A=1$ denotes an intervention, then CQTEs answer the prediction question: given an individual's covariates, how would intervening affect the 10%-worst possible outcomes.
A related quantity is the difference between the $\tau$ and $1-\tau$ conditional quantiles across treatments: $\omega(X;\tau)=q(F_{Y(1)\mid X};\tau)-q(F_{Y(0)\mid X};1-\tau)$. This quantity can bound the (unobservable) individual treatment effect: $$\ts\Prb{\omega(X;\tau)\leq Y(1)-Y(0)\leq \omega(X;1-\tau)\mid X}\geq 1-4\tau.$$ For the sake of brevity and uniform treatment we focus on CDTEs (\ie, using the same statistic for both potential outcome distributions), but our results readily extend to quantities like this as well.
Our second example is the conditional super-quantile treatment effect (CSQTE). Given $\tau\in(0,1)$, the super-quantile (also known as conditional value-at-risk or tail expectation) at level $\tau$ of a distribution $F$ is defined as
The corresponding CDTE is called the CSQTE at level $\tau$.
Note that a $\beta$ realizing the above infimization is the quantile at level $\tau$, $q(F;\tau)$. Therefore, given positive density at the quantile, the super-quantile is given in (ref) by setting $m=1$ and
The super-quantile $\mu(F;\tau)$ is the largest-possible subpopulation average among all subpopulations comprising a $1-\tau$ fraction of the population of values described by $F$. When $F(q(F;\tau))=\tau$, this can be phrased as the average of values above the $\tau$ quantile. If we are interested in the left tail (average below a quantile), we can simply consider the the negative CSQTE in the negative outcomes. Unlike the quantile, the super-quantile is a coherent risk measure artzner1999coherent.
In particular, while the CQTE captures the impact on the outcomes at a single probability level, they can fail to provide a full picture of the risk profile as they are indifferent to anything beyond the threshold of the quantile. In contrast, the CSQTE captures effects beyond quantile breakpoint and quantifies the impact on the average, say, 10%-worst (or, best) outcomes. Crucially, since super-quantiles are a coherent risk measure, making individual decisions based on CSQTEs (\eg, intervene when the CSQTE is positive) is rational with respect to the coherent-risk axioms.
Our third example is a whole class of coherent risk measures. Coherent risk measures quantify how bad/good a random loss/reward is while satisfying certain axioms (monotonicity, sub-additivity, homogeneity, and translational invariance). A key result is that coherent risk measures are equivalent to distributionally robust optimization ruszczynski2006optimization. In this section, we focus on the conditional $f$-risk treatment effect (C$f$RTE), a family of coherent risk measures generated by $f$-divergences.
Given a convex $f:\Rl\to\Rl$ with $f(1)=0$, we define the conditional $f$-risk at level $\delta\geq0$ as:
where $D^f(G\|F)=\EE_F[f(dG/dF)]$. For example, the $f$-divergence for $f(x)=x\log x$ is the Kullback–Leibler divergence, and the $f$-risk is known as entropic value-at-risk (EVaR) which upper bounds the super-quantile at level $1-e^{-\delta}$ ahmadi2012entropic. The $f$-risk admits a dual formulation given by the following convex optimization problem rockafellar1974conjugate:
where $f^*(x^*)=\sup_{x\in\RR}~xx^*-f(x)$ is the convex conjugate of $f$.
Therefore, under appropriate regularity, $R^f(F;\delta)$ is given by (ref) with
The corresponding CDTE is called the C$f$RTE at level $\delta$. Frequently employed in actuarial science and finance, coherent risk measures are a fundamental tool for assessing risk. C$f$RTE can therefore be used to perform risk-benefit analyses on the treatment effect profiles of affected groups.
Learning CATEs. CATE estimation is a central problem in causal learning and finds uses in both the analysis of causal interventions and decision support for personalization. Going beyond fully-parametric models, early advances relied on semiparametric models that imposed structure on the CATE function robins1992estimating, van2003unified, vansteelandt2014structural. Recently there has been a surge of interest in leveraging machine learning for CATE estimation. These flexible methods either employ specific machine learning models such as Bayesian regression trees hill2011bayesian, hahn2020bayesian, random forests (RFs) causalforest, oprescu2019orthogonal, neural networks johansson2016learning, atan2018deep, shi2019adapting, or allow for arbitrary blackbox meta-learners kunzel2019metalearners,nie2021quasi by leveraging efficient influence functions robins2017minimax, kennedy2020optimal, curth2020estimating and Neyman orthogonality chernozhukov2018double, foster2019orthogonal. This paper is closest to the works on efficient influence functions and blackbox meta-learners, and we add to this literature by considering distributional statistics beyond averages.
Double machine learning and orthogonal statistical learning. Another vein of related literature is learning with nuisances. Both our work and the above methods based on regressing efficient influence functions may be phrased within the wider framework of orthogonal statistical learning foster2019orthogonal, which extends double machine learning (DML) chernozhukov2018double from estimation to minimizing loss functions involving unknown nuisances subject to Neyman orthogonality neyman1959optimal. Neyman orthogonal losses arise naturally from efficient influence functions ichimura2022influence, a fact that has been leveraged by foster2019orthogonal, kennedy2020optimal, curth2020estimating,athey2021policy to tackle both CATE and policy learning. These methods have learning rates that adapt to the complexity of the target rather than that of the nuisances, but they largely focus on conditional averages. Our work builds on this line of research and is most similar to kennedy2020optimal in that we propose nuisance agnostic CDTE estimators with guarantees beyond those given by Neyman orthogonality.
Distributional treatment effects. The literature on DTEs can be split into two categories: (i) estimating cumulative distribution functions of potential outcomes and (ii) directly estimating distributional parameters of interest. In the first category, the main approach is to model conditional counterfactual distributions using distribution regression chernozhukov2013inference, chernozhukov2020network or fully flexible approaches such as neural networks ge2020conditional, zhou2021estimating and mean kernel embeddings park2021conditional. Since these methods rely on plug-in estimation, they can be slow and biased when concerned with a particular DTE. The second category focuses on estimating a particular DTE. Existing orthogonal/efficient methods focus on unconditional DTEs firpo,belloni,ldml. Existing methods that tackle CDTEs rely on plug-in estimation park2021conditional or parametric methods hohberg2020treatment. Our work bridges the gap by proposing robust, agnostic, and flexible CDTE learning.
In this section, we propose an algorithm for learning $\operatorname{CDTE}$s from data. As a first step, consider the following naive, but straightforward estimation procedure:
where $\widehat{\kappa}_a(\cdot)$ are estimates for $\kappa_a(\cdot)$ (see (ref) regarding how these may be constructed). This estimator is known as a “plug-in” estimator or, in the context of CATE learning, a “T-learner” kunzel2019metalearners. Unfortunately, this approach has several drawbacks. One concern is that the treatment effect signal can be easily masked by noise in the baseline predictors. In particular, while many variables may strongly predict baseline response, only a few strongly modulate effect. The effect function may often be simpler, sparser, and/or smoother than each baseline function. Therefore, the plug-in estimator can suffer from excessive bias inherent in fitting baseline estimators in high-dimensions or using flexible models. If, on the other hand, we seek to use a simple model such as a linear fit, we will find that differencing the best linear predictors of baseline outcomes does not yield the best linear predictor of effect. For these reasons, it is imperative to learn CDTEs in a direct, model-agnostic, and robust way.
To address the short-comings of the plug-in estimator, we consider the plug-in prediction on each data point $\operatorname{CDTE}^{\text{Plugin}}(X_i)$, debias it using the observed action and outcome $A_i,Y_i$, and finally regress the debiased prediction on $X$ again. Our first task is to propose an appropriate debiased pseudo-outcome for CDTE learning.
We refer to $e^*,\alpha^*,\nu^*$ as nuisance functions, as they are unknown functions needed to construct our pseudo-outcome. One initial motivation for (ref) is that, by iterated expectations, we have $\E[\psi(Z,e^*,\alpha^*,\nu^*)\mid X]=\operatorname{CDTE}(X)$. That is, if we had plugged in the true nuisances, then the regression of our pseudo-outcome on $X$ yields exactly what we want. This is, however, not so surprising because if we plug in $\nu=\nu^*$, the last term in (ref) just has zero conditional expectation, given $X$, so we are left with $\kappa^*_1(X) - \kappa^*_0(X)$. The reason why (ref) is special is that, as we will show, if we make small errors in the nuisances, the impact on the conditional expectation of our pseudo-outcome is even smaller, leading to robustness guarantees. This is in contrast to the plug-in approach, where errors in $\widehat\kappa_a(X)$ propagate directly to $\operatorname{CDTE}^{\text{Plugin}}$, which is just their difference.
The origin of (ref) is that $\psi(Z,e^*,\alpha^*,\nu^*)$ is in fact the efficient influence function for the estimand $\EE[\operatorname{CDTE}(X)]$ (for a review of influence functions see kennedy2022semiparametric,ichimura2022influence). In particular, if our statistic is simply the mean ($\rho(y,\kappa)=y-\kappa$) then $\psi(Z,e^*,\alpha^*,\nu^*)$ reduces to the familiar doubly-robust influence function, $\E[Y\mid X,A=1]-\E[Y\mid X,A=0]+\frac{A-e^*(X)}{e^*(X)(1-e^*(X))}(Y-\E[Y\mid X,A])$ that produces the CATE when regressed on $X$ (as studied by kennedy2020optimal). As we never use the fact that $\psi(Z,e^*,\alpha^*,\nu^*)$ is the efficient influence function, we do not prove this or impose the necessary regularity conditions for this to actually hold precisely. Instead, we simply use a perturbation argument to essentially guess the form of (ref), given which we directly prove our robustness and inference guarantees.
We now describe our algorithm, which is summarized in (ref). We first split the data into $K$ even folds. We then construct a pseudo-outcome for each data point by plugging in estimates of the nuisances into (ref). The nuisances at a point are fit on data excluding the fold that the data point belongs to. This ensures that the data point and the nuisance estimates are independent without splitting the data into two and instead only using parts of it for each task. Finally, we regress the pseudo-outcome on $X$ using a given regressor. We use $\widehat{\EE}_n[W\mid X=x]$ to denote the function learned by the given regression method when regressing $W$ on $X$ given $n$ data points $(X_i,W_i)$, $i\in\overline{1, n}$. This notation affords us significant generality. For example, the regression method may be to minimize the sum of squared errors over some function class (\eg, linear or neural nets) or it may be given by local polynomial regression or RFs.
(ref) requires nuisance estimators as inputs. Exactly how these nuisances are estimated may depend on the particular scenario. Let us first discuss the propensity $e^*(x)$. If it is known, as in experimental settings, we may simply set it as our estimate. Otherwise, we can estimate it using probabilistic classification (\eg, logistic regression or neural nets with softmax output).
Next, we discuss $\nu_a^*(x)$. Most generally, since it is defined by solving the conditional moment restriction $\EE[\rho(Y,\nu_a^*(X))\mid X]=0$, this nuisance may be learned by employing methods made for solving such models ai2003efficient,chen2009efficient, bennett2019deep, bennett2020variational,athey2019generalized,khosravi2019non, dikkala2020minimax. However, in some specific examples, more direct methods may be applicable. For both CQTE and CSQTE, $\nu_a^*(x)$ includes a conditional quantile function, which can be learned using any quantile regression method, whether minimizing the check loss koenker1978regression or using forests meinshausen2006quantile. For CSQTE, we additionally need to fit the conditional super-quantile. Per (ref), that nuisance is given by the regression $\E[(1-\tau)^{-1}Y\II[Y\geq q_a(X;\tau)]\mid X,A=a]$. Therefore, one possibility is to split the training data (being all data excluding the $k\thh$ fold) into two halves, fit a conditional quantile estimate $\hat q^{(k)}_a(x;\tau)$ on one, set $\omega_i=(1-\tau)^{-1}Y_i\II[Y_i\geq \hat q^{(k)}_a(X_i;\tau)]$ on the other, and return $\hat\mu^{(k)}_a(x;\tau)=\widehat\EE_{(1-1/k)n/2}[\omega\mid X=x,A=a]$. In particular, a given $X$-regression method can simply be applied once to $A=0$ and once to $A=1$. Appendix A of dorn2021doubly provides guarantees for this procedure when the regression learner minimizes squared error over a class with bracketing entropy and olma2021nonparametric when using local linear regression.
Lastly, we discuss $\alpha_a^*(x)$. In some cases, it is a known function that need not be estimated. For C$f$RTE, it is equal to $(-1,0, 0)$. In other cases, it is given directly by other nuisances. For CSQTE, it is equal to $(-1,(1-\tau)^{-1}q^*_a(x;\tau))$ so we can simply re-use the estimate we constructed for $\nu_a^*$. In yet other cases, it is another nuisance that must be estimated. For CQTE, it is equal to $1/f_{Y\mid X=x,A=a}(q^*_a(x;\tau))$, the reciprocal of the density of $Y\mid X=x,A=a$ at the conditional quantile. One way to fit this suggested by leqi2021median is to split the training data into two halves, fit a conditional quantile estimate $\hat q^{(k)}_a(x;\tau)$ on one, set $\omega_i=K((Y_i-q^{(k)}_a(X_i;\tau))/b_n)/b_n$ on the other, where $K(u)$ is a kernel function such as the standard normal density, and return $\hat\alpha^{(k)}_a(x)=\widehat\EE_{(1-1/k)n/2}[\omega\mid X=x,A=a]$.
In this section, we study the finite sample error rates for Algorithm (ref) with arbitrary first- and second-stage estimators. For this and the next section, let us fix some $\operatorname{CDTE}$ with pseudo-outcome as in Definition (ref) and estimation procedures input to Algorithm (ref). We require the following boundedness conditions.
The first condition in (ref) ensures that both treatments and controls can be observed for any $X$ with some fixed probability. This is guaranteed in a randomized trial if $e^*(X)$ is a constant and otherwise is a standard assumption in observational studies. The other conditions encode the cross-term structure of the derivatives of our moments. In most examples, the requirement of $\rho$ being bounded amounts to requiring $Y$ to be bounded. Similar boundedness assumptions are often made in debiased machine learning for ATE and CATE to control remainder terms.
We can then prove the following conditional Neyman orthogonality for our pseudo-outcomes, which is the key step for learning guarantees.
The result shows that whether we use the pseudo-outcome with oracle nuisances or estimated nuisances as a regression target, the difference is bounded by a quadratic form in the nuisance errors, wherein $G,H$ govern which pairwise error products appear. Thus, even if nuisances are estimated slowly, the impact is marginal, as the error is squared.
We will show that, up to $\mathcal{E}=\sum_{k=1}^K\mathcal{E}(\hat e^{(k)},\hat \alpha^{(k)},\hat \nu^{(k)})$, our estimate $\widehat{\operatorname{CDTE}}$ produced by (ref) behaves the same as if we applied our regression method to the pseudo-outcome with oracle nuisances, $\widetilde{\operatorname{CDTE}}(x)=\widehat{E}_n[\psi(Z,e^*,\alpha^*, \nu^*)\mid X=x]$. In particular, if $\mathcal E$ is generally smaller than the error $\|\widetilde{\operatorname{CDTE}}-\operatorname{CDTE}\|$, then the leading behavior of $\widehat{\operatorname{CDTE}}$ will be equivalent to that of $\widetilde{\operatorname{CDTE}}$. The particular statement of this will depend on what last-stage regression method we are using.
The significance is that (ref) will behave like regressing an unbiased observation of $\operatorname{CDTE}$ on $X$, even though we used estimated nuisances. First, this means we can learn $\operatorname{CDTE}$ at rates that match the complexity of that function. Second, this implies that we can directly approximate $\operatorname{CDTE}$ and get model-agnostic best-in-class guarantees. For example, if we use linear regression, we would get the best linear approximation for $\operatorname{CDTE}$ at a rate of $n^{-1/2}$. This is especially important if we seek an interpretable model. In the next section, we show this also enables inference.
First, we consider an empirical risk minimization algorithm (nonparametric least squares): given a class $\Fcal\subset[\Xcal\to\Rl]$,
The rate $O_p(n^{-1/(2+r)})$ is generally the rate for regressing a known target using nonparametric least squares over $\Fcal$ with such bracketing entropy. For example, if $\Fcal$ is H\"older functions of smoothness $\beta$ in $d$-dimensional inputs then it satisfies the entropy condition with $r=d/\beta$ vdv_wellner and $n^{-\beta/(2\beta+d)}$ is the optimal rate for regression in such a class stone1982optimal. Thus, if $\mathcal E=O_p(n^{-1/(2+r)})$, our regression behaves as though we supplied it with the oracle nuisances.
Obtaining such a rate for $\mathcal E$ is generally lax because it is quadratic in nuisance-estimation errors. For example, if nuisances are estimated at the much slower rate $O_p(n^{-1/(4+2r)})$, the condition is ensured. The special structure in (ref) further permits some trade-off between the rates of different nuisances. In Appendix (ref), we give the pseudo-outcome for each of the examples in Section (ref) and instantiate (ref) to exactly characterize the trade-offs.
Going beyond empirical risk minimizers, leveraging (ref) and invoking a result of kennedy2020optimal \footnote{The result appears as Theorem 1 in v2 of the arxiv preprint.} we can characterize the behavior when using a last-stage regression method satisfying certain stability properties.
We next study how (ref) can be used for inference on the best linear projection of $\operatorname{CDTE}$: $$\ts \gamma^*\in\argmin_{\gamma\in\RR^p}\|\operatorname{CDTE}-\gamma\tr\phi(\cdot)\|, $$ for a given $\phi:\Xcal\to\RR^p$. In particular, $\phi$ may be subset just some of the features. Linear projections offer an interpretable view into distributional-effect heterogeneity. In the previous section, we showed (ref) can perform well at learning linear projections. Next, we show that we can further conduct inference on the coefficients, which can facilitate interpretation and credible conclusions.
(ref) implies that we can just use out-of-the-box OLS with the built-in OLS inference as the final stage regression in (ref). The inference results would remain valid, provided nuisances are estimated slowly but not too slowly. In particular, we should generally use robust (aka sandwich or Huber–White) standard errors, as we do not expect the projection to have homoskedastic errors.
In this section, we demonstrate our method and theoretical results by applying (ref) to learn CSQTEs ((ref)). We first benchmark its performance in simulated data and then illustrate its use in a study of 401(k) eligibility and its effect on wealth accumulation. We provide additional results for CQTEs and C$f$RTEs in (ref). Replication code is available at \url{https://github.com/CausalML/CDTE}.
We sample from the following data generating process:
where $\sigma$ is the logistic sigmoid. The coefficients in the expression for $A$ are chosen such that the true propensity lies in the range $[0.05, 0.95]$. We seek to measure CSQTE at level $\tau=0.75$, \ie, the conditional average of values above the $0.75$ conditional quantile. For this DGP, the true CSQTE at $\tau=0.75$ is given by $\mu_1(X; \tau)-\mu_0(X; \tau)= 1.29(e^{X_0+X_1}-e^{X_0})$, which is heterogeneous in $X$.
We estimate ${e}(X)$ using logistic regression and ${q}_a(X;\tau)$ using a quantile random forest (QRF) meinshausen2006quantile. We consider three options for estimating $\mu_a(X;\tau)$. First, we consider a flexible learner we term the superquantile RF (SQRF). SQRF uses a RF to calculate weights (like QRF) and then computes (ref), replacing the expectation by the weighted average over the data. Second, we consider doing the same using a Gaussian kernel for calculating weights, choosing the bandwidth by Silverman's rule silverman2018density. We refer to this as the slow learner because it suffers badly from the curse of dimension. Third, we consider estimating $\mu_a(X;\tau)$ using ordinary least squares for the regression $\frac{1}{1-\tau}\EE[Y\II[Y\geq \widehat{q}_a(x;\tau)]\mid X=x]$ using the estimated quantile $\widehat{q}_a$. We term this estimator a misspecified learner since it is unlikely to fully capture the complexity of the superquantile (which is an exponential function of the features). For the final stage of (ref), we use either an ordinary least squares model (CSQTE+OLS) or a RF (CSQTE+RF). We set $K=5$ as the number of folds required by (ref). All forests use scikit-learn scikit-learn defaults except for the minimum leaf size which is set to $n/20$ to control overfitting.
We compare the out-of-sample mean squared error (MSE) of the CSQTE estimator with that of the naive plug-in estimator from Eq. (ref). To account for possible smoothing by the last stage regressor, we also construct Plugin+OLS and Plugin+RF given by running an additional OLS/RF model on the cross-fitted plug-in predictions. And, when the second stage algorithm is OLS, we check whether the 95%-confidence interval OLS returns for the $X_1$ coefficient contains the coefficient from the true projection. The results are shown in (ref). We run $100$ simulations for each $n=100,200,\ldots,12800$ and evaluate MSE over a fixed set of 500 random $X$ values. Our CSQTE learner provides uniformly strong MSE performance, and the results show this is not just a consequence of the second-stage regression. For inference, we achieve good coverage whereas plug-in approaches yield little to no coverage. These findings confirm our theoretical results: the quadratic dependence on nuisance errors enables oracle rates when the nuisances are estimated slowly or are misspecified ((ref), (ref)), and we can obtain valid inference when projecting onto linear spaces ((ref)).
We apply the CSQTE estimator to study the impact of 401(k) eligibility on net financial assets. We use the dataset from chernozhukov2004effects, which is based on the 1991 Survey of Income and Program Participation. The data contains 9{,}915 observations with 9 covariates such as age, income, education, family size, marital status, IRA participation, {\etc} While the eligibility for 401(k) (the treatment $A$) is not assigned at random, poterba1994401,chernozhukov2004effects argue that unconfoundedness can be assumed conditional on the observed covariates. The outcome of interest ($Y$) is the net financial assets of an individual, defined as the sum of 401(k) balance, bank accounts and interest-earning assets minus non-mortgage debt.
We apply the CSQTE estimator to understand effect heterogeneity beyond the conditiona mean. Specifically, we estimate the average treatment effects on the bottom and top $25\%$ of financial asset holders, conditional on covariates. We analyze the results alongside CATE estimates given by a DR-Learner kennedy2020optimal as implemented by econml. For nuisance estimation, we use RF models with hyperparameters as in chernozhukov2018double. The superquantile model is SQRF described above, the quantile model is QRF, the outcome learner for CATE is a RF regression and the propensity model is a RF classifier.
We consider two options for second-stage regressions. First, we consider using an RF model using all 9 covariates. We train the three estimators (CATE and upper/lower CSQTEs), and plot the distribution of predicted conditional effects on the 9{,}915 observations in (ref). We observe that 401(k) eligibility has a much higher financial impact on the high end of the conditional net worth distributions as compared to those on the lower end. While the CATE distribution reassuringly lies in between the bottom 25% and top 25% distributions, it fails to capture this disparity in effects.
To understand the heterogeneity in CSQTEs and CATEs, we use an OLS projection in the final stage. We choose the top three features that the RF last stage picked up as most important (see (ref) in Appendix (ref)): income, age and education. The resulting coefficients and 95% confidence intervals are depicted in (ref). For the bottom 25%, income and education are the main (statistically significant) drivers of the average effects. The positive income coefficient suggests that 401(k) eligibility provides more gains to those who potentially have the funds to invest. Likewise, the negative education coefficient means that the effects are higher among less educated earners. We hypothesize that higher educated earners might have a more comprehensive financial education and would save and invest regardless of 401(k) eligibility, whereas 401(k) availability provides lower educated earners a default investment option. For the top 25% asset holders, higher age and education rather than income have the largest impact on 401(k) eligibility returns. This group also displays more variability as none of the coefficients are statistically significant. The CATEs provide middle-of-the-road estimates that miss out on trends at the two ends of the spectrum. Thus, CSQTEs are a powerful tool for uncovering different behaviours across the conditional distribution that a simple average would obscure.
We provide new tools to assess distributional effect heterogeneity through agnostic and robust learning of CDTEs in a generic framework that includes quantiles, super-quantiles, and $f$-risk measures. These tools can be used to analyze treatments as well as to support personalized decisions that take risk into account. We provide strong guarantees for both learning and inference, and demonstrate our methods in both synthetic and real data. We further discuss some limitations of our work in Appendix (ref).
This material is based upon work supported by the National Science Foundation under Grant No. 1939704 and the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, under Award Number DE-SC0023112.