EconBase
← Back to paper

Stratification Trees for Adaptive Randomization in Randomized Controlled Trials

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.

94,361 characters · 13 sections · 79 citation commands

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

Stratification Trees for Adaptive Randomization in Randomized Controlled Trials

abstractThis paper proposes an adaptive randomization procedure for two-stage randomized controlled trials. The method uses data from a first-wave experiment in order to determine how to stratify in a second wave of the experiment, where the objective is to minimize the variance of an estimator for the average treatment effect (ATE). We consider selection from a class of stratified randomization procedures which we call stratification trees: these are procedures whose strata can be represented as decision trees, with differing treatment assignment probabilities across strata. By using the first wave to estimate a stratification tree, we simultaneously select which covariates to use for stratification, how to stratify over these covariates, and the assignment probabilities within these strata. Our main result shows that using this randomization procedure with an appropriate estimator results in an asymptotic variance which is minimal in the class of stratification trees. Moreover, our results are able to accommodate a large class of assignment mechanisms within strata, including stratified block randomization. In a simulation study, we find that our method, paired with an appropriate cross-validation procedure, can improve on ad-hoc choices of stratification. We conclude by applying our method to the study in karlan2017, where we estimate stratification trees using the first wave of their experiment.

KEYWORDS: randomized experiments; decision trees; adaptive randomization\\ JEL classification codes: C14, C21, C93

Introduction

This paper proposes an adaptive randomization procedure for two-stage randomized controlled trials (RCTs). The method uses data from a first-wave experiment in order to determine how to stratify in a second wave of the experiment, where the objective is to minimize the variance of an estimator for the average treatment effect (ATE). We consider selection from a class of stratified randomization procedures which we call stratification trees: these are procedures whose strata can be represented as decision trees, with differing treatment assignment probabilities across strata.

Stratified randomization is ubiquitous in randomized experiments. In stratified randomization, the space of available covariates is partitioned into finitely many categories (i.e. strata), and randomization to treatment is performed independently across strata. Stratification has the ability to decrease the variance of estimators for the ATE through two parallel channels. The first channel is from ruling out treatment assignments which are potentially uninformative for estimating the ATE. For example, if we have information on the sex of individuals in our sample, and average outcomes vary with sex, then performing stratified randomization over this characteristic can reduce variance (we present an example of this for the standard difference-in-means estimator in Appendix (ref)). The second channel through which stratification can decrease variance is by allowing for differential treatment assignment probabilities across strata. For example, if we again consider the setting where we have information on sex, then it could be the case that for males the outcome under one treatment varies much more than under the other treatment. As we show in Section (ref), this can be exploited to reduce variance by assigning treatment according to the Neyman Allocation, which in this example would assign more males to the more variable treatment. Our proposed method leverages insights from supervised machine-learning to exploit both of these channels, by simultaneously selecting which covariates to use for stratification, how to stratify over these covariates, as well as the optimal assignment probabilities within these strata, in order to minimize the variance of an estimator for the ATE.

Our main result shows that using our procedure results in an “optimal" (to be made precise later) stratification of the covariate space, where we restrict ourselves to stratification in a class of decision trees. A decision tree partitions the covariate space such that the resulting partition can be interpreted through a series of yes or no questions (see Section (ref) for a formal definition and some examples). We focus on strata formed by decision trees for several reasons. First, since the resulting partition can be represented as a series of yes or no questions, it is easy to communicate and interpret, even with many covariates. This feature could be particularly important in many economic applications, because many RCTs in economics are undertaken in partnership with external organizations (for example, many RCTs described in karlan2016 were undertaken in this way), and thus clear communication of the experimental design could be crucial. Second, as we explain in Section (ref), using partitions based on decision trees gives us theoretical and computational tractability. Third, as we explain in Section (ref), using decision trees allows us to flexibly address the additional goal of minimizing the variance of estimators for subgroup-specific effects. Lastly, decision trees naturally encompass the type of stratifications usually implemented by practitioners. The use of decision trees in statistics and machine learning goes back at least to the work of Breiman breiman1984, gyorfi1996, and has seen a recent resurgence in econometrics athey2016, athey2017.

An important feature of our theoretical results is that we allow for the possibility of so-called restricted randomization procedures within strata. Restricted randomization procedures limit the set of potential treatment allocations, in order to force the true treatment assignment proportions to be close to the desired target proportions efron1971, wei1978, antognini2004, kuznetsova2011, zelen1974. Restricted randomization induces dependence in the assignments within strata, which complicates the analysis of our procedure. By extending techniques recently developed in bugni2017, our results accommodate a large class of restricted randomization schemes, including stratified block randomization, which as we discuss in Example (ref) is a popular method of randomization.

Although our main focus is on increasing efficiency, stratified randomization has additional practical benefits beyond reducing the variance of ATE estimators. For example, when a researcher wants to analyze subgroup-specific effects, stratifying on these subgroups serves as a form of pre-analysis registration. To that end, we also present results on how to extend our procedure when targeting subgroup-specific effects.

The literature on design and inference in RCTs is vast (references in athey2017survey, cox2000, glennerster2013, pukelsheim2006, rosenberger2015, and from a Bayesian perspective, ryan2016, provide an overview). The classical literature on optimal randomization, going back to the work of smith1918 silvey2013, maintains a parametric relationship for the outcomes with respect to the covariates, and targets efficient estimation of the model parameters. In contrast, our paper follows a recent literature which instead maintains a non-parametric model of potential outcomes, and targets efficient estimation of treatment effect parameters. This recent literature can be broadly divided into “one-stage" procedures, which do not use prior data on all experimental outcomes to determine how to assign treatment kallus2013, kasy2016, barrios2014,aufenanger2017,quistorff2020, and “multi-stage" procedures, of which our method is an example. Multi-stage procedures use prior data on the experimental outcomes to determine how to randomize. For example, they may use response information from previous experimental waves to determine how to randomize in subsequent waves of the experiment. We will call these procedures response-adaptive. Although response adaptive methods typically require information from a prior experiment, such settings do arise in economic applications. First, many social experiments have a pilot phase or multi-stage structure. For example, simester2006, karlan2008, and karlan2017 all feature a multi-stage structure, and karlan2016 advocate the use of pilot experiments to help avoid potential implementation failures when scaling up to the main study. Second, many research areas have seen a profusion of related work which could be used as a first wave of data in a response-adaptive procedure hahn2011. The study of response-adaptive methods to inform many aspects of experimental design, including how to randomize, has a long history in the literature on clinical trials, both from a frequentist and Bayesian perspective hu2006,sverdlov2015,cheng2003, as well as in the literature on bandit problems bubeck2012.

Three papers which propose response-adaptive randomization methods in a framework similar to ours are hahn2011, chambaz2014 and bai2019 viviano2020. hahn2011 develop a procedure which uses the information from a first-wave experiment to compute the propensity score that minimizes the asymptotic variance of an ATE estimator, over a discrete set of covariates (i.e. they stratify the covariate space ex-ante). They then use the resulting propensity score to assign treatment in a second-wave experiment. In contrast, our method computes the optimal assignment proportions over a data-driven discretization of the covariate space. chambaz2014 propose a multi-stage procedure which uses data from previous experimental waves to compute an optimal propensity score, where the propensity score is constrained through entropy restrictions. However, their method requires the selection of several tuning parameters, as well as additional regularity conditions, and their optimal target depends on these features in a way that may be hard to assess in practice. Their results are also derived in a framework where the number of experimental waves goes to infinity, which may not be an appropriate framework for many settings encountered in economics. Moreover, the results in both hahn2011 and chambaz2014 assume that assignment was performed completely independently across individuals in a given wave. In contrast, we reiterate that our results will accommodate a large class of stratified randomization schemes. bai2019 derives the MSE-optimal blocking of an experimental sample for the difference-in-means estimator, given a fixed assignment proportion, and shows that this blocking takes the form of a “matched-pairs" style design. He then proposes procedures which use information from a first-wave experiment to approximate the optimal blocking in a second-wave experiment. In contrast, in this paper we maintain an asymptotic framework where the number of strata is fixed, which precludes the type of blocking designs considered in bai2019. However, as explained in bai2019, it is in fact possible to combine his procedure with the one proposed in this paper, by implementing his optimal blocking within each stratum produced by our method. See Remark (ref) for further discussion.

The paper proceeds as follows: In Section (ref), we provide a motivating discussion, set up the notation, and formally define the set of randomization procedures we consider. In Section (ref), we present the formal results underlying the method as well as several relevant extensions. In Section (ref), we perform a simulation study to assess the performance of our method in finite samples. In Section (ref), we consider an application to the study in karlan2017, where we estimate stratification trees using the first wave of their experiment. Section (ref) concludes.

Preliminaries

In this section we discuss some preliminary concepts and definitions. Section (ref) presents a series of simplified examples which we use to motivate our procedure. Section (ref) establishes some notation and provides the definition of a stratification tree, as well as our notion of a randomization procedure.

Motivating Discussion

We present a series of simplified examples which we use to motivate our proposed method. First we study the problem of optimal experimental assignment without covariates. We work in a standard potential outcomes framework: let $(Y(1),Y(0))$ be potential outcomes for a binary treatment $A \in \{0, 1\}$, and let the observed outcome $Y$ for an individual be defined as

equation[equation omitted — 53 chars of source]

Let $$\mu_a := E[Y(a)], \sigma^2_a := \text{Var}(Y(a))~,$$ for $a \in \{0, 1\}$. Our quantity of interest is the average treatment effect $$\theta := \mu_1 - \mu_0~.$$ Suppose we perform an experiment to obtain a size $n$ sample $\{(Y_i, A_i)\}_{i=1}^n$, where the sampling process is determined by $\{(Y_i(1),Y_i(0))\}_{i=1}^n$, which are i.i.d, and the treatment assignments $\{A_i\}_{i=1}^n$, where exactly $n_1 := \sum_{i=1}^nA_i = \lfloor{n\pi\rfloor}$ individuals are randomly assigned to treatment $A = 1$, for some $\pi \in (0,1)$ (however, we emphasize that our results will accommodate other methods of randomization). Given this sample, consider estimation of $\theta$ through the standard difference-in-means estimator: $$\hat{\theta}^{(1)} := \frac{1}{n_1}\sum_{i=1}^{n}Y_iA_i - \frac{1}{n - n_1}\sum_{i=1}^nY_i(1-A_i)~.$$ It can then be shown that $$\sqrt{n}(\hat{\theta}^{(1)} - \theta) \xrightarrow{d} N(\theta, V^{(1)})~,$$ where $$V^{(1)} := \frac{\sigma_1^2}{\pi} + \frac{\sigma_0^2}{1 - \pi}~.$$ Our goal is to choose $\pi$ to minimize the variance of $\hat{\theta}$. Solving this optimization problem yields the following solution: $$\pi^* := \frac{\sigma_1}{\sigma_1 + \sigma_0}~.$$ This allocation, known as the Neyman Allocation, assigns more individuals to the treatment which is more variable. Note that when $\sigma_0^2 = \sigma_1^2$, so that the variances of the potential outcomes are equal, the optimal proportion is $\pi^* = 0.5$, which corresponds to a standard equal treatment allocation. In general, implementing $\pi^*$ is infeasible without knowledge of $\sigma^2_0$ and $\sigma^2_1$. In light of this, if we had prior data $\{(Y_j, A_j)\}_{j=1}^m$ which allowed us to estimate $\sigma^2_0$ and $\sigma^2_1$, then we could use this data to estimate $\pi^*$, and then use this estimate to assign treatment in a subsequent wave of the study. The idea of sequentially updating estimates of unknown population quantities using past observations, in order to inform experimental design in subsequent stages, underlies many procedures developed in the literatures on response adaptive experiments and bandit problems, and is the main idea underpinning our proposed method.

remarkAlthough the Neyman Allocation minimizes the variance of the difference-in-means estimator, it is entirely agnostic on the welfare of the individuals in the experiment itself. In particular, the Neyman Allocation could assign the majority of individuals in the experiment to the inferior treatment if that treatment has a much larger variance in outcomes (see hu2006 for relevant literature in the context of clinical trials, as well as narita2018 for recent work on this issue in econometrics). While this feature of the Neyman Allocation may introduce ethical or logistical issues in some relevant applications, in this paper we focus exclusively on the problem of estimating the ATE as accurately as possible.

Next we repeat the above exercise with the addition of a discrete covariate $S \in \{1, 2, ... , K\}$ over which we stratify. We perform an experiment which produces a sample $\{(Y_i, A_i, S_i)\}_{i=1}^n$, where the sampling process is determined by i.i.d draws $\{(Y_i(1), Y_i(0), S_i)\}_{i=1}^n$ and the treatment assignments $\{A_i\}_{i=1}^n$. For this example suppose that the $\{A_i\}_{i=1}^n$ are generated as follows: for each $k$, exactly $n_1(k) := \sum_{i=1}^n{\bf 1}\{S_i = k\}A_i = \lfloor{n(k)\pi(k)\rfloor}$ individuals are randomly assigned to treatment $A = 1$, with $n(k) := \sum_{i=1}^n{\bf 1}\{S_i = k\}$.

Note that when the assignment proportions $\pi(k)$ are not equal across strata, the difference-in-means estimator $\hat{\theta}^{(1)}$ is no longer consistent for $\theta$. Hence we consider the following weighted estimator of $\theta$: $$\hat{\theta}^{(2)} := \sum_k\frac{n(k)}{n}\hat{\theta}(k)~,$$ where $\hat{\theta}(k)$ is the difference-in-means estimator for $S = k$: $$\hat{\theta}(k) := \frac{1}{n_1(k)}\sum_{i=1}^{n}Y_iA_i{\bf 1}\{S_i = k\} - \frac{1}{n(k) - n_1(k)}\sum_{i=1}^nY_i(1-A_i){\bf 1}\{S_i = k\}~.$$ In words, $\hat{\theta}^{(2)}$ is obtained by computing the difference in means for each $k$ and then taking a weighted average over each of these estimates. Note that when $K = 1$ (i.e. when $S$ can take only one value), this estimator simplifies to the difference-in-means estimator. It can be shown under appropriate conditions that $$\sqrt{n}(\hat{\theta}^{(2)} - \theta) \xrightarrow{d} N(0, V^{(2)})~,$$ where $$V^{(2)} := \sum_{k=1}^K P(S = k) \left[ \left(\frac{\sigma_0^2(k)}{1-\pi(k)} + \frac{\sigma_1^2(k)}{\pi(k)}\right) + \left(E[Y(1) - Y(0)|S = k] - E[Y(1)-Y(0)]\right)^2\right]~,$$ with $\sigma^2_a(k) := E[Y(a)^2|S = k] - E[Y(a)|S = k]^2$. The first term in $V^{(2)}$ is the weighted average of the conditional variances of the difference in means estimator for each $S = k$. The second term in $V^{(2)}$ arises due to the additional variability in sample sizes for each $S = k$. We note that this variance takes the form of the semi-parametric efficiency bound derived by hahn1998 for estimators of the ATE which use the covariate $S$. Following a similar logic to what was proposed above without covariates, we could use first-wave data $\{(Y_j, A_j, S_j)\}_{j=1}^m$ to form a sample analog of $V^{(2)}$, and choose $\{\pi^*(k)\}_{k=1}^K$ to minimize this quantity.

Now we introduce the setting that we consider in this paper: suppose we observe covariates $X \in \mathcal{X} \subset \mathbb{R}^d$, so that our covariate space is now multi-dimensional with potentially continuous components. How could we practically extend the logic of the previous examples to this setting? A natural solution is to discretize (i.e. stratify) $\mathcal{X}$ into $K$ categories (strata), by specifying a mapping $S:\mathcal{X} \rightarrow \{1, 2, 3, ..., K\}$, with $S_i := S(X_i)$, and then proceed as in the above example. As we argued in the introduction, stratified randomization is a popular technique in practice, and possesses several attractive theoretical and practical properties. In this paper we propose a method which uses first-wave data to estimate (1) the optimal stratification, and (2) the optimal assignment proportions within these strata. In other words, given first-wave data $\{(Y_j, A_j, X_j)\}_{j=1}^m$, where $X \in \mathcal{X} \subset \mathbb{R}^d$, we propose a method which selects $\{\pi(k)\}_{k=1}^K$ and the function $S(\cdot)$, in order to minimize the variance of our estimator $\hat{\theta}^{(2)}$. In particular, our proposed solution selects a randomization procedure amongst the class of what we call stratification trees, which we introduce in the next section.

remarkOur focus on the minimization of asymptotic variance is motivated by standard asymptotic optimality results for regular estimators van1998, and our procedure directly impacts the size of confidence sets and the power of tests constructed using our normal approximation. However, accurate estimation of the ATE is not the only objective one could consider when designing an RCT. For example, we could instead consider designing the RCT with the ultimate goal of finding a treatment allocation which maximizes population welfare (see for example manski2004, kitagawa2017 and references therein). Some recent work kasy2021 has focused on adaptive assignment mechanisms with this objective in mind. To what extent these objectives can be incorporated into our framework would be an interesting direction for future work.

Notation and Definitions

In this section we establish our notation and define the class of randomization procedures that we will consider. As in Section (ref) let $(Y(1), Y(0))$ be potential outcomes for a binary treatment $A \in \{0,1\}$ and let $X \in \mathcal{X} \subset \mathbb{R}^d$ denote a vector of observed pre-treatment covariates. Let $Q$ denote the distribution of $(Y(1), Y(0), X)$. Throughout the paper we assume that all of our observations are generated by i.i.d draws from $Q$. We restrict $Q$ as follows:

assumptionQ satisfies the following properties: \begin{itemize}[topsep = 1pt] • $Y(a) \in [-M, M]$ for some $M < \infty$, for $a \in \{0, 1\}$. • $X \in \mathcal{X} = \bigtimes_{j=1}^d [b_j, c_j]$, for some $\{b_j, c_j\}_{j=1}^d$ finite. • $X = (X_C, X_D)$, where $X_C \in \mathbb{R}^{d_1}$ for some $d_1 \in \{0, 1, 2, ..., d\}$ is continuously distributed with a bounded, strictly positive density. $X_D \in \mathbb{R}^{d - d_1}$ is discretely distributed with finite support. \end{itemize}
remarkThe assumptions imposed on $(Y(1),Y(0),X)$ in Assumption (ref) are used frequently throughout the proofs of our results. However, they may be stronger than desirable in some applications. For example, the assumption that $X$ has only discrete or continuous components which are supported on a rectangle may fail in certain practical examples (see for example the set of covariates considered in Section (ref)). However, in the simulations presented in Section (ref) and Appendix (ref) we consider designs where $Y(a)$ has unbounded support and $X$ is not supported on a rectangle, and these results suggest that Assumption (ref) could be reasonably weakened. We further note that although a user does not need to specify a choice of $M$ to implement our procedure, they are required to specify a choice of $\mathcal{X}$. We illustrate this point in the application presented in Section (ref).

Our quantity of interest is the average treatment effect (ATE) given by: $$\theta := E[Y_i(1) - Y_i(0)]~.$$

An experiment on an i.i.d sample $\{(Y_i(1), Y_i(0), X_i)\}_{i=1}^n$ produces the following data: $$\{W_i\}_{i=1}^n:= \{(Y_i, A_i, X_i)\}_{i=1}^n~,$$ whose joint distribution is determined by $Q$, the potential outcomes expression ((ref)), and the randomization procedure which generates $\{A_i\}_{i=1}^n$. We focus on the class of stratified randomization procedures: these randomization procedures first stratify according to baseline covariates and then assign treatment status independently across each of these strata. However, we attempt to make minimal assumptions on the exact specification of the randomization procedure, and in particular we do not require the treatment assignment within each stratum to be independent across observations.

We will now describe the structure we impose on the class of possible strata we consider. For $L$ a positive integer, let $K = 2^L$ and let $[K] := \{1, 2, ... ,K\}$. Consider a function $S: \mathcal{X} \rightarrow [K]$, then $\{S^{-1}(k)\}_{k=1}^K$ forms a partition of $\mathcal{X}$ with $K$ strata. For a given positive integer $L$, we work in the class $S(\cdot) \in \mathcal{S}_L$ of functions whose partitions form tree partitions of depth $L$ on $\mathcal{X}$, which we now define. Our definition is recursive, so we begin with the definition for a tree partition of depth one:

definitionLet $\mathcal{X} = \bigtimes_{j=1}^d [b_j, c_j]$, and let $x = (x_1, x_2, ..., x_d) \in \mathcal{X}$. A tree partition of depth one on $\mathcal{X}$ is a partition $\{\mathcal{X}_D(j,\gamma), \mathcal{X}_U(j,\gamma)\}$ of $\mathcal{X}$, where $$\mathcal{X}_D(j,\gamma) := \{x \in \mathcal{X}: x_j \le \gamma\}~,$$ $$\mathcal{X}_U(j,\gamma) := \{x \in \mathcal{X}: x_j > \gamma\}~,$$ for some $j \in [d]$ and $\gamma \in (b_j, c_j)$. We call $\mathcal{X}_D(j,\gamma)$ and $\mathcal{X}_U(j,\gamma)$ leaves (or sometimes terminal nodes).
exampleFigure (ref) presents two different representations of a tree partition of depth one on $[0,1]^2$. The first representation we call graphical: it depicts the partition on a square drawn in the plane. The second depiction we call a tree representation: it illustrates how to describe a depth one tree partition as a yes or no question. In this case, the question is “is $x_1$ less than or greater than 0.5?".
figure[figure omitted — 1,063 chars of source]

Next we define a tree partition of depth $L > 1$ recursively:

definitionA tree partition of depth $L > 1$ on $\mathcal{X} = \bigtimes_{j=1}^d [b_j, c_j]$ is a partition $\{\mathcal{X}^{(L-1)}_D(j, \gamma), \mathcal{X}^{(L-1)}_U(j, \gamma)\}$ of $\mathcal{X}$, where $$\mathcal{X}^{(L-1)}_D(j, \gamma) \hspace{2mm} \text{is a tree partition of depth $L - 1$ on $\mathcal{X}_D(j, \gamma)$}~,$$ $$\mathcal{X}^{(L-1)}_U(j, \gamma) \hspace{2mm} \text{is a tree partition of depth $L - 1$ on $\mathcal{X}_U(j, \gamma)$}~,$$ for some $j \in [d]$ and $\gamma \in (b_j, c_j)$. We call $\mathcal{X}^{(L-1)}_D$ and $\mathcal{X}^{(L-1)}_U$ left and right subtrees, respectively.
exampleFigure (ref) depicts two representations of a tree partition of depth two on $[0,1]^2$.
figure[figure omitted — 1,673 chars of source]

We focus on strata that form tree partitions for several reasons. First, relative to more “flexible" partitions of $\mathcal{X}$, these types of strata are easy to represent and interpret, especially in higher dimensions, via their tree representations or as a series of yes or no questions. We argued in the introduction that this could be of particular importance in economic applications. Second, as we explain in Remarks (ref) and (ref), restricting ourselves to tree partitions helps with computational and theoretical tractability. In particular, computing an optimal stratification function is a difficult discrete optimization problem, but restricting ourselves to tree partitions allows us to employ an effective search heuristic known as an evolutionary algorithm. Third, the recursive aspect of tree partitions makes the targeting of subgroup-specific effects convenient, as we explain in Section (ref).

For each $k \in [K]$, define $\pi := (\pi(k))_{k=1}^K$ to be the vector of target proportions of units assigned to treatment $1$ in each stratum. A stratification tree is a pair $(S,\pi)$, where $S(\cdot)$ forms a tree partition, and $\pi$ specifies the target proportions in each stratum. We denote the set of stratification trees of depth $L$ as $\mathcal{T}_L$.

remarkTo be precise, any element $T = (S, \pi) \in \mathcal{T}_L$ is equivalent to another element $T' = (S', \pi') \in \mathcal{T}_L$ whenever $T'$ can be realized as a re-labeling of $T$. For instance, if we consider Example (ref) with the labels $1$ and $2$ reversed, the resulting tree is identical to the original except for this re-labeling. $\mathcal{T}_L$ should be understood as the quotient set that results from this equivalence.
exampleFigure (ref) depicts a representation of a stratification tree of depth two. Note that the terminal nodes of the tree have been replaced with labels that specify the target proportions in each stratum.
figure[figure omitted — 941 chars of source]

Note that in our definition, a stratification tree of depth $L$ has exactly $K = 2^L$ leaves, so that our trees are “perfectly balanced". For the optimization problem we consider in our paper this is without loss of generality, because it can be shown bugni2021 that the asymptotic variance which results from a given stratification tree can be made weakly smaller by taking any leaf and subdividing it arbitrarily into two leaves with the same assignment target as their parent.

We further impose that the set of trees cannot have arbitrarily small cells, nor can they have arbitrarily extreme treatment assignment targets:

assumptionWe constrain the set of stratification trees $T = (S, \pi) \in \mathcal{T}_L$ such that, for some fixed $\nu > 0$ and $\delta > 0$, $\pi(k) \in [\nu, 1- \nu]$ and $P(S(X)=k) > \delta$.

We impose the restriction in Assumption (ref) to ensure, for example, that $1/\pi(k)$ is uniformly bounded for all $T \in \mathcal{T}_L$. For technical reasons relating to the potential non-measurability of our estimator, we will impose one additional restriction on $\mathcal{T}_L$.

assumptionLet $\mathcal{T}_L^{\dagger} \subset \mathcal{T}_L$ be a countable, closed subset of the set of stratification trees\footnotemark[1]. We then consider the set of stratification trees restricted to this subset. By an abuse of notation, we continue to denote the set of stratification trees we will consider as $\mathcal{T}_L$.
remarkWe emphasize that this assumption is only used as a sufficient condition to guarantee measurability, in order to invoke Fubini's theorem. Note that, in practice, restricting the set of stratification trees to those constructed from a finite grid satisfies Assumption (ref). However, our results will also apply more generally.

\footnotetext[1]{Here “closed" is with respect to an appropriate topology on $\mathcal{T}_L$, see Appendix (ref) for details.}

For each $T \in \mathcal{T}_L$, and given an i.i.d sample $\{(Y_i(0),Y_i(1), X_i)\}_{i=1}^n$ of size $n$, an experimental assignment is described by a random vector $(A_i(T))_{i=1}^n$. For our purposes a randomization procedure (or randomization scheme) is a family of such random vectors indexed by $T = (S, \pi) \in \mathcal{T}_L$. For $T = (S, \pi)$, let $S_i:= S(X_i)$, $(S_i)_{i=1}^n$ be the random vector of stratification labels of the observed data. We impose two assumptions on the randomization procedure $(A_i(T))_{i=1}^n$.

First, we require the following exogeneity assumption:

assumptionThe randomization procedure is such that, for each $T = (S, \pi) \in \mathcal{T}_L$, $$\left[\{(Y_i(0), Y_i(1), X_i)\}_{i=1}^n \perp \!\!\! \perp (A_i(T))_{i=1}^n \right]\bigg| (S_i)_{i=1}^n~.$$

This assumption asserts that the randomization procedure can depend on the observables only through the strata labels. Next, let $p(k;T) := P(S_i = k)$ be the population proportions of each stratum, then we also require that the randomization procedure satisfy the following “consistency" property:

assumptionThe randomization procedure is such that $$\sup_{T \in \mathcal{T}_L} \left|\frac{n_1(k;T)}{n} - \pi(k)p(k;T)\right| \xrightarrow{p} 0~,$$ for each $k \in [K]$, where $$n_1(k;T) := \sum_{i=1}^n{\bf 1}\{A_i(T) = 1, S_i = k\}~.$$

This assumption asserts that the assignment procedure must approach the target proportion asymptotically, and do so in a uniform sense over all stratification trees in $\mathcal{T}_L$.

Other than Assumptions (ref) and (ref), we do not require any additional assumptions about how assignment is performed. Examples (ref) and (ref) illustrate two randomization schemes which satisfy these assumptions and are popular in economics. bugni2017 make similar assumptions for a fixed stratification and show that they are satisfied for a wide range of assignment procedures, including procedures often considered in the literature on clinical trials: see for example efron1971, wei1978, antognini2004, and kuznetsova2011. In Proposition (ref) below, we verify that Assumptions (ref) and (ref) hold for stratified block randomization (see Example (ref)), which is a common assignment procedure in economic applications.

exampleSimple random assignment assigns each individual within stratum $k$ to treatment via a coin-flip with weight $\pi(k)$. Formally, for each $T$, $(A_i(T))_{i=1}^n$ is a vector with independent components such that $$P(A_i(T) = 1|S_i = k) = \pi(k)~.$$ Simple random assignment is theoretically convenient, and features prominently in papers on adaptive randomization. However, it is considered unattractive in practice because it results in a “noisy" assignment for a given target $\pi(k)$, and hence could be far off the target assignment for any given random draw. Moreover, this extra noise increases the finite-sample variance of ATE estimators relative to other assignment procedures which target $\pi(k)$ more directly kasy2013.
exampleStratified block randomization (SBR) assigns a fixed proportion $\pi(k)$ of individuals within stratum $k$ to treatment $1$. Formally, let $n(k)$ be the number of units in stratum $k$, and let $n_1(k)$ be the number of units assigned to treatment 1 in stratum $k$. In SBR, $n_1(k)$ is given by $$n_1(k) = \lfloor n(k) \pi(k) \rfloor~.$$ SBR proceeds by randomly assigning $n_1(k)$ units to treatment $1$ for each $k$, where all $${n(k) \choose n_1(k)}~,$$ possible assignments are equally likely. This assignment procedure has the attractive feature that it targets the proportion $\pi(k)$ as directly as possible. An early discussion of SBR can be found in zelen1974.

We conclude this section by showing that Assumptions (ref) and (ref) are satisfied by SBR:

propositionSuppose randomization is performed through SBR (see Example (ref)), then Assumptions (ref) and (ref) are satisfied.

Results

In this section we formally define our proposed procedure and present results about its asymptotic behavior. Section (ref) sets up the problem and presents the main results about the asymptotic normality of our estimator. Section (ref) considers several extensions: a cross-validation procedure to select the depth $L$ of the stratification tree, asymptotic results for a “pooled" estimator of the ATE, and extensions for the targeting of subgroup specific effects.

Main Results

In this section we describe our procedure and present our main formal results. Recall our discussion at the end of Section (ref): given first-wave data, our goal is to estimate a stratification tree which minimizes the asymptotic variance in a certain class of ATE estimators, which we now introduce. For a fixed $T \in \mathcal{T}_L$, let $\{(Y_i, A_i, X_i)\}_{i=1}^n$ be an experimental sample generated from a randomized experiment with randomization procedure $(A_i(T))_{i=1}^n$. Consider estimation of the following equation by OLS: $$Y_i = \sum_k \alpha(k ;T) {\bf 1}\{S_i = k\} + \sum_k \beta(k ;T) {\bf 1}\{A_i = 1, S_i = k\} + u_i~.$$ Then our ATE estimator is given by $$\hat{\theta}(T) := \sum_k\frac{n(k; T)}{n}\hat{\beta}(k; T)~,$$ where $n(k; T) := \sum_i {\bf 1}\{S_i = k\}$. In words, this estimator takes the difference in means between treatments within each stratum, and then averages these over the strata. Given appropriate regularity conditions, the results in bugni2017 establish asymptotic normality for a fixed $T = (S, \pi) \in \mathcal{T}_L$: $$\sqrt{n}(\hat{\theta}(T) - \theta) \xrightarrow{d} N(0,V(T))~,$$ where $$V(T) := \sum_{k=1}^K P(S(X) = k) \left[\left(E[Y(1) - Y(0)|S(X)=k] - E[Y(1)-Y(0)]\right)^2 + \left(\frac{\sigma_0^2(k)}{1-\pi(k)} + \frac{\sigma_1^2(k)}{\pi(k)}\right)\right]~,$$ and $$\sigma_a^2(k) := E[Y(a)^2|S(X)=k] - E[Y(a)|S(X)=k]^2~.$$

Again we remark that this variance takes the form of the semi-parametric efficiency bound of hahn1998 amongst all estimators that use the strata indicators as covariates. We propose a two-stage adaptive randomization procedure which asymptotically achieves the minimal variance $V(T)$ across all $T \in \mathcal{T}_L$. For the rest of the paper, we denote the sample size of the first wave by $m$ and the sample size of the second wave by $n$. We index first-wave observations by $j = 1, \ldots, m$ and second-wave observations by $i = 1, \ldots, n$. In the first stage, we use first-wave data $\{(Y_j,A_j,X_j)\}_{j=1}^m$ to estimate some “optimal" tree $\hat{T}_m$ which is designed to minimize $V(T)$. In the second stage, we perform a randomized experiment using stratified randomization with $(A_i(\hat{T}_m))_{i=1}^n$ to obtain second-wave data $\{(Y_i,A_i,X_i)\}_{i=1}^n$. Finally, to analyze the results of the experiment, we consider both the “unpooled" estimator $\hat{\theta}(\hat{T}_m)$ defined above, which uses only the second-wave data to estimate the ATE, as well as a “pooling" estimation strategy, which use both waves of data to construct an ATE estimator (see Section (ref)).

remarkThe depth $L$ of the set of stratification trees will remain fixed but arbitrary throughout Section (ref). The primary reason for this is technical: in order to allow for a wide variety of possible assignment procedures we leverage and extend the results in bugni2017, but these results are derived in an asymptotic framework where the number of strata is fixed. Considering extensions of their results to settings where the number of strata grows is beyond the scope of this paper. However, we provide three arguments for why considering a fixed-$L$ asymptotic framework may not be a major limitation in our setting: (1) accurately estimating stratification trees with many strata is both computationally difficult and may result in poor finite-sample performance unless the first-wave sample size is unrealistically large (we return to the question of how to choose $L$ with this consideration in mind in Section (ref)), (2) in Appendix (ref) we provide some preliminary simulation evidence which suggests that there are decreasing returns to increasing $L$, and (3) a simple compromise for practitioners who wish to stratify more finely is to perform ad-hoc stratification within each of the strata (with accompanying assignment proportions) produced by our method. By arguing as in Theorem 6.1 in bugni2021 it can be shown that this is guaranteed to weakly decrease the asymptotic variance of our estimator, and moreover, as long as the ad-hoc stratification is not too fine, Proposition (ref) establishes that our proposed inference procedure will still be valid. If a practitioner wishes to stratify very finely, then one possibility is to perform the optimal blocking procedure derived in bai2019 within each of the strata produced by our method. However in this case our proposed inference procedure may no longer be appropriate. See bai2019 for details.

We now present the main theoretical properties of our method. First, we establish conditions under which the estimator $\hat{\theta}(\hat{T}_m)$ constructed using the second wave of data is asymptotically normal, with minimal variance in the class of estimators defined above. Additionally, we provide a consistent estimator of the asymptotic variance of our estimator, and establish a form of “robustness" of our estimator to potential inconsistency of $\hat{T}_m$.

From now on, to be concise, we will call data from the first-wave the pilot data, and data from the second-wave the main data. As in the paragraph above, denote the pilot data as $\{W_j\}_{j=1}^m := \{(Y_j,X_j,A_j)\}_{j=1}^m$. Given this pilot sample, we require the following high-level consistency property for our estimator $\hat{T}_m$:

assumptionThe estimator $\hat{T}_m$ is a $\sigma\{(W_j)_{j=1}^m\}/\mathcal{B}(\mathcal{T}_L)$ measurable function of the pilot data\footnotemark[2] and satisfies $$|V(\hat{T}_m) - V^*| \xrightarrow{p} 0~,$$ where $$V^* := \inf_{T \in \mathcal{T}_L} V(T)~,$$ as $m \rightarrow \infty$.

\footnotetext[2]{$\mathcal{B}(\mathcal{T}_L)$ is the Borel-sigma algebra on $\mathcal{T}_L$ generated by an appropriate topology and $\sigma\{(W_i)_{i=1}^m\}$ is the sigma-algebra generated by the pilot data. See Appendix (ref) for details.}

Note that Assumption (ref) does not require that $V^*$ is uniquely minimized at some $T \in \mathcal{T}_L$. Moreover, Assumption (ref) imposes no explicit restrictions on how $\hat{T}_m$ is constructed, or even on the nature of the pilot data itself. In Appendix (ref), Lemma (ref) we show that Assumption (ref) is sufficient to guarantee that the sequence of trees $\hat{T}_m$ “approaches" the set of minimizers of $V(\cdot)$, but this does not guarantee that this sequence converges to any fixed tree within that set. Similar results have been derived in a maximum-likelihood context in redner1981.

In Proposition (ref) below, we establish sufficient conditions on the pilot data under which an appropriate $\hat{T}_m$ can be constructed by solving the following empirical minimization problem: $$\hat{T}_m^{EM} \in \arg\min_{T \in \mathcal{T}_L} \widetilde{V}_m(T)~,$$ where \[\widetilde{V}_m(T) := \sum_{k=1}^K\frac{m(k;T)}{m}\left[\left(\hat{E}[Y(1) - Y(0)|S(X) = k] - \hat{E}[Y(1) - Y(0)]\right)^2 + \left(\frac{\hat{\sigma}^2_{0,S}(k)}{1 - \pi(k)} + \frac{\hat{\sigma}^2_{1,S}(k)}{\pi(k)}\right)\right]\] with

align*[align* omitted — 296 chars of source]

In general, computing $\hat{T}_m^{EM}$ involves solving a nonlinear discrete optimization problem. Although this problem does not have a unique solution, Proposition (ref) establishes conditions such that any (measurable) sequence of minimizers which results from solving this minimization problem will satisfy Assumption (ref), due to the fact that the empirical objective $\widetilde{V}_m(\cdot)$ approximates $V(\cdot)$ uniformly well as $m \rightarrow \infty$.

remarkIn Appendix (ref) we describe an evolutionary algorithm which effectively performs a stochastic search for the global minimizer of the empirical minimization problem, and provide rough guidelines for implementation. We make a few comments here about the effectiveness of this algorithm in practice: first, the algorithm finds the global minimum in simple verified examples. Second, the algorithm always returns the same tree in repeated runs of the algorithm (up to negligible perturbations), with appropriate tuning parameters. Third, our current implementation of the algorithm (implemented in Julia 1.6) terminates fairly quickly for moderate depths and sample sizes: typically in less than one hour on a personal computer.

In Proposition (ref) we verify Assumption (ref) for $\hat{T}_m^{EM}$ when the pilot data comes from a stratified RCT with equal assignment proportions across strata:

propositionSuppose the pilot data come from a RCT performed using stratified randomization with stratification function $\zeta: \mathcal{X} \rightarrow \{1, 2, ..., Z\}$ (where $Z < \infty$ is a fixed number) such that $P(\zeta(X) = z) > 0$ for all $z = 1, \ldots, Z$. Let $m(z; \zeta) := \sum_{j = 1}^m {\bf 1}\{\zeta(X) = z\}$, $m_a(z; \zeta) := \sum_{j = 1}^m{\bf 1}\{A_j = a, \zeta(X) = z\}$. Suppose that the randomization procedure for the pilot data satisfies: \[\left[\{(Y_j(0), Y_j(1), X_j)\}_{j=1}^m \perp \!\!\! \perp (A_j)_{j =1}^m\right] \Bigg| (\zeta_j)_{j =1}^m ~,\] \[\frac{m_1(z; \zeta)}{m(z;\zeta)} \xrightarrow{p} \pi \hspace{1mm} \text{for some $\pi \in (0, 1)$, for all $z = 1, \ldots Z$}~.\] Under Assumptions (ref), (ref), and (ref), Assumption (ref) is satisfied for $\hat{T}_m^{EM}$.

We now state the first main result of the paper: an optimality result for the estimator $\hat{\theta}(\hat{T}_m)$. In Remark (ref) we comment on some of the technical challenges that arise in the proof of the result.

theoremGiven Assumptions (ref), (ref), (ref), (ref), (ref), and (ref), we have that $$\sqrt{n}(\hat{\theta}(\hat{T}_m) - \theta) \xrightarrow{d} N(0,V^*)~,$$ as $m, n \rightarrow \infty$.
remarkHere we comment on some of the technical challenges that arise in proving Theorem (ref). First, we develop a theory of convergence for stratification trees by defining a novel metric on $\mathcal{S}_L$ based on the Frechet-Nikodym metric, and establish basic properties about the resulting metric space. In particular, we use this construction to show that a set of minimizers of $V(T)$ exists given our assumptions, and that $\hat{T}_m$ converges to this set of minimizers in an appropriate sense. For these results we frequently exploit the fact that for a fixed index $k \in [K]$, the class of sets $\{S^{(-1)}(k): S \in \mathcal{S}_L\}$ consists of rectangles, and hence forms a VC class. Next, because Assumptions (ref) and (ref) impose so little on the dependence structure of the randomization procedure, it is not clear how to apply standard central limit theorems. When the stratification is fixed, bugni2017 establish asymptotic normality by essentially re-writing the sampling distribution of the estimator as a partial-sum process. In our setting the stratification is random, and so to prove our result we generalize their construction in a way that allows us to re-write the sampling distribution of the estimator as a sequential empirical process van1996. We then exploit the asymptotic equicontinuity of this process to establish asymptotic normality (see Lemma (ref)).

Next we construct a consistent estimator for the variance $V^*$. Let $$\widehat{V}_H(T) := \sum_{k=1}^K \frac{n(k;T)}{n}\left(\hat{\beta}(k;T) - \hat{\theta}(T)\right)^2~,$$ and let $$\widehat{V}_Y(T) := \hat{R}'(T)\hat{V}_{hc}(T)\hat{R}(T)~,$$ where $\hat{V}_{hc}(T)$ is the robust variance estimator for the parameters in the saturated regression, and $\hat{R}(T)$ is following vector with $K$ “leading" zeros: $$\left(\hat{R}(T)\right)' := \left[0, 0, 0, \ldots, 0, \frac{n(1; T)}{n}, \ldots, \frac{n(K; T)}{n}\right]~.$$

We obtain the following consistency result:

theoremGiven Assumptions (ref), (ref), (ref), (ref), (ref), and (ref), then $$\widehat{V}(\hat{T}_m) \xrightarrow{p} V^*~,$$ where $$\widehat{V}(T) := \widehat{V}_H(T) + \widehat{V}_Y(T)~,$$ as $m, n \rightarrow \infty$.

Although Theorem (ref) guarantees that $\hat{\theta}(\hat{T}_m)$ is asymptotically normal as $m, n \rightarrow \infty$, we may be concerned about the validity of this approximation when conducting inference in settings where the pilot sample size is not large. Accordingly, we finish this section by presenting a result about the asymptotic validity of hypothesis tests constructed from $\hat{\theta}(\hat{T}_m)$ and $\widehat{V}(\hat{T}_m)$ when $\hat{T}_m$ is not necessarily itself consistent in the sense of Assumption (ref). Consider the problem of testing

equation[equation omitted — 120 chars of source]

at level $\alpha \in (0,1)$, using a standard test given by \[\phi_n(T) := {\bf 1}\{\left|W(T)\right| > z_{1 - \frac{\alpha}{2}}\}~,\] where \[W(T) := \frac{\sqrt{n}(\hat{\theta}(T) - \theta_0)}{\sqrt{\hat{V}(T)}}~,\] and $z_{1-\frac{\alpha}{2}}$ is the $1 - \frac{\alpha}{2}$ quantile of a standard normal random variable.

propositionLet $\widetilde{T}_m$ be any sequence of trees constructed from the pilot data. Suppose that \[V(T) > 0 \hspace{2mm} \text{for all} \hspace{2mm} T \in \mathcal{T}_L~.\] Given Assumptions (ref), (ref), (ref), (ref), and (ref), under the null hypothesis given by ((ref)), \[\lim_{m,n \rightarrow \infty} E[\phi_n(\widetilde{T}_m)] = \alpha~.\]

Note that Proposition (ref) also accommodates the case where the pilot sample size is fixed at some number $m'$ by simply defining the sequence $\widetilde{T}_m$ to be equal to $\widetilde{T}_{m'}$ for $m \ge m'$. We conclude from Proposition (ref) that, regardless of whether or not $\hat{T}_m$ is consistent for an optimal tree, we can use $W(\hat{T}_m)$ and the critical values from a standard normal distribution to conduct asymptotically valid inference. Indeed, we will see in the simulations of Section (ref) that even in situations where $\hat{T}_m$ is a very poor estimate of an optimal tree, the coverage of a confidence interval or size of a test constructed using $\phi_n(\hat{T}_m)$ are close to the nominal level.

Extensions

In this section we present some extensions to the main results. First we present a version of $\hat{T}_m$ whose depth is selected by cross-validation. Second, we describe a method to combine estimates of the ATE from both waves of data, and establish properties of the resulting “pooled" estimator. Finally, we explain how to accommodate the targeting of subgroup-specific effects.

Cross-validation to select $L$

In this subsection we present a method to help select the depth $L$ in practice. To choose $L$, we revisit the first-stage estimation problem via the lens of the general model-selection paradigm described in, for example, arlot2010. The tradeoff which arises when choosing between various choices of $L$ in the first-stage estimation problem can be framed as a classical tradeoff between approximation error and estimation error, as we now describe. For each $L$, let \[V^*_L := \min_{T \in \mathcal{T}_L}V(T)~.\] On one hand, using a larger $L$ allows us to attain a (weakly) lower value for the asymptotic variance $V^*_L$. On the other hand, using a larger $L$ makes the set of trees $\mathcal{T}_L$ more complex, and thus makes estimation of the optimal tree more difficult in a finite sample, since we run the risk of “overfitting". More formally, let $\bar{L}$ be some upper bound on the depth of trees to be considered (in practice, this would correspond to some computational limit, or potentially to some exogenous logistical constraint), and let $[\bar{L}] = \{0, 1, 2, \ldots, \bar{L}\}$, where we understand $L = 0$ to mean no stratification. Let $\hat{T}_m^{(L)}$ be a stratification tree of depth $L$ estimated from the pilot data. We can then decompose the excess variance obtained by using $\hat{T}_m^{(L)}$ relative to $V^*_{\bar{L}}$ as \[V(\hat{T}_m^{(L)}) - V^*_{\bar{L}} = \left(V^*_{L} - V^*_{\bar{L}}\right) + \left(V(\hat{T}_m^{(L)})- V^*_{L}\right)~.\] The first term on the right-hand-side of this expression can be understood as the approximation error that results from optimizing in the class of trees $\mathcal{T}_L$ instead of the class of trees $\mathcal{T}_{\bar{L}}$. The second term on the right-hand-side of the expression can be understood as the estimation error that results from using the estimated tree $\hat{T}_m^{(L)}$ instead of the optimal tree for the class $\mathcal{T}_L$. Using this decomposition we see immediately that by Assumption (ref) we are guaranteed that setting $L = \bar{L}$ achieves the smallest possible excess variance asymptotically, however, our goal here is to attempt to balance these two tradeoffs in finite samples. In particular, the ideal “oracle depth" $L^*_m$ for a pilot sample of size $m$ is given by \[L^*_m := \arg\min_{L \in [\bar{L}]} \left(V(\hat{T}_m^{(L)}) - V^*_{\bar{L}}\right)~,\] which exactly balances the tradeoff between the estimation and approximation errors. Of course, this is infeasible, and so instead we propose selecting a depth $\hat{L}_m$ via a standard “$B$-fold" cross validation procedure which is designed to balance the estimation and approximation tradeoffs described above.

For simplicity we describe $2$-fold cross validation, but we comment on other choices of $B$ in Remark (ref) below. The cross-validation procedure proceeds as follows. First, split the pilot sample randomly into two halves and denote these by $\mathcal{D}_1$ and $\mathcal{D}_2$. For each $L$, let $\hat{T}^{(L,1)}_m$ and $\hat{T}^{(L,2)}_m$ be stratification trees of depth $L$ estimated on $\mathcal{D}_1$ and $\mathcal{D}_2$, respectively. Let $\widetilde{V}^{(1)}_m(\cdot)$ and $\widetilde{V}^{(2)}_m(\cdot)$ be the empirical variances computed on $\mathcal{D}_1$ and $\mathcal{D}_2$ (where, in the event that a cell in the tree partition is empty, we assign a value of infinity to the empirical variance). Then we define the following cross-validation criterion: \[\widetilde{V}^{CV}_L := \frac{1}{2}\left(\widetilde{V}^{(1)}_m\left(\hat{T}^{(L,2)}_m \right) + \widetilde{V}^{(2)}_m\left(\hat{T}^{(L,1)}_m\right)\right)~.\] In words, for each $L$, we estimate a stratification tree on each half of the sample, compute the empirical variance of these estimates by using the other half of the sample, and then average the results. Intuitively, as we move from small values of $L$ to large values of $L$, we would expect that this cross-validation criterion should generally decrease with $L$, and then eventually increase, in accordance with the tradeoff between estimation error and approximation error. Given $\widetilde{V}^{CV}_L$, we select our depth $\hat{L}_m$ as follows: $$\hat{L}_m := \arg\min_{L\in [\bar{L}]} \widetilde{V}^{CV}_L~,$$ where in the event of a tie we choose the smallest such $L$. Accordingly, our cross-validated stratification tree is defined as \[\hat{T}_m^{CV} := \hat{T}^{(\hat{L}_m)}_m~,\] i.e. $\hat{T}_m^{CV}$ is chosen to be the stratification tree whose depth minimizes the cross-validation criterion $\widetilde{V}^{CV}_L$.

We assess the finite-sample performance of $\hat{T}_m^{CV}$ via simulation in Section (ref), and note there that trees constructed using the cross-validation procedure can outperform trees constructed using $\bar{L}$ when the pilot is small, and perform similarly when the pilot is large (we prove a formal result about the large pilot behavior of the cross-validation procedure in Appendix (ref)). As a result we recommend that researchers fix some maximum allowable depth $\bar{L}$ (again, this will frequently correspond to a computational limit), and then use cross-validation to select an appropriate depth between $0$ and $\bar{L}$. In Section (ref), we use this cross-validation procedure to select the depth of the stratification trees we estimate for the experiment undertaken in karlan2017.

remarkOur description of cross-validation above defines $B$-fold cross-validation for $B = 2$. It is straightforward to extend this to general $B$, where the dataset is split into $B$ folds. In many statistical applications $5$ or $10$ folds has become the practical standard. However, we illustrate in Appendix (ref) that, when the first-wave sample size is small, using more than $2$ folds can potentially reduce performance. With that in mind, increasing the number of folds can be beneficial when the first-wave sample size is not too small; in particular, using more folds results in more stability, in the sense that $\hat{L}$ fluctuates less across different splits of the data.

A pooling estimator of the ATE

In this subsection we study an estimator which allows us to “pool" data from both datasets when estimating the ATE. Pooling may be particularly useful in formal two-stage randomized experiments where the first wave sample-size is large relative to the total sample-size (for example, in the application we consider in Section (ref)).

Let $\hat{\theta}_1$ be an estimator of the ATE constructed from the pilot data, and let $\hat{\theta}(\hat{T}_m)$ be the estimator defined in Section (ref). We impose the following high level assumption on the asymptotic behavior of $\hat{\theta}_1$:

assumption$\hat{\theta}_1$ is an asymptotically normal estimator for the ATE: \[\sqrt{m}(\hat{\theta}_1 - \theta) \xrightarrow{d} N(0, V_1)~,\] as $m \rightarrow \infty$.

Assumption (ref) holds for a variety of standard estimators under various assignment schemes: see for example the results in bugni2015, bugni2017, and bai2018. We also impose the following assumption on the relative rates of growth of the pilot and main sample.

assumptionLet $m$ be the pilot data sample size, $n$ the main data sample size, and $N = m + n$. We assume that \[\frac{m}{N} \rightarrow \lambda~,\] for some $\lambda \in [0, 1]$.

We propose the following sample-size weighted estimator: \[\hat{\theta}_{MW} := \widehat{\lambda}\hat{\theta}_1 + (1 - \widehat{\lambda})\hat{\theta}(\hat{T}_m)~,\] where $\hat{\lambda} := m/N$. Theorem (ref) derives the limiting distribution of this estimator:

theoremGiven Assumptions (ref), (ref), (ref), (ref), (ref), (ref), (ref), and (ref), we have that \[\sqrt{N}(\hat{\theta}_{MW} - \theta) \xrightarrow{d} N(0, V^*_\lambda)~,\] where $N := n+m$ and $V^*_\lambda := \lambda V_1 + (1 - \lambda)V^*$, as $m, n \rightarrow \infty$.

In words, we see that the pooled estimator $\hat{\theta}_{MW}$ has an asymptotic variance which is a weighted combination of the optimal variance and the variance from estimation in the pilot experiment, with weights which correspond to their relative sizes. In the asymptotic regime where $\lambda = 0$, $V^*_\lambda = V^*$, and hence pooling has no impact on the asymptotic behavior of the estimator. In contrast, in an asymptotic regime where $\lambda \ne 0$, $\hat{\theta}_{MW}$ and $\hat{\theta}(\hat{T}_m)$ are computed on sample sizes which differ asymptotically. To compare their variances, note that \[\text{Var}\left(\hat{\theta}_{MW}\right) \approx \frac{V^*_\lambda}{N}~,\] whereas \[\text{Var}\left(\hat{\theta}(\hat{T}_m)\right) \approx \frac{V^*}{n}~,\] so that, asymptotically, pooling will be beneficial when $(1 - \lambda)V^*_\lambda < V^*$. In practice, we expect that this will often be the case when the pilot data make up a large proportion of the total sample size (as in for example the application in Section (ref)).

remarkThe pooled estimator we present in this section is myopic, in the sense that $\hat{T}_m$ and $\hat{\theta}_{MW}$ are estimated as if the researcher did not anticipate that they would pool the data in the second stage. This has the benefit of being straightforward to analyze under very general assumptions on the pilot experiment. However, we could also consider “smart" versions of pooling, where the researcher estimates $\hat{T}_m$ taking into account that pooling will occur in the second stage. We present a preliminary discussion of such a strategy in Appendix (ref), but due to the increased technical complications of this approach we do not pursue a formal analysis of the procedure in this paper.

Stratification Trees for Subgroup Targeting

In this subsection we explain how the method can flexibly accommodate the problem of variance reduction for estimators of subgroup-specific ATEs, while still minimizing the variance of the unconditional ATE estimator in a restricted set of trees. It is common practice in RCTs for the strata to be specified such that they are the subgroups that a researcher is interested in studying glennerster2013. This serves two purposes: the first is that it enforces a pre-specification of the subgroups of interest, which guards against ex-post data mining. Second, it allows the researcher to improve the efficiency of the subgroup specific estimates.

Let $S' \in \mathcal{S}_{L'}$ be a tree of depth $L' < L$, whose terminal nodes represent the subgroups of interest. Suppose these nodes are labelled by $g = 1, 2, ..., G$, and that $P(S'(X) = g) > 0$ for each $g$. The subgroup-specific ATEs are defined as follows: $$\theta^{(g)} := E[Y(1) - Y(0)|S'(X) = g]~.$$ We introduce the following new notation: let $\mathcal{T}_L(S') \subset \mathcal{T}_L$ be the set of stratification trees of depth $L$ which can be constructed as extensions of $S'$. For a given $T \in \mathcal{T}_L(S')$, let $\mathcal{K}_g(T) \subset [K]$ be the set of terminal nodes of $T$ which pass through the node $g$ in $S'$ (see Figure (ref) for an example).

figure[figure omitted — 1,616 chars of source]

Given a tree $T \in \mathcal{T}_L(S')$, a natural estimator of $\theta^{(g)}$ is then given by $$\hat{\theta}^{(g)}(T) := \sum_{k \in \mathcal{K}_g}\frac{n(k; T)}{n'(g)}\hat{\beta}(k; T)~,$$ where $n'(g) = \sum_{i=1}^n {\bf 1}\{S'(X_i) = g\}$ and $\hat{\beta}(k)$ are the regression coefficients of the saturated regression over $T$. It is straightforward to see from the recursive structure of stratification trees that choosing $T$ as a solution to the following problem: $$\min_{T \in \mathcal{T}_L(S')} V(T)~,$$ will minimize the asymptotic variance of the subgroup specific estimators $\hat{\theta}^{(g)}$, while still minimizing the variance of the global ATE estimator $\hat{\theta}$ in the restricted set of trees $\mathcal{T}_L(S')$. Moreover, to compute a minimizer of $V(T)$ over $\mathcal{T}_L(S')$, it suffices to compute the optimal tree for each subgroup, and then append these to $S'$ to form the stratification tree.

In Section (ref) we illustrate the application of this idea to the setting in karlan2017. In their paper, they study the effect of information about a charity's effectiveness on subsequent donations to the charity, and in particular the treatment effect heterogeneity between large and small prior donors. For their application we specify $S'$ to be a tree of depth $1$, whose terminal nodes correspond to the subgroups of large and small prior donors. We then estimate the optimal tree for each of these subgroups and append them to $S'$ to form a stratification tree which simultaneously minimizes the variance of the subgroup-specific estimators, while still minimizing the variance of the global estimator in this restricted class.

Simulations

In this section we analyze the finite sample behaviour of our method via a simulation study, and in particular analyze the performance of the cross-validation procedure presented in Section (ref). We consider three DGPs in the spirit of the designs considered in athey2016. We emphasize that although these designs are artificial, they highlight several interesting qualitative patterns. For all three designs in this section, the outcomes are specified as follows: $$Y_i(a) = \kappa_a(X_i) + \nu_a(X_i)\cdot\epsilon_{a,i}~.$$ Where the $\epsilon_{a,i}$ are i.i.d $N(0, 0.1)$, and $\kappa_a(\cdot)$, $\nu_a(\cdot)$ are specified individually for each DGP below. In all cases, $X_i \in [0,1]^d$, with components independently and identically distributed as $Beta(2,5)$. The specifications are given by:

{\bf Model 1}: $d = 2$, $\kappa_0(x) = 0.2$, $\nu_0(x) = 5$, $$\kappa_1(x) = 10x_1^2{\bf 1}\{x_1 > 0.4\} - 5x_2^2{\bf 1}\{x_2 > 0.4\}~,$$ $$\nu_1(x) = 1 + 10x_1^2{\bf 1}\{x_1 > 0.6\} + 5x_2^2{\bf 1}\{x_2 > 0.6\}~.$$ This is a “low-dimensional" design with two covariates. The first covariate is given a higher weight than the second in the outcome equation for $Y(1)$.

{\bf Model 2}: $d = 10$, $\kappa_0(x) = 0.5$, $\nu_0(x) = 5$, $$\kappa_1(x) = \sum_{j = 1}^{10} (-1)^{j-1}10^{-j+2}x_j^2{\bf 1}\{x_j > 0.4\}~,$$ $$\nu_1(x) = 1 + \sum_{j = 1}^{10} 10^{-j+2}x_j^2{\bf 1}\{x_j > 0.6\}~.$$ This is a “moderate-dimensional" design with ten covariates. Here the first covariate has the largest weight in the outcome equation for $Y(1)$, and the weight of subsequent covariates decreases quickly.

{\bf Model 3}: $d = 10$, $\kappa_0(x) = 0.2$, $\nu_0(x) = 9$, $$\kappa_1(x) = \sum_{j = 1}^{3} (-1)^{j-1}10x_j^2\cdot{\bf 1}\{x_j > 0.4\} + \sum_{j = 4}^{10}(-1)^{j-1}5x_j^2\cdot{\bf 1}\{x_j > 0.4\}~,$$ $$\nu_1(x) = 1 + \sum_{j = 1}^{3}10x_j^2\cdot{\bf 1}\{x_j > 0.6\} + \sum_{j = 4}^{10}5x_j^2\cdot{\bf 1}\{x_j > 0.6\}~.$$ This is a “moderate-dimensional" design with ten covariates. Here the first three covariates have similar weight in the outcome equation for $Y(1)$, and the next seven covariates have a smaller but still significant weight.

In each case, $\kappa_0(\cdot)$ is calibrated so that the average treatment effect is close to $0.1$, and $\nu_0(\cdot)$ is calibrated so that $Y_i(1)$ and $Y_i(0)$ have similar unconditional variances (see Appendix (ref) for details). In each simulation we test six different methods of stratification. In all cases, when we stratify we consider a maximum of $8$ strata (which corresponds to a stratification tree of depth 3). In all cases we use SBR to perform assignment. We consider the following methods of stratification:

itemize[topsep = 1pt] • No Stratification: Here we assign the treatment to half the sample, with no stratification. • Ad-hoc: Here we stratify in an “ad-hoc" fashion and then assign treatment to half the sample in each stratum. To construct the strata we iteratively select a covariate and stratum at random, and stratify on the midpoints of the currently defined stratum. • Ad-hoc $+$ Neyman: Here we split the sample and perform a pilot experiment to estimate the Neyman allocation for each stratum defined in Ad-hoc. We then use the resulting stratification to assign treatment in the second wave. • Stratification Tree: Here we split the sample and perform a pilot experiment to estimate a stratification tree. We then use this tree to assign treatment in the second wave. • Cross-Validated Tree: Here we estimate a stratification tree as above, while selecting the depth via 2-fold cross validation. If the depth of the resulting tree is less than 3, then we perform ad-hoc stratification within each of the computed leaves as described in Remark (ref). • Infeasible Optimal Tree: Here we estimate an “optimal" tree by using a large auxiliary sample (see Appendix (ref) for details). We then split the sample and perform a pilot experiment in the first wave, while using the optimal tree to assign treatment in the second wave.

We perform the simulations with a sample size of $5,000$, and consider three different splits of the total sample for the pilot experiment and main experiment. The pilot experiment was performed using SBR with ad-hoc stratification. To estimate the stratification trees we minimize an empirical analog of the asymptotic variance as described in Section (ref). The estimator of the ATE we use throughout is the pooled estimator described in Section (ref).

We assess the performance of the randomization procedures through the following criteria: the empirical coverage of a $95\%$ confidence interval formed using a normal approximation, the percentage reduction in average length of the $95\%$ CI relative to no stratification, the power of a $t$-test for an ATE of 0, and the percentage reduction in root mean-squared error (RMSE) relative to no stratification. For each design we perform $6,000$ Monte Carlo iterations. Tables (ref), (ref), and (ref) below present our simulation results.

remarkIn Appendix (ref) we present additional simulation results. In particular, we explore alternative choices of $B$ in our $B$-fold cross validation procedure and alternative choices for the maximum number of strata.
table[table omitted — 1,745 chars of source]
table[table omitted — 1,737 chars of source]
table[table omitted — 1,736 chars of source]

For all three designs, we find that both the stratification tree and CV tree generally outperform no stratification and ad-hoc stratification, with particularly sizable gains for Model 2. However, when using a small pilot, the performance of the stratification tree without cross validation is often found to be worse than not stratifying at all. In such settings, we find that the CV tree effectively protects against overfitting. With larger sized pilots, we see that both trees perform comparably to the optimal tree in all three designs. Overall, we conclude that our proposed cross-validation procedure does a good job of protecting against overfitting, and we recommend that practitioners use cross validation to help select the depth of their trees in practice. With that in mind, we would still caution against using our method with small pilots even when using cross validation to select the depth.

An Application

In this section we study the behavior of our method in an application, using the experimental data from karlan2017. First we provide a brief review of the empirical setting: karlan2017 study how donors to the charity Freedom from Hunger respond to new information about the charity's effectiveness. The experiment, which proceeded in two separate waves corresponding to regularly scheduled fundraising campaigns, randomly mailed one of two different marketing solicitations to previous donors, with one solicitation emphasizing the scientific research on FFH's impact, and the other emphasizing an emotional appeal to a specific beneficiary of the charity. The outcome of interest was the amount donated in response to the mailer. karlan2017 found that, although the effect of the research insert was small and insignificant, there was substantial heterogeneity in response to the treatment: for those who had given a large amount of money in the past, the effect of the research insert was positive, whereas for those who had given a small amount, the effect was negative. They argue that this evidence is consistent with the behavioral mechanism proposed by kahneman2003, where small prior donors are driven by a “warm-glow" of giving (akin to Kahneman's System I decision making), in contrast to large prior donors, who are driven by altruism (akin to Kahneman's System II decision making). However, the resulting confidence intervals of their estimates are wide, and often contain zero karlan2017. The covariates available in the dataset for stratification are as follows:

itemize[topsep = 1pt] • Total amount donated prior to mailer • Amount of most recent donation prior to mailer (denoted {\tt pre gift} below) • Amount of largest donation prior to mailer • Number of years as a donor (denoted {\tt \# years} below) • Number of donations per year (denoted {\tt freq} below) • Average years of education in census tract • Median zipcode income • Prior giving year (either 2004/05 or 2006/07) (denoted {\tt p.year} below)

As a basis for comparison, Figure (ref) depicts the stratification used for the first wave in karlan2017.

figure[figure omitted — 991 chars of source]

We estimate two different stratification trees using data\footnotemark[3] from the first wave of the experiment (with a sample size of $10,869$), that illustrate stratifications which could have been used to assign treatment in the second wave. We compute the trees by minimizing an empirical analog of the variance, as described in Section (ref). The first tree is fully unconstrained, and hence targets efficient estimation of the unconditional ATE estimator, while the second tree is constrained in accordance with Section (ref) to efficiently target estimation of the subgroup-specific effects for large and small prior donors (see below for a precise definition). In both cases, the depth of the stratification tree was selected using $2$-fold cross validation as described in Section (ref), with a maximal depth of $\bar{L} = 5$ (which corresponds to a maximum of $32$ strata). When computing our trees, given that some of these covariates do not have upper bounds a-priori, we impose an upper bound on the allowable range for the strata to be considered (we set the upper bound as roughly the 97th percentile in the dataset, although in practice this could be set using historical data). \footnotetext[3]{Replication data is available by request from Innovations for Poverty Action. Observations with missing data on median income, average years of education, and those receiving the “story insert" were dropped.}

Figure (ref) depicts the unrestricted tree estimated via cross-validation. We see that the cross-validation procedure selects a tree of depth one, which may suggest that the covariates available to us for stratification are not especially relevant for decreasing the variance of the estimator. However, we do see a wide discrepancy in the assignment proportions for the selected strata. In words, the subgroup of respondents who have been donors for more than $16$ years have a larger variance in outcomes when receiving the research mailer than the control mailer. In contrast the subgroup of respondents who have been donors for less than $16$ years have roughly equal variances in outcomes under both treatments.

figure[figure omitted — 657 chars of source]

Next, we estimate the restricted stratification tree which targets the subgroup-specific treatment effects for large and small prior donors. We specify a large donor as someone who's most recent donation prior to the experiment was larger than $\$100$. We proceed by estimating each subtree using cross-validation. Figure (ref) depicts the estimated tree. We see that the cross-validation procedure selects a stratification tree of depth 1 in the left subtree and a tree of depth 0 (i.e. no stratification) in the right subtree, which further reinforces that the covariates we have available may be uninformative for decreasing variance.

figure[figure omitted — 856 chars of source]

These results are not necessarily surprising given the nature of the experiment: with very high probability, a recipient of either mailer is likely to make no donation at all, and hence we might expect limited heterogeneity in the potential outcomes with respect to our observable characteristics.

remarkIn Appendix (ref) we repeat the simulation exercise of Section (ref) with an application-based simulation design. There we find that our method obtains very modest gains in precision relative to the stratification used in karlan2017, which may not be surprising given the nature of the experiment.

Conclusion

In this paper we proposed an adaptive randomization procedure for two-stage randomized controlled trials, which uses the data from a first-wave experiment to assign treatment in a second wave of the RCT. Our method uses the first-wave data to estimate a stratification tree: a stratification of the covariate space into a tree partition along with treatment assignment probabilities for each of these strata.

Going forward, there are several extensions of the paper that we would like to consider. First, although we have argued throughout this paper that we find tree partitions to be a natural and convenient constraint, there are serious theoretical and empirical questions about whether or not weakening this restriction could lead to large decreases in asymptotic variance. A potentially easy compromise would be to consider oblique tree partitions, where each split of the tree is determined by a linear-index of the covariates. It would be interesting to know to what extent our results generalize to this setting. Second, many RCTs are performed as cluster RCTs, that is, where treatment is assigned at a higher level of aggregation such as a school or city. Extending the results of the paper to this setting could be a worthwhile next step. Similarly, we could extend the results to settings with non-compliance by leveraging recent results in bugni2021. Another avenue to consider would be to combine our randomization procedure with other aspects of the experimental design. For example, carneiro2016 set up a statistical decision problem to optimally select the sample size, as well as the number of covariates to collect from each participant in the experiment, given a fixed budget. It may be interesting to embed our randomization procedure into a similar decision problem. Finally, although our method employs stratified randomization, we assumed throughout that the experimental sample is an i.i.d sample. Further gains may be possible by considering a setting where we are able to conduct stratified sampling in the second wave as well as stratified randomization. To that end, song2014 develop estimators and semi-parametric efficiency bounds for stratified sampling which may be useful.

{ Acknowledgments}

I am grateful for advice and encouragement from Ivan Canay, Joel Horowitz, and Chuck Manski. I would also like to thank three anonymous referees, Eric Auerbach, Yuehao Bai, Lori Beaman, Stephane Bonhomme, Federico Bugni, Ivan Fernandez-Val, Hidehiko Ichimura, Sasha Indarte, Seema Jayachandran, Vishal Kamat, Dean Karlan, Cynthia Kinnan, Dennis Kristensen, Ryan Lee, Eric Mbakop, Matt Masten, Francesca Molinari, Denis Nekipelov, Sam Norris, Susan Ou, Azeem Shaikh, Mikkel Solvsten, Imran Rasul, Alex Torgovitsky, Chris Udry, Takuya Ura, Andreas Wachter, Joachim Winter, and seminar participants at many institutions for helpful comments and discussions. This research was supported in part through the computational resources and staff contributions provided for the Quest high performance computing facility at Northwestern University, and the Acropolis computing cluster at the University of Chicago.

small