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.
102,052 characters · 26 sections · 49 citation commands
Finding Subgroups with Significant Treatment Effects
\xdef\@thefnmark\@footnotetext{ Jann Spiess, Graduate School of Business, Stanford University, [email removed]; Vasilis Syrgkanis, Department of Management Science & Engineering, Stanford University, [email removed]; Victor Yaneng Wang, Global Priorities Institute, University of Oxford, [email removed]. We thank Saarthak Gupta for expert research assistance and audiences at Wharton, UT Austin, and CLeaR for helpful comments. This manuscript supersedes an earlier version that was selected for oral presentation at the 1st Conference on Causal Learning and Reasoning (CLeaR 2022) under the title “Evidence-Based Policy Learning” (Jann Spiess and Vasilis Syrgkanis). }
Keywords: Policy learning, heterogeneous treatment effects, randomized experiments, recursive partitioning, causal machine learning.
Randomized controlled trials (RCTs) are a standard way of inferring the causal effect of an intervention on outcomes of interest. However, such outcomes are often noisy, and estimated overall effects in RCTs can be small or imprecise. One could attempt to overcome these issues by collecting more data to increase sample sizes, yet in practice, this may be infeasible or too costly. We present an alternative solution: we develop tools that provide reliable evidence of the efficacy of the intervention by finding subgroups with statistically significant treatment effects, even if overall effects are noisy or small.
Our tool for finding subgroups with statistically significant effects builds upon a fast-growing literature on causal machine-learning methods, including estimators for heterogeneous treatment effects and personalized treatment-assignment policies. These methods allow us not only to estimate whether an intervention works on average but also to analyze for whom the intervention is particularly effective. However, available tools are not typically optimized for finding evidence of treatment effects in very noisy data. In particular, our article addresses two related shortcomings of existing approaches. First, existing machine-learning methods for the estimation of heterogeneous treatment effects and personalized treatment assignments tend to perform poorly in data with small and noisy effects. Second, unlike existing methods, we fundamentally incorporate significance testing into our procedure. We show that this motivates a particular objective function, which in turn requires new tools to implement.
We formalize the task of finding a subgroup with significant treatment effects in (ref), and show in (ref) that it translates into an objective that differs from the standard goal of finding those observations who have high treatment effects. We consider procedures that learn assignment rules based on training data from a randomized experiment. These assignment rules map observable characteristics of observations in the experiment to a binary indicator of whether an observation should be included in a subsequent hypothesis test. The chosen assignment rule is then deployed on held-out data, where we test whether there is a significant treatment effect on the chosen subgroup. We optimize for maximizing the probability of rejecting the null hypothesis of a zero effect, corresponding to finding reliable evidence of treatment effects. Under a Normal large-sample approximation, we show that this goal is equivalent to maximizing the ratio of mean to standard error of treatment effect estimates in the hold-out data. When conditional variances differ across covariates, then the solution to this goal can differ markedly from subgroups formed solely by those observations that have a predictably positive or large treatment effect.
Having motivated an objective function for finding subgroups with statistically significant effects, in (ref) we discuss an implementation based on a “significance tree” algorithm that performs greedy recursive partitioning. In (ref), we then consider approximate solutions based on submodular optimization. We show that empirical optimization of our target goal approximates an optimal assignment and reduces to a submodular minimization problem, for which we present approximation bounds for finite support under a sparsity assumption on conditional treatment effects and conditional variances.
In the simulation study in (ref), we show the gain from using our purpose-built significance tree for finding subgroups with statistically significant treatment effects, as well as the value of submodular (re-)optimization. In the simulation study, we compare the performance of our algorithm to tree-based algorithms that form subgroups based on positive estimated treatment effects alone. We demonstrate the improvement in test performance and highlight cases in which our assignment tests better than a simple overall test in the full training and hold-out sample combined.
We can illustrate the main idea behind our approach in a simple example with just one covariate, which we expand on in our simulation study. This example shows the main intuition for why optimizing for significant treatment effects can improve power, compared to testing effects using the full sample or selecting those observations with positive or high treatment effects. (ref) plots a draw of treatment and control outcomes on the $y$-axis against their univariate covariate on the $x$-axis. In this simulated example, there are three parts of the covariate space: the left ($x \in [0,1]$) has negative treatment effects and high variance, the middle ($x \in [1,2]$) has positive average treatment effect and high noise, and the right ($x \in [2,3]$) has positive treatment effects with small noise.
In our example, we now compare different methods of selecting observations for testing. If we used the full sample for testing, we would include the left area with negative treatment effects, which would go against our goal of finding evidence for a positive treatment effect. We therefore hope to do better by using data to select a subset on which we can achieve higher power. A standard machine-learning method that aims to find observations with positive treatment effects would often separate the leftmost from the remaining parts (dashed blue line). This would improve the power of the test, but still include the middle part, yielding a test with low power on a hold-out dataset because of the noisy outcomes in this region. In this example, our significance tree (blue solid line) selects instead only the right-most segment into the subgroup that gets tested, yielding better evidence of a positive effect. This is because the right-most part presents the clearest evidence of the efficacy of the treatment, since it combines a positive effect with small noise.
Building upon our baseline algorithm, (ref) provides two modifications that can further improve the power of finding evidence for positive treatment effects. First, we consider a continuous relaxation that tests for significant effects based on a weighted average of observations, rather than a simple subgroup. Second, we aggregate individual significance trees into a forest that can capture more flexible subgroups. We demonstrate the corresponding performance gains in a simulation.
Finally, (ref) describes additional extensions of our setup and algorithm. For instance, we could instead address the problem of finding significant heterogeneity in treatment effects. Alternatively, we could consider a mixed criterion that incorporates both the probability of passing a hypothesis test on the hold-out sample and the value of assigning the chosen observations to treatment conditional on passing that test. In addition, we consider alternative outcomes and reference assignments, estimated propensity scores, and cross-evaluated tests.
\paragraph{Related work.}
We build upon an extensive machine-learning literature on learning personalized policies from experimental or observational data dudik2011doubly,swaminathan2015counterfactual,nie2017quasi,kallus2018policy,kitagawa2018should. A large volume of works uses ideas from the doubly robust literature in semi-parametric inference rubin2005general,rubin2007doubly,diaz2013targeted,van2014targeted,kennedy2017non,kennedy2019robust to construct objectives to be optimized over a space of policies. However, the majority of these works optimize the mean policy value and do not focus on the objective of succeeding in a subsequent hypothesis test on a fresh sample. Closest in spirit to our approach is kallus2021more, which improves the efficiency of policy-learning estimators in data with limited overlap by a retargeting procedure that assigns a higher weight to regions with good overlap. Somewhat analogously, our paper shows that for data where we struggle to obtain statistically significant overall effects, we can extract more reliable evidence by focusing on subgroups in which effects are large and measured more precisely. However, our approach aims to maximize performance on a hold-out hypothesis test, whereas kallus2021more instead aims to improve the efficiency of a policy-learning estimator.
Our work also builds upon work on recursive partitioning and treatment effect estimation in the causal machine-learning literature kunzel2017meta,athey2016recursive,athey2016generalized,nie2017quasi,athey2017efficient,zhou2018offline,oprescu2018orthogonal,friedberg2018,chernozhukov2017orthogonal,chernozhukov2018plugin,foster2019orthogonal,Bertsimas2019-ms. Unlike this treatment-effect literature, we do not care about learning an accurate model of how the effect varies, but rather only learn a model of which subgroup to select for further testing.
Moreover, the majority of the heterogeneous treatment effect literature constructs subgroups mainly based on some mean heterogeneity criterion. Arguably closest to us is the work of powers2018some which analyzes a tree algorithm whose criterion instead uses a standardized heterogeneity of the form $(\tau(S_1)-\tau(S_2))^2/(\sigma(S_1)^2 + \sigma(S_2)^2)$ athey2016recursive, where $S_1, S_2$ is a candidate split, $\tau(S_i)$ is the mean effect on each subgroup and $\sigma(S_i)$ the standard deviation. However, this criterion tries to identify the modes of heterogeneity and does not target the selection of subgroups directly.
Our work also relates to recent work on false discovery rate control in multiple and adaptive hypothesis testing ramdas2019unified,lei2017star. Unlike our setting, the FDR literature typically considers the statistically more daunting task of testing on the same data that was used to construct the adaptive hypothesis. Here we consider a sample-splitting setting, where testing will be constructed on a separate data set and we specifically focus on the out-of-sample power of the test. Arguably closest to our work from the FDR literature is the concurrent work of duan2021interactive, which instantiates the FDR line of work on the heterogeneous treatment effect problem and presents an interactive algorithm that controls the false discovery rate of the constructed subgroup to be treated. The work of Armstrong2015-gc considers inference on optimal treatment assignments, where the goal is to obtain a subgroup for which all conditional average treatment effects are positive with high confidence. Relative to this work, we instead provide a subgroup that aims to maximize evidence for a positive average subgroup effect on hold-out data, thus sidestepping the multiple-testing challenge in Armstrong2015-gc.
Our setup and motivation are very closely related to concurrent and independent work in Talisa2021-ud that also considers the objective of finding subgroups with statistically significant treatment effects. Relative to their approach, we construct recursive-partitioning and model-based solutions where model selection itself is directly optimized for maximizing power. Relatedly, McFowland2018-wo considers the problem of identifying most-affected subgroups, which it operationalizes as a problem of detecting anomalous patterns. In a similar spirit, Ladhania2020-ln considers the selection of subgroups with consistently large treatment effects, which are then tested on separate data.
Recent work on providing robust versions of off-policy learning is similar in spirit to our motivation. Dwivedi2020 aims to identify subgroups with heterogeneous treatment effects that are stable (i.e. do not vary a lot) and are well-calibrated out of sample. Finally, concurrent work in leqi2021median provides a way of estimating assignment policies that maximize the median policy value, instead of the expected policy value, which typically is more robust to outlier observations. Both of these works could lead to better out-of-sample power of statistical tests since they lead to more robust subgroup heterogeneity identification.
We aim to find a subgroup of observations for which we can provide evidence of a positive treatment effect based on data from a randomized experiment. We learn this subgroup from a training sample of $n$ iid draws $(Y_i,W_i,X_i)$ from a population distribution of random variables $(Y,W,X)$, where the outcome $Y$ is real-valued, treatment $W \in \{0,1\}$ is assigned randomly with probability $\mathop{}\!\textnormal{E}[W] = p \in (0,1)$, and there are additional covariates $X \in \mathcal{X}$ available that are not affected by the treatment. We assume that $Y = Y^{(W)}$ with potential outcomes $(Y^{(1)},Y^{(0)})$ that represent the outcomes an observation would have in treatment ($W=1$) and control ($W=0$), respectively. We then define the treatment effect on an observation as $Y^{(1)} - Y^{(0)}$. Randomization of $W$ allows us to identify from the data not only the average treatment effect $\tau = \mathop{}\!\textnormal{E}[Y^{(1)} - Y^{(0)}]$, but also conditional average treatment effects $\tau(x) = \mathop{}\!\textnormal{E}[Y^{(1)} - Y^{(0)}|X=x]$.
We formalize the selection of a subgroup by an assignment policy $a: \mathcal{X} \rightarrow \{0,1\}$ that maps a covariate $x \in \mathcal{X}$ to the indicator $a(x) \in \{0,1\}$ of whether an observation with covariates $x$ is in the chosen subgroup. For a subgroup of the population defined in this way, we can define the total effect of assigning observations in this subgroup to treatment as \[ u(a) = \mathop{}\!\textnormal{E}[a(X) \: Y^{(1)} + (1-a(X)) \: Y^{(0)}] - \mathop{}\!\textnormal{E}[Y^{(0)}] = \mathop{}\!\textnormal{E}[a(X) \: \tau(X)]. \] In this article, we are concerned with finding such subgroups for which we can provide evidence that $u(a) > 0$, i.e. that the group-wise average treatment effect $\mathop{}\!\textnormal{E}[\tau(X_i)|a(X_i) {=} 1] = u(a) / \mathop{}\!\textnormal{P}(a(X_i) = 1)$ is positive. We could also interpret this goal within the context of policy learning as follows: we seek a policy $a$ that assigns a subgroup of observations to treatment, where we have evidence that if we deploy this assignment policy on future observations from the same population, this leads to positive lift $u(a)$ relative to assigning everybody to control.
Existing methods in the policy-learning literature are typically focused on finding a subgroup assignment $a$ that maximizes the expected lift $u(a)$. Note that a subgroup assignment that maximizes $u(a)$ is $\tilde{a}(x) = \mathbbm{1}\left(\tau(x) > 0\right)$, suggesting two related standard strategies. The first strategy is to estimate heterogeneous treatment effects $\tau(x)$ directly by $\hat{\tau}(x)$ and then assign based on $\hat{a}(x) = \mathbbm{1}\left(\hat{\tau}(x) > 0\right)$. The second strategy is to perform cost-sensitive classification so as to maximize the empirical objective $\frac{1}{n} \sum_{i=1}^n a(X_i) \: \hat{\tau}_i$. This latter approach would be based on an unbiased estimate $\hat{\tau}_i$ of $Y_i^{(1)} - Y_i^{(0)}$ for each observation $i$, such as the estimate $\tilde{Y}_i$ we consider in (ref) below.
We depart from existing methods by instead fundamentally incorporating hypothesis testing into our target criterion. In practice, after estimating an assignment $\hat{a}$, we often have to validate that the chosen subgroup is indeed is associated with a positive effect $u(\hat{a})> 0$ by providing evidence in the form of a statistical hypothesis test. We therefore assume that we also have access to an additional hold-out dataset with $N$ draws $(Y_j,W_j,X_j)$ from the same distribution as the experiment, on which we perform a test of the form
where $\hat{u}(\hat{a})$ is an estimate of $u(\hat{a})$, $\widehat{\mathop{}\!\textnormal{Var}}(\hat{u}(\hat{a}))$ is an estimate of the measurement noise $\mathop{}\!\textnormal{Var}(\hat{u}(\hat{a})|\hat{a})$ from the hold-out data, and $z_{1-\alpha}$ is the critical value of a one-sided Normal test with significance level $\alpha$ (that is, the $1-\alpha$-quantile of a standard Normal). The test is justified under standard regularity conditions and for appropriate estimators $\hat{u}(a),\widehat{\mathop{}\!\textnormal{Var}}(\hat{u}(a))$ by
under the null hypothesis $u(\hat{a}) = 0$, where we assume throughout that $\hat{a}$ assigns a sufficiently large part of the sample to treatment such that the limit remains well-behaved.
Our goal is to maximize the probability of passing the test, so we want to choose a subgroup assignment $\hat{a}$ with maximal expected probability \[ U(\hat{a})=\mathop{}\!\textnormal{E}[\widehat{T}(\hat{a})|\hat{a}] = \mathop{}\!\textnormal{P}\left(\frac{\hat{u}(\hat{a})}{\sqrt{\widehat{\mathop{}\!\textnormal{Var}}(\hat{u}(\hat{a}))}} \geq z_{1-\alpha}\middle|\hat{a}\right), \] where expectation and probability are over the hold-out sample. For an assignment policy that chooses nobody for the subgroup, we can in practice assume that the test statistic is standard Normal (that is, $p$-values are uniformly distributed), yielding baseline utility $U(\hat{a}) = \alpha$. By splitting the sample, optimizing for such a goal does not distort inference as long as the hold-out data is only used for the test.
In the previous section, we motivated the objective function
for the expected lift $u(a) = \mathop{}\!\textnormal{E}[a(X) \: \tau(X)]$ from the subgroup chosen by $a$. In this section, we consider a specific estimator for estimating the lift $u(a)$ and constructing a test. We then solve for the theoretically optimal (but generally infeasible) oracle maximizers of our testing objective when we apply a large-sample approximation for the training sample.
In order to derive results on maximizing the power of the hold-out test, we now consider a specific class of estimators $\hat{u}(a)$ for the lift $u(a)$ and their variance $\mathop{}\!\textnormal{Var}(\hat{u}(a)|a)$ in the hold-out sample $(Y_j,W_j,X_j)_{j=1}^N$. An unbiased estimator for $u(a)$ based on inverse propensity score weighting is
Even though it is unbiased, this estimator can have high variance, which we can typically reduce by appropriate residualization. To this end, note that the estimator
based on the centered outcome $Y - \hat{c}(X)$ is still unbiased for $u(a)$, given any fixed centering function $\hat{c}:\mathcal{X} \rightarrow \mathbb{R}$ that does not use any data from the hold-out set. The variance of $\hat{u}(a)$ can then be estimated by $\widehat{\mathop{}\!\textnormal{Var}}(\hat{u}(\hat{a})) = \frac{1}{N (N-1)} \sum_{j=1}^N (\hat{a}(X_j) \:\tilde{Y}_j - \hat{u}(\hat{a}))^2$.
While our theory remains applicable for any choice of centering function $\hat{c}$ that does not depend on hold-out data, we consider the specific choice that estimates \[ c^*(x) = (1-p) \mathop{}\!\textnormal{E}[Y| X=x, W=1] + p \mathop{}\!\textnormal{E}[Y| X=x, W=0]. \] Choosing $c^*$ minimizes the variance of $u(\hat{a})$ and leads to the standard semi-parametric efficient estimator Wu2017. Furthermore, spiessoptimal shows that this optimal centering function can be written as the minimizer of the weighted least squares regression problem \footnote{The first-order condition for each $x$ is: $\frac{1}{p} \mathop{}\!\textnormal{E}[Y| X=x, W=1] + \frac{1}{1-p} \mathop{}\!\textnormal{E}[Y| X=x, W=0] = \left(\frac{1}{p} + \frac{1}{1-p}\right) c(x)$. Re-arranging yields $c(x) = (1-p) \mathop{}\!\textnormal{E}[Y| X=x, W=1] + p \mathop{}\!\textnormal{E}[Y| X=x, W=0]$ as desired. Wu2017 proposes a leave-one-out estimator of average treatment effects in experiments based on residualization by this score. spiessoptimal argues that this specific score solves the weighted prediction problem in (ref). opper2021improving also considers an implementation for average effects in experiments using cross-fitted centering. }
which can subsequently be estimated by arbitrary non-parametric machine-learning methods in a data-adaptive manner. Across our implementations, we estimate $c^*$ in this way from the training data, yielding some estimate $\hat{c}$. We then perform tests based on the centered outcome $\tilde{Y} = Y - \hat{c}(X)$.
We now solve for optimal subgroup assignments given the above treatment-effect estimators. Expanding on the notation from (ref), we write
for the conditional average treatment effect and the variance in its estimation at covariate value $x \in \mathcal{X}$, where $Y-c(X)$ is the centered outcome for some centering function $c: \mathcal{X} \rightarrow \mathbb{R}$. In the case of the optimal centering function $c^*$, we obtain the minimal variance $\sigma^2_{*}(x) = \frac{\mathop{}\!\textnormal{Var}(Y^{(1)}|X{=}x)}{p} + \frac{\mathop{}\!\textnormal{Var}(Y^{(0)}|X{=}x)}{1-p}$, where we write $\sigma^2_{*}(x) = \sigma^2_{c^*}(x)$. We focus on the case where treatment effects $\tau(x)$ are small and of the order of magnitude $\frac{1}{\sqrt{N}}$, represented with the notation $ \tau(x) = \frac{\mu(x)}{\sqrt{N}}. $ This assumption allows us to model non-trivial testing in the limit, and expresses the use case where data is so noisy or effects so small that testing is meaningful.
In the local-to-zero limiting regime and for the estimators outlined above, for any fixed rule $a$ with $\mathop{}\!\textnormal{E}[a(X)] > 0$ and any centering function $\hat{c}$ we have that
where we write $\sigma^2 = \sigma^2_{\hat{c}}$. Since the asymptotic power of the resulting test is thus
(with the Normal cdf $\Phi$), maximizing asymptotic power corresponds to maximizing
The result holds more generally for continuous relaxations $a: \mathcal{X} \rightarrow [0,\infty)$ of the assignment problem, to which we return in (ref). In the binary case we study here, $a^2(x) = a(x)$, so the criterion becomes $\mathop{}\!\textnormal{E}[a(X) \: \tau(X)] / \sqrt{\mathop{}\!\textnormal{E}[a(X) \: \sigma^2(X)]}$. In additional results in (ref), we show that there are cases in which optimizing for the empirical equivalent of (ref) can lead to strictly more powerful tests than standard approaches from policy assignment that aim to maximize the expected lift $u(a)$.
In (ref), we characterized an oracle solution to the goal of finding an assignment that optimizes for passing a hypothesis test on hold-out data in terms of its population objective. In principle, one could estimate such an assignment using any estimation method for treatment effects $\tau(x)$ and variances $\sigma^2(x)$ on the training data. In this and the following section, we study how to optimize the sample analog of the criterion in (ref) in the training data. Specifically, our goal is now to find an assignment $\hat{a}$ based on $n$ observations $(Y_i,W_i,X_i)_{i=1}^n$ from the training sample that approximately maximizes the population objective (ref).
A popular approach in data science is recursive partitioning via binary decision trees athey2016recursive. Given the success of greedy tree-based heuristics in the statistical estimation literature, we focus on such heuristic solutions to this maximization problem. In particular, we define a greedy tree construction algorithm that at every step investigates the expansion of each leaf into two child nodes by splitting on a particular feature and at a particular threshold. Subsequently, we investigate the combination of the two possible treatment configurations of the two candidate child nodes and evaluate the empirical analog
of the population objective from (ref) for each candidate assignment $a$, choosing the candidate split that maximizes the statistic. This criterion corresponds to the training-sample equivalent of the hold-out test statistic. We note that unlike typical tree criteria, our $t$-statistic criterion cannot be solely computed based on the observations in a given node, since the standard deviation in the denominator is a global property. Thus at every candidate split we need to use the whole assignment policy to evaluate the performance of a new split. The details of our heuristic are presented in (ref).
\FloatBarrier
In this section, we return to the objective formalized in (ref), and outline an implementation based on submodular optimization. In order to do so, we first derive an alterative representation of an empirical analog to the objective from (ref). We then provide submodular optimization solutions separately for finite- and infinite-valued covariates. Throughout, our results remain applicable at the given sample size $n$ of the training sample, and not only asymptotically for the local-to-zero regime of (ref).
Again considering centered outcomes $Y_i - \hat{c}(X_i)$ with centering function $c$ (which can itself be estimated using cross-fitting on the training sample), we now solve for an assignment by maximizing the empirical criterion
in the training sample. In order to see the connection of the empirical objective from this section to the population objective from (ref), note that maximizing (ref) is equivalent to maximizing
among all assignments with $\sum_{i=1}^n a(X_i) \: \tilde{Y}_i > 0$ (assuming that there is at least one positive $\tilde{Y}_i$), where maximizing the last expression is then also equivalent to maximizing $\widehat{T}(a)$ from (ref) in the training sample.
To understand the properties of maximizing the empirical objective (ref), it will be helpful to define the population analog \[ V(a) = \frac{\mathop{}\!\textnormal{E}[a(X) \tau(X)]}{\sqrt{\mathop{}\!\textnormal{E}[a(X) \: \ensuremath{\nu}^2(X)]}} \] where $\ensuremath{\nu}^2_c(x) = \mathop{}\!\textnormal{E}[\tilde{Y}^2|X=x]$ for $\tilde{Y} =\left(\frac{W}{p} - \frac{1-W}{1-p}\right) (Y - c(X))$ and we write $\ensuremath{\nu}^2(x) = \ensuremath{\nu}_{\hat{c}}^2(x)$. To connect this objective to its empirical analog, the following properties will later be helpful.
Note here that the objective $\frac{\mathop{}\!\textnormal{E}[a(X) \tau(X)]}{\sqrt{\mathop{}\!\textnormal{E}[a(X) \sigma^2(X)] + \mathop{}\!\textnormal{Var}(a(X) \tau(X))}}$ for binary $a(x)$ only differs from the population objective in (ref) by the $\mathop{}\!\textnormal{Var}(a(X) \tau(X))$ term in the denominator, which is of lower order in the local-to-zero asymptotic approximation from (ref).
We now consider the problem of minimizing the empirical objective from (ref) when the covariates only take on a finite number of values. Observe that for a finite set $\mathcal{X}$, the empirical optimization problem of maximizing $\widehat{V}(a)$ is equivalent to finding a subset $S\subseteq\mathcal{X}$ that maximizes $\frac{\sum_{x\in S} w(x)}{\sqrt{\sum_{x\in S} v(x)}}$ for $v(x)= \frac{1}{n} \sum_{i \in \{1,\ldots,n\}: X_i=x} \tilde{Y}_i^2 \geq 0$ and $w(x)=\frac{1}{n} \sum_{i\in \{1,\ldots,n\}: X_i=x} \tilde{Y}_i$ (where we assume that the value if the denominator vanishes). Without loss of generality, when solving the empirical maximization problem, we can discard features $x\in \mathcal{X}$ for which $w(x) \leq 0$ (which includes all those with $v(x) = 0$). Thus when solving the optimization problem, it suffices to consider the case when $v(x), w(x)>0$, for all $x \in \mathcal{X}$.
The problem of maximizing $\frac{\sum_{x\in S} w(x)}{\sqrt{\sum_{x\in S} v(x)}}$ over $S\subseteq \mathcal{X}$ with $w(x), v(x) >0 $ for all $x \in \mathcal{X}$ is equivalent to minimizing the ratio of a monotone submodular function over a monotone modular function, noting that we converted our original maximization problem to a minimization problem by taking the reciprocal of the objective function. Corollary 3.3 of bai2016algorithms shows that the latter can be solved up to a $(1+\epsilon)$-approximation with $O\left(\log(1/\epsilon)\right)$ calls to a submodular minimization process which can be solved in polynomial time schrijver2000combinatorial; this process minimizes $\sqrt{\sum_{x\in S} v(x)} - \lambda \sum_{x\in S} w(x)$, which is also submodular, for some appropriately chosen $\lambda$ in each call. We note however that the result in bai2016algorithms requires solving a constrained submodular minimization, subject to the constraint of a non-empty solution. This constraint is non-trivial and not the standard submodular minimization problem solved by existing solvers (covering constraints can render submodular minimization NP-hard). In (ref) and (ref) below, we show how to modify the result and the algorithm in bai2016algorithms to only require an unconstrained submodular minimization oracle.
(ref) provides an approximation guarantee on the empirical criterion. However, for finite $\mathcal{X}$ of size $k$, a simple uniform deviation bound implies that the empirical objective converges to the population objective for any assignment $a: \mathcal{X}\to \{0,1\}$, i.e. w.p. $1-\delta$, \[ \max\left( \sup_{a} \left| \frac{1}{n} \sum_{i=1}^n a(X_i)\ \tilde{Y}_i - \mathop{}\!\textnormal{E}[a(X) \tau(X)]\right|, \sup_{a} \left| \frac{1}{n} \sum_{i=1}^n a(X_i)\ \tilde{Y}_i^2 - \mathop{}\!\textnormal{E}[a(X)\ensuremath{\nu}^2=(X_i)]\right| \right) \leq \epsilon(n, \delta) \] for any fixed centering function $\hat{c}$ and for $\epsilon(n, \delta)= O\left(\sqrt{\frac{k\log(1/\delta)}{n}}\right)$. Moreover, by (ref), \[ \ensuremath{\nu}^2(x) - \ensuremath{\nu}^2_{*}(x) = \frac{(\hat{c}(x) - c^*(x))^2}{p(1-p)}. \] In fact, we can invoke tighter uniform deviation bounds to obtain the following proposition.
Thus, for finite $\mathcal{X}$, we can approximate the optimal assignment arbitrarily well with respect to our objective function. Moreover, the latter function is a special case of a submodular function that can be solved much more efficiently than generic submodular minimization Stobbe2010.
This reduction also draws strong connections between our approach and the variance penalization literature in policy learning. In particular, swaminathan2015counterfactual proposes that one maximizes an objective of the form (after bringing it to our notation) \[ \sum_{x\in \mathcal{X}} a(x) w(x) - \lambda_n \left(\sum_{x\in \mathcal{X}} a(x) v(x) - \left(\sum_{x\in \mathcal{X}} a(x) w(x)\right)^2\right) \] with a vanishing $\lambda_n$ at a roughly $O(n^{-1/2})$ rate when $a$ belongs to a VC class. We see that unlike variance penalization we balance the standard deviation with the mean policy outcome and we do not necessarily have a vanishing penalty $\lambda$. Even closer to our objective is the second moment penalization in foster2019orthogonal, which recommends maximizing \[ \sum_{x\in \mathcal{X}} a(x) w(x) - \lambda \sum_{x\in \mathcal{X}} a(x) v(x). \] Again this objective has the same discrepancy with our objective in that we balance the root of the second moment and we do not necessarily have a vanishing $\lambda$. Moreover, in our objective the $\lambda$ is auto-tuned so as to minimize the empirical $t$-statistic.
Above, we considered maximizing the empirical objective (ref) for finite $\mathcal{X}$. However, we in many applications $\mathcal{X}$ may not be finite, but either uncountable or exponential in some implicit representation. For instance, if $x$ is represented by $J$ binary features, then $|\mathcal{X}|=2^J$. Moreover, typically, $x$ will be represented by $J$ continuous features. In this setting we do not only need to find an appropriate subset for a set of finite pre-defined subgroups, but we also need to construct how these subgroups are defined. In other words, we need to partition the $\mathcal{X}$ space in a finite set of partitions and then find the optimal assignment with respect to this partition. However, the choice of partitions also needs to happen in a manner that will lead to a large $t$-statistic.
As a first step towards maximizing our target objective for complex $\mathcal{X}$, we observe that it generally suffices to consider assignments that treat observations with covariates $x$ equivalently if neither their mean treatment effect $\tau(x)$ nor $\ensuremath{\nu}^2(x)$ changes.
Specifically, if we can find a partition $\mathcal{X} = P_1 \cup \ldots \cup P_k$ such that for all $\ell \in \{1,\ldots,k\}$ and for all $x,x'\in P_i$ we have that $\tau(x)=\tau_\ell = \tau(x')$ and $\ensuremath{\nu}^2(x)=\ensuremath{\nu}^2_\ell=\ensuremath{\nu}^2(x')$, then it suffices to optimize over assignments $a(x) = r(\pi(x))$, where $\pi(x)$ denotes the index of the partition that $x$ lies in. Equivalently, it suffices to identify the subset $S \subseteq \{1,\ldots,k\}$ of partitions to treat, with value $\frac{\sum_{\ell \in S} p_\ell \tau_\ell}{\sum_{\ell \in S} p_\ell \ensuremath{\nu}_\ell^2}$ for $p_\ell = \mathop{}\!\textnormal{P}(X\in P_\ell)$. Subsequently, in that case, we can apply a reduction to submodular minimization for finite elements $\mathcal{X}$, but now with the partitions $\{P_1, \ldots, P_k\}$ taking the role of the finitely many elements of $\mathcal{X}$. Hence, a $(1+\epsilon)$-approximate optimal policy can be found with $O\left(\log(1/\epsilon)\right)$ oracle calls to a submodular minimization oracle for decomposable submodular functions, as in (ref).
We now expand this idea of selecting groups of covariates with the same $\tau(x)$ and $\ensuremath{\nu}^2(x)$ to selecting groups of units that are similar in terms of their $\tau(x)$ and $\ensuremath{\nu}^2(x)$. Unlike a policy that defines similarity based on the full covariate vector $x$, such an assignment only relies on similarity in the low-dimensional space of $\tau(x)$ and $\ensuremath{\nu}^2(x)$.
Under the assumptions of the proposition, the ability to compute optimal assignments as in (ref) thus extends to the case where $\tau(x)$ and $\ensuremath{\nu}^(x)$ are only approximately constant within sets $P_\ell$. Specifically, (ref) implies that we can again apply a reduction to submodular minimization for finite elements $\mathcal{X}$, and translate a feasible approximate optimal policy for that case to one that still provides a good solution for maximizing $\widehat{V}$ in the case of more complex $\mathcal{X}$.
As a concrete application of our framework, we now focus on the case when $\mathcal{X}=\{0,1\}^J$. Here, we assume that both $\tau(x)$ and $\ensuremath{\nu}^2_{*}(x)$ are $r$-sparse functions, in the sense that they can be written as $\tau(x) = \tilde{\tau}(x_R)$ and $\ensuremath{\nu}^2_{*}(x) = \tilde{\ensuremath{\nu}}^2(x_R)$ for some subset $R\subseteq \{1,\ldots,J\}$ (where $x_R = (x_i)_{i\in R}$), with $|R|=r$. Furthermore, for simplicity of exposition we assume that $\tau(x), \ensuremath{\nu}^2_{*}(x) \in [-\sfrac{1}{2},+\sfrac{1}{2}]$ and that $\tilde{Y}_i \in [-1, +1]$. Our results can trivially be extended to any finite bound $H$ on the random variables at an extra multiplicative cost of $H$ in our bounds.
Importantly, we assume that the variables in $R$ are $\beta$-strictly relevant, in the sense that the variance reduction incurred by conditioning on one of the variables in $R$ is larger than conditioning on any variable in $\{1,\ldots,J\} \setminus R$, by an additive factor $\beta$ . This condition is formally expressed in the following definition.
We now apply the above sparsity condition to the functions $\tau(x)$ and $\ensuremath{\nu}^2_{*}(x)$. We have the following corollary of the results in syrgkanis2020estimation (in particular Theorem C.10.2).
Thus by taking the intersection of the partitions defined by two greedy decision trees with depth $d=r$, we guarantee the property required by (ref). This partition can also be thought of as the leaves of a decision tree with depth $d=2\, r$.
Subsequently, for any policy value defined on the leaves of a binary tree of depth $d$, adaptively chosen based on the data, we want to understand the concentration of the numerator and denominator of our objective $\widehat{V}$. If we let $A$ denote the space of binary policies defined on the leaves of an adaptively built binary decision tree of depth $d$, then we need to control the uniform deviations, that is, we want w.p.\ $1-\delta$, \[ \max\left( \sup_{a \in A} \left| \frac{1}{n} \sum_{i=1}^n a(X_i)\ \tilde{Y}_i - \mathop{}\!\textnormal{E}[a(X) \tau(X)]\right| \leq \epsilon(n, \delta), \sup_{a \in A} \left| \frac{1}{n} \sum_{i=1}^n a(X_i)\ \tilde{Y}_i^2 - \mathop{}\!\textnormal{E}[a(X)\ensuremath{\nu}^2_c(X_i)]\right|\right) \leq \epsilon(n, \delta). \] Assuming that $\tilde{Y}_i$ is bounded in $[-1, +1]$, then by standard results in statistical learning theory, both of the above deviations are satisfied with $\epsilon(n,\delta) = O\left(\sqrt{\frac{\textnormal{VC}(A)}{n}} + \sqrt{\frac{\log(1/\delta)}{n}}\right)$ shalev2014understanding, where $\textnormal{VC}(A)$ is the VC-dimension of policy space $A$. By existing results in statistical learning theory, the VC dimension of a binary classification tree of depth $d$ with $J$ features is $O\left(d\, 2^d \log(J)\right)$ (see Mansour97pessimisticdecision for the first claim of this fact and e.g. Section A.2 of syrgkanis2020estimation for a proof, or Lemma 4 of zhou2018offline, or Leboeuf2020 for a more extensive treatment). This yields the following lemma.
We can now combine these results into a final proposition about the approximation provided by a binary decision tree combined with submodular optimization on sparse data.
Finally, in the latter proposition we note that we can remove the $\log(n)$ factor in the leading term of the error, if the regression trees for learning $\tau(x)$ and $\ensuremath{\nu}^2(x)$ are also trained on a separate sample. This will lead to a leading error term of $\sqrt{2^{r+2}\log(2/\delta)/n}$, which could lead to an improvement for large $n$.
In this section, we take the empirical implementations outlined in (ref) and (ref), and evaluate experimentally how they compare to alternative methods in the policy learning and heterogeneous treatment effect literature. We consider two experimental designs. In the first setting, we consider a stylized data-generating process that highlights vividly the regime where our algorithm substantially outperforms alternative approaches. Subsequently, we consider a semi-synthetic experiment where we use the features and treatment assignments from a real-world dataset of an HIV trial duflo2019hiv. This was a four-armed randomized trial among youth in western Kenya, assessing the effect of various interventions on herpes type 2 (HSV-2) prevalence. We find that in both settings our significance tree achieves strictly higher or comparable out-of-sample $p$-values, as compared to the treatment subgroups identified by the alternative methods. This holds across multiple signal-to-noise regimes, with the difference being very stark in some conditions.
We compare the performance of different methods by using them to estimate a subgroup assignment on the training data, and then evaluating this assignment on hold-out data from the same distribution. The methods we compare are:
All of our methods are judged by calculating the out-of-sample $p$-value of the chosen subgroup assignment $\hat{a}$. In other words, we use the optimally centered policy value $\widehat{V}(\hat{a})$ on the hold-out and transform it into a $p$-value. The centering function $\hat{c}$ for the hold-out is estimated on the training sample using a weighted random forest with loss function as in (ref). We always estimate $\hat{c}$ this way across all methods, regardless of how the methods use $\hat{c}$ in their training procedure.
We first consider a very simple setting designed to showcase the advantage of our method over alternative approaches. We consider the case of a single feature $X$ that is uniformly distributed on $]0, 3]$. The feature space is partitioned in three intervals $]0,1]$, $]1,2]$, $]2,3]$, with parameters varying by interval. In each interval $t$, the value of the outcome can be described by
where $\lceil X \rceil \in \{1,2,3\}$ denotes the smallest integer not smaller than $X$. Here, we use as parameters for the three regions $\tau=(0, 2, 1)$, $\beta=(0, 10, 100)$, and $\sigma=(5, 10, 1)$.
This simple setting has the following characteristics. In the first region, the treatment effect is null, and observations from this region only add unnecessary noise to the tested subgroup. In the second region, the treatment effect is very positive, but observations from that region also have high variance (which can be seen as unobserved heterogeneity). Thus even though with many draws we would have clear evidence that observations from this region contribute a positive treatment effect, in a small sample treating this region runs the risk of adding more noise than signal and getting a statistically insignificant result. The final region has a modest treatment effect, but is much less noisy. Hence, including data from this final region should improve the power of a test. Moreover, we observe that these three regions operate in baseline outcome regimes that differ by orders of magnitude. Thus without correct centering of the outcomes, noise in the outcome can hurt the power of any of the methods.
In this setting, a successful method should not include the first region, only include the second region when the sample size is large enough, and definitely include the third region in the tested subset. Furthermore, centering the outcome should only improve power. Among the methods we analyze, only those optimized for the power of the test fulfill these requirements, while standard tools from policy assignment tend to include regions with excess noise. In our experiments, we therefore find that the significance tree algorithm with optimal centering evidenceres, which is built to trade off signal and noise effectively, significantly out-performs all other benchmarks for a large range of sample sizes. The results are presented in (ref).
We next present results from a semi-synthetic simulation based on a randomized trial on HIV prevention from duflo2019hiv. In the trial, youths in Kenya were randomized into one of four arms. We use one of the three treatment arms as our treatment along with the control arm, and select 2456 units that do not contain missing values in the treatment or outcome. For each unit, we have 44 observable characteristics, which we collect in a vector $X$. The outcome of interest is the prevalence of the herpes simplex virus 2 (HSV-2) two years after the treatment (binary outcome $Y \in \{0, 1\}$). Throughout, our simulation is informative about the performance of our method on data similar to that from duflo2019hiv, but it does not (and is not designed to) provide any new insight into that study.
We analyze the performance of our algorithm in a semi-synthetic experiment. For all 2456 units, we use the raw feature vector $X_i \in \mathbb{R}^{44}$ and treatment indicator $W_i \in \{0,1\}$ from the real data. We then train a $\ell_1$-penalized (sparse) logistic regression outcome model separately on the treated and untreated population in order to learn functions $f_w(x) = \mathop{}\!\textnormal{P}(Y=1| X=x, W=w) = \mathop{\textnormal{expit}}(\theta_w' x)$ for the logistic sigmoid function $\mathop{\textnormal{expit}}(z) = (1+\exp(-z))^{-1}$. Subsequently, we train un-penalized logistic regressions with the selected regressors. The learned coefficients are depicted in (ref). We see that there are three significant factors of heterogeneity, with two factors increasing the effect and one factor decreasing the effect. Next, we generate new outcome data from several manipulations of the learned models. In particular, we generate binary outcome data based on the probabilistic models
where $\overline{X}=\frac{1}{n} \sum_{i=1}^n X_i$. The constant “cov-strength” controls the signal/heterogeneity contained in the features, while the constant “separation” controls the baseline effect of the treatment.
Our experimental results are depicted in Figures (ref), (ref) and (ref). We find that our evidenceres and submod methods outperform the benchmarks in the medium signal regimes where the signal in the features is relatively low. Meanwhile, evidenceres and submod remain comparable with alternative methods in the two extreme regimes (very low/very high signal). Morever, the performance of both algorithms is comparable, with evidenceres having a slight advantage. We also see that running submodular minimization at the end of the significance tree has little to no effect. This implies not only that the greedy tree algorithm found good subgroups, but also that the assignment that it constructed in a step-by-step manner was appropriate for these subgroups.
The simulation in (ref) is based on binary outcomes. We also simulate a variant of the semi-synthetic experiment with continuous outcome data. Here we use the group structure from (ref), but define groups based on the real covariates from the HIV trial. Moreover, we use the treatment variable from the real data. The continuous outcome in each experiment is generated as follows:
Our experimental results are depicted in (ref). We find that our submod method outperforms the benchmarks, correctly identifying not to include high-variance subgroups, even though they might have a slightly positive effect. Moreover, the evidenceres method comes a close second, albeit in this setting we see that submod is strictly better.
In this section we consider two modifications that aim to improve the power of our empirical implementation. In (ref) we modify our target criterion, by generalizing our setup to permit weighted rather than binary selection into subgroups. In (ref) we then consider a forest-based extension of our significance-tree approach. We illustrate in (ref) how these modifications can improve power in a simple simulation.
In our baseline setup, we focused on obtaining a binary assignment $a: \mathcal{X} \rightarrow \{0,1\}$ that selects whether we assign an observation with covariate vector $x$ to our testing subgroup. We can relax this assumption and consider instead continuous weights $a: \mathcal{X} \rightarrow \mathbb{R}_{\geq 0}$, where our target is the expected lift $u(a) = \mathop{}\!\textnormal{E}[a(X) \: \tau(X)]$. Testing for $\mathop{}\!\textnormal{E}[a(X) \: \tau(X)] \geq 0$ can now be understood in two ways.
Firstly, consider the most general case where assignment weights can take any value in $\mathbb{R}_{\geq 0}$. This can be interpreted as a test of the null hypothesis that there is no subgroup with significant positive effects. Should we find significant continuous weights, this allows us to reject this null hypothesis. We can view the weights here as a technical tool for implementing this test. Likewise, subgroups with larger weights can be interpreted as contributing stronger evidence to our test. Note that this interpretation holds under any arbitrary re-scaling of the weights.
An alternative interpretation in the context of policy learning is possible if we impose the additional requirement $a: \mathcal{X} \rightarrow [0,1]$. This can be achieved by appropriate scaling of the weights; for instance, we could divide all weights by the largest weight. $a$ can now be understood as a probabilistic assignment policy that assigns observations to treatment with probability $a(X)$. We can then test whether this assignment policy passes a hypothesis test on a hold-out sample.
Even under this latter interpretation, note that we do not actually randomize, since the hold-out data is already given. Instead, we estimate the weighted estimand $u(\hat{a})$ without adding additional randomness, which would be statistically inefficient. Only if we were to deploy the assignment policy $\hat{a}(x)$ in the population to select who should get treated would we then randomize for an observation with covariate $x$ based on the probability $\hat{a}(x)$.
The oracle result in (ref) remains applicable for the continuous relaxation. Indeed, the continuous relaxation allows for the direct solution of optimal weights in the framework of (ref):
This particularly tractable formulation implies a criterion that can be used to partition the covariate space recursively using trees with binary splits as in (ref). Specifically, let $\bar{\tau}_\ell$ be the average treatment effect of a leaf and let $\bar{\sigma}^2_\ell = \mathop{}\!\textnormal{Var}(\hat{\tau}_\ell)$ be the variance of the corresponding estimator $\hat{\tau}_\ell$. Then we aim to find a partition that maximizes $ \sum_{\ell} \sfrac{\bar{\tau}^2_\ell}{\bar{\sigma}^2_\ell}, $ corresponding to weights $\bar{a}_\ell \propto \frac{\bar{\tau}_\ell}{\bar{\sigma}^2_\ell}$ on $\hat{\tau}_\ell$. When we also write $\hat{\sigma}^2_\ell$ for an estimator of $\bar{\sigma}^2_\ell$, this suggests maximizing the sample analog $ \sum_{\ell} \sfrac{\hat{\tau}^2_\ell}{\hat{\sigma}^2_\ell} $ over partitions of the data, which can be achieved recursively, and then setting $\hat{a}_\ell \propto \frac{\hat{\tau}_\ell}{\hat{\sigma}^2_\ell}$. Note here that partitions can be computed locally, in contrast to the setting with binary assignment weights. For illustration, (ref) compares a greedily fitted tree that maximizes the criterion in (ref) over binary assignments with a greedily fitted tree that recursively maximizes this weighted criterion.
There may be two distinct reasons why we may prefer the continuous relaxation over binary assignment weights. The first one is in terms of power. Binary assignment weights are a special case of continuous assignment weights. As such, we would expect the continuous formulation to perform weakly better than the binary version with regard to our testing goal (i.e. maximizing the probability of passing a hypothesis test on hold-out data). This continuous relaxation modifies the hypothesis test to examine whether a subgroup with significant positive effects could be found, even if we do not now find such a subgroup. By contrast, the binary assignment tests for a specific subgroup, and is thus more restrictive. The continuous formulation may lead to an improvement in power, though it has the drawback of not finding a specific subgroup. The second reason for preferring the continuous version is computational. The tree implementation for the continuous relaxation computes splits locally, whereas with binary assignments the splits have to be computed globally. Thus, the implementation of the continuous formulation is more computationally efficient.
As decision trees are typically of limited expressiveness and high variance, we could alternatively improve performance by bagging trees into a significance forest following the random-forest construction of breiman2001random. This procedure can be used both for the binary assignment version and the continuous relaxation from (ref). A potential implementation is as follows:
In this section, we compare the performance of the continuous relaxation and forest variants to the baseline implementation from (ref). To provide a comprehensive comparison, we include nine types of subgroup-assignment algorithms in this simulation experiment:
Unlike the comparison in (ref), we use the same centering strategy in training and hold-out for each of the methods. This means that those methods that do not employ centering also do not use centering for the test itself, thus treating training and hold-out consistently and highlighting the importance of residualization.
For this comparison, we employ a specific data-generating process that highlights the differences between these different implementations. Specifically, we assume that $X = (X_1, X_2)$ is uniformly distributed on $[-2, +2]^2$, and that the outcome is generated as
We set $a = 1, b = 0$, and $c= 0.2$. $d$ then controls the strength of the treatment effect and its heterogeneity. We run our simulations on four different settings for $d$:
As $X_1 \in [-2, +2]$, there are several regions where the treatment effect is negative, and indeed the average treatment effect is zero. This captures in a stylized sense our motivating example of an RCT where overall effects are small in a noisy environment. The extent of this noise is governed by $X_2$, with larger values of $X_2$ corresponding to noisier regions.
For each treatment setting we run 1000 iterations of our simulation. Within each iteration, we draw separate training and test datasets based on the above data-generating process. Each such dataset is of size 400. We fit all the assignment policies on the training dataset, and then evaluate their respective performances by calculating their corresponding $p$-values on the test dataset.
For subgroup-assignment algorithms with centered outcomes, we compute the optimal centering function with a random forest regressor. We train one version of the centering function on the whole training dataset, which is then used to center outcomes in the hold-out test dataset. We also train another version of the centering function using cross-fitting in the training dataset: this is used to center outcomes within the training dataset when constructing our subgroup-assignment algorithms.
When using single significance trees, we control complexity by setting the minimum leaf size and maximum tree depth to 10. When aggregating these into significance forests, we permitted greater complexity in the individual trees: we set each tree to now have a minimum leaf size of 1 per arm, and did not impose a maximum tree depth.
Our results are presented in (ref) and (ref). As expected, we see that assigning everyone to treatment does relatively poorly, but our proposed algorithms deliver substantial improvements. Our results also demonstrate the performance gains in this setting from using residualized outcomes; using significance forests rather than single trees; and using continuous rather than binary assignment weights.
We now further investigate the difference between these methods. To represent our subgroup-assignment algorithms graphically, we focus on a single iteration of the simulations. That is, we take one draw of a training and test dataset, with each dataset of sample size 400. We then fit all the algorithms to the training dataset. Finally, we plot the corresponding assignment weights on the test dataset. We use the medium treatment regime with $d=1$. The resulting binary assignment algorithms are shown in (ref), and continuous assignment algorithms are shown in (ref). Across methods, the algorithms tend to select areas where $X_1$ is large but where $X_2$ is not too large. This is as we would expect, since this region should correspond to a good ratio of treatment effect to noise. While simple trees tend to only differentiate along the $X_1$ axis (along which the treatment effect varies), the forest-based implementations can capture variation along $X_2$ (noise variation) as well, effectively trading off both.
\FloatBarrier
To wrap up, we discuss natural extensions of our framework to different objectives and settings.
In our baseline setup, we focus on finding evidence of subgroups with positive treatment effects. We can use a similar approach to finding evidence of heterogeneous treatment effects. This extension is similar to how Crump2008-cr and Armstrong2015-gc extend tests for non-constant and positive effects to tests for heterogeneous effects. Specifically, we now consider an assignment
where units are assigned to a high-effect ($+1$), neutral ($0$), and low-effect ($-1$) group. We then evaluate whether the average treatment effects in the high group are above that of the low group by a hypothesis test. Towards this end, we consider the estimands
where $a_- = (- a)_+ \geq 0$ denotes the negative part. Both $\Delta_1(a) >0$ and $\Delta_2(a) > 0$ provide evidence of treatment-effect heterogeneity. We focus here on collecting evidence that $\Delta_2(a) > 0$, which itself is maximized for $a^*(x) = \text{sign}(\tau(X) {-} \mathop{}\!\textnormal{E}[\tau(X)]))$. This test is similar to the regression-based calibration test employed by chernozhukov2018generic to evaluate treatment-effect heterogeneity.
A test for $\Delta_2(a) > 0$ is based on $ \widehat{Z}(a) = \frac{\hat{\Delta}_2(a)}{\sqrt{\widehat{\mathop{}\!\textnormal{Var}}(\hat{\Delta}_2(a))}} $ where we set
We now consider a local-to-zero approximation for which $ \tau(x) = \bar{\tau} + \frac{\mu(x)}{\sqrt{N}}. $ Under regularity conditions,
and we can adapt our above approach to maximize
either over discrete assignments $a: \mathcal{X} \rightarrow \{-1,0,+1\}$ or the continuous relaxation $a: \mathcal{X} \rightarrow \mathbb{R}$.
The testing goal ignores the actual outcome associated with our chosen subgroup. We can instead consider a second goal that also incorporates the expected outcome subject to passing the test, namely
This goal would correspond for example to a researcher wanting to find subgroups whose treatment effects are not only statistically significant, but also large in magnitude. It mixes the expected outcome $u(\hat{a})$ with the testing goal $\mathop{}\!\textnormal{E}[\widehat{T}(\hat{a})|\hat{a}]$ to trade off the value of the subgroup's outcomes with how easy it is to pass the test. In the approximation of (ref), the optimal subgroup assignment would then solve
In the main article, we considered a treatment effect that is identified in an experiment by $\tau(x) = \mathop{}\!\textnormal{E}\left[\tilde{Y}\middle|X=x\right] = \mathop{}\!\textnormal{E}\left[\frac{W - p}{p(1-p)} (Y-c(X))\middle|X=x\right]$ and can be estimated accordingly. We can extend the same techniques to a family of parameters $\tau(x)$ defined via the linear conditional moment equation \[ \tau(x) = \mathop{}\!\textnormal{E}[\psi(Z; h)| X=x] \] for data $Z$, target causal parameter $\tau(x)$, and nuisance functions $h$. For $\tau(x)$ defined in this way, we aim to obtain an assignment that provides a significant test against the null that $u(a)=\mathop{}\!\textnormal{E}[a(X) \: \tau(X)] \leq 0$. If the moment function satisfies the conditional Neyman orthogonality property chernozhukov2018double,foster2019orthogonal,oprescu2018orthogonal, and these nuisance estimates are estimated in a cross-fitting manner and are $o(n^{1/4})$ estimable, then their error is asymptotically negligible. In this case, all our algorithms can be directly generalized by simply renaming the pseudo-outcome to \[ \tilde{Y}_i = \psi(Z_i;\hat{h}). \] The estimation and optimization algorithms as well as the analysis of their finite-sample performance easily generalizes to this general pseudo-outcome setting. Moreover, Neyman orthogonality for this setting can be easily achieved via an automated-double machine learning process as outlined in the recent work of chernozhukov2018learning,chernozhukov2020adversarial.
As an example, we can generalize the case of conditional average treatment effects to cases where the propensity score is unknown and has to be estimated, in which case $h(w,x) = (c(w, x), p(x))$, with $c(w,x) = \mathop{}\!\textnormal{E}[Y| W=w, X=x]$, $p(x)=\mathop{}\!\textnormal{E}[W| X=x]$ and $\psi(Y,W,X; (c(\cdot, \cdot),p(\cdot))) = c(1,X)-c(0,X) + (Y-c(W,X)) \left(\frac{W}{p(X)} - \frac{1-W}{1-p(X)}\right)$. Even more straightforwardly, we can also consider the simple prediction setup where we measure the treatment effect (or outcome of interest) $Y$ directly, in which case we write $\tau(x) = \mathop{}\!\textnormal{E}[Y|X=x]$. Here, $\psi(y,x) = Y$.
Our results generalize to testing effects relative to alternative null assignments other than assigning all observations to control. In the case of a reference assignment $b(X)$,
In that case, all results go through if we relabel $Y_b^{(w)} = Y^{(|w-b(x)|)}$ for observations with $X=x$.
So far, we have focussed on separating the training of a subgroup-assignment algorithm from its evaluation. While this may be the most applicable case in many heavily regulated settings, we may gain power through re-using training and hold-out samples by switching their roles.
In this article, we address the issue of obtaining reliable evidence about an intervention of interest, when overall estimated effects from RCT data are small or noisy. To do so, we propose new machine-learning tools that are optimized for finding subgroups with statistically significant treatment effects. We first formalize the objective of maximizing the probability of passing a hypothesis test on additional hold-out data. We then provide empirical implementations based on decision trees and submodular minimization. Finally, we document the performance of these empirical implementations in a set of simulation studies, discuss modifications, and point to practically relevant extensions.