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.
126,329 characters · 22 sections · 102 citation commands
Optimal Stratification of Survey Experiments
{Keywords: Matched Pairs, Blocking, Survey Sampling, Robust Standard Error, Treatment Effects.} \\
{JEL Codes: C10, C14, C90}
\onehalfspacing
Randomized controlled trials (RCTs) are now common in economics research, with thousands of active experiments in the AEA RCT registry spanning a wide range of fields. A key objective of experimental design is to reduce the variance of treatment effect estimation, helping applied researchers make the most efficient use of their limited resources. One way to do this is by covariate-adaptive treatment assignment, which balances observed covariates between the treatment and control group at design-time. This paper contributes to theory of covariate-adaptive treatment assignment, but also models a new dimension of experimental design: the selection of the experimental participants.
The selection of participants, also known as the sampling frame, is an important step in designing an experiment. For example, abaluck2021 run an experiment to estimate the effect of mask distribution on covid infection rates in Bangladesh. From a pool of $1000$ eligible villages, they first randomly sample $600$ to be included in the experiment, then assign the sampled villages to various interventions that promote mask usage. Similarly, breza2021 estimate the effect of Facebook ads discouraging holiday travel on covid infection rates. Since their budget for running ads is finite, they first sample a small set of counties in which to run ads and collect outcome data, then randomly assign these sampled counties to low or high intensity of treatment. We show how to increase the efficiency of treatment effect estimation by sampling experimental units that are representative of the broader target population, and provide new inference methods that take full advantage of these precision gains.
To do so, this paper introduces a new family of finely stratified randomization procedures that can be used for both representative sampling of the experimental units and finely balanced treatment assignment. When used for assignment, our method generalizes the principle of matched pairs randomization to propensities $\propfn(x) \not = 1/2$, allowing discrete or continuous stratification variables in general dimension. The basic building block is a new algorithm that matches the experimental units into homogeneous groups of $k$ by minimizing an objective function directly linked to estimation efficiency. This matching algorithm also enables finely stratified sampling of the experimental participants. For example, suppose $1000$ people respond to an advertisement to participate in an experiment, but the experimental budget only allows for $200$ participants. Using observed covariate information, we match the units into homogeneous $5$-tuples, sampling $\propselect = 1/5$ of the units in each tuple to participate, uniformly at random. By finely representing the distribution of treatment effect heterogeneity in the population into our smaller experiment, this sampling procedure reduces the variance due to treatment effect heterogeneity. More generally, we provide finely stratified designs implementing heterogeneous sampling propensities $\propselect(x)$.
We study a two-stage procedure that (1) samples participants then (2) assigns treatments to the sampled units, using finely stratified designs at both stages. Under finely stratified treatment assignment alone, the difference-of-means estimator $Y \sim 1 + D$ achieves the hahn1998 variance bound for the average treatment effect (ATE), effectively doing nonparametric regression adjustment “by design.” Our analysis shows that finely stratified sampling provides an additional nonparametric variance reduction, dampening the variance due to treatment effect heterogeneity. In particular, representative sampling makes this variance component scale with the the number of units we sample from, rather than the smaller true experiment size, boosting the effective sample size for this component of the variance. Extending recent results in bai2020pairs, we characterize the optimal stratification variables for sampling and assignment, showing that for sampling one wants to stratify on covariates that are most predictive of treatment effect heterogeneity. In an extension, we also study estimation of the sample average treatment effect (SATE) over the eligible population using design-based asymptotics.
Building on these asymptotic results, we formalize and solve an optimal design problem with fixed experimental budget and heterogeneous costs. In development economics, the cost of including a village in an experiment can vary widely based on observable characteristics such as its distance from the urban center, village size and so on. This forces applied researchers to choose a tradeoff between sample size and sample representativeness when they select where to experiment. We propose a new formalization of this tradeoff, deriving the jointly optimal sampling intensity $\propselectopt(x)$ and assignment propensity $\propfnopt(x)$ under finely stratified randomization. Our results provide simple heuristics for optimal sampling of the experimental units, analogous to classical results on sample allocation for coarsely stratified survey design (cochran1977). Under homoskedasticity, for instance, the optimal sampling propensity $\propselectopt(x) \propto \cost(x)^{-1/2}$, where $\cost(x)$ is the cost of including a unit of type $X=x$ in the experiment. We show that an oracle design that implements discretizations of $\propselectopt(x)$ and $\propfnopt(x)$ using fine stratification minimizes the asymptotic variance of our estimator over all stratified designs, subject to the experimental budget constraint.
We also briefly investigate exact optimality in finite samples. For fixed propensity $\propfn = 1/2$, we prove that the globally optimal covariate-adaptive randomization takes the form of a novel “alternating design”, assigning a certain optimal allocation vector $(d_i^*)_{i=1}^n$ and its mirror image $(1-d_i^*)_{i=1}^n$ each with probability $1/2$. The optimal allocation $(d_i^*)_{i=1}^n$ solves the well-known Max-Cut graph partitioning problem (rendl2008), with edge weights related to the smoothness of the outome functions.
We apply our optimal stratification results to the problem of two-wave design using data from a pilot experiment, providing the first fully efficient solution to this problem. The basic idea is to estimate the optimal sampling and assignment propensities $\propselectopt(x)$ and $\propfnopt(x)$ using the pilot, then implement these estimates in the main experiment using fine stratification. Under large pilot asymptotics, this strategy is as efficient as the oracle design, achieving the budget-constrained minimal asymptotic variance. We also provide results under fixed pilot asymptotics, and briefly discuss potential robustifications. In the case without sampling, the problem of design using a pilot has received considerable attention in the recent literature, see for instance hahn2012, tabord-meehan2020, and bai2020pairs, and we give a detailed comparison with these results.
Finally, we provide novel asymptotically exact inference for the average treatment effect under joint finely stratified sampling and assignment, using a collapsed-strata\footnote{See also abadie2008 and bai2021inference for related results in the context of matched pairs assignment.} type estimator (hansen1953). These inference methods allow experimenters to fully exploit the efficiency gains from both finely stratified sampling and assignment. The use of non-constant sampling proportions $\propselect(x)$ produces discontinuities in the propensity-weighted outcome functions, introducing new technical challenges relative to previous work. Simulations and an empirical application to $N=9$ papers recently published in top journals in economics demonstrate the value of our proposed methods.
Our sampling model is related to the classical literature on survey sampling, e.g.\ as surveyed in cochran1977 and lohr2021. In contemporaneous work, yang2021 propose a two-stage design using rerandomization for both sampling and assignment. Under rerandomization, difference-of-means estimation is asymptotically slightly less efficient than ex-post linear covariate adjustment. By contrast, we show that two-stage fine stratification is asymptotically equivalent to nonparametric covariate adjustment for the imbalances in both the sampling and assignment variables. Proposition (ref) provides a formal equivalence statement.
For an overview of experimental design theory, see rosenberger2016book or athey2017survey. A representative sample of recent work on stratified treatment assignment includes imai2009, bugni2018inference, fogarty2018, wang2021, bai2021inference, dechaisemartin21, bai2020pairs, and tabord-meehan2020. For treatment assignment, our work is most related to bai2020pairs, who introduces finely stratified designs for constant propensity $\propfn = a/k$ and univariate stratification variables. Aside from stratification, other recent proposals for balanced treatment assignment include kasy2016, kallus2017balance, li2018rerandomization, krieger2019, and harshaw2021. We explicitly compare with some of these methods in Remark (ref).
Our results on design using a pilot study is related to previous results in hahn2012, bai2020pairs, tabord-meehan2020, and kasy2021adaptive. We provide detailed comparisons in Section (ref) below. Our inference results are related to the method of collapsed-strata in hansen1953 and its modern variants studied in abadie2008 and bai2021inference.
The rest of the paper is organized as follows. Section (ref) introduces notation and discusses our matching algorithm. Section (ref) states our main asymptotic results, including the equivalence with nonparametric adjustment. Section (ref) formalizes and solves the optimal stratification and finite sample optimal design problems. Section (ref) discusses design using a pilot study. Section (ref) provides our inference methods. Our empirical results are presented in Section (ref), and recommendations for practice in Section (ref).
Consider running an experiment to estimate the average treatment effect (ATE). There are $n$ eligible units, with observed baseline covariates $(X_i)_{i=1}^n$. We wish to sample proportion $\propselect \in (0, 1]$ of these units to participate in the experiment. Denote $\Ti = 1$ if a unit is sampled and $\Ti=0$ otherwise. Sampled units are then assigned to treatment or control $\Di \in \{0,1\}$. Let $Y_i(d)$ denote the potential outcome of unit $i$ for $d \in \{0,1\}$. Since outcomes are only observed for participating units\footnote{If control outcomes are costlessly observed for all units, the sampling problem becomes trivial. We still contribute novel assignment designs in this case.} we may write \[ Y_i = \Ti[\Di Y_i(1) + (1-\Di) Y_i(0)]. \] We focus on estimation and inference for the population $\ate = E[Y(1) - Y(0)]$, modeling the $n$ eligible units $(X_i, Y_i(0), Y_i(1))_{i=1}^n \sim F$ as an iid sample from a superpopulation of interest. For example, the eligible units could be $n = 1000$ respondents to an online advertisement to participate in an experiment with a budget constraint of $100$ participants. In this context, the variable $\Ti \in \{0, 1\}$ models which of the $n=1000$ units we choose to include in the experiment. In some applications, the eligible units may comprise the entire population of interest. For example, the units $i=1, \dots, n$ may be the entire population of villages in a country, which we sample to obtain the participating villages. To accommodate such applications, Section (ref) in the appendix presents design-based versions of our main results, targeting the sample average treatment effect $\sate = n\inv \sum_{i=1}^n Y_i(1) - Y_i(0)$ in the eligible population. Here, we define the $\sate$ over the entire eligible population $i = 1, \dots, n$ that we are allowed to sample from, not just the smaller set of units that are chosen to participate in the experiment $\{i: \Ti = 1\}$.
Our goal is to sample a representative subset of the eligible units and assign them to treatment and control in a way that finely balances the baseline covariates $(X_i)_{i=1}^n$. To do so, we introduce a new family of finely stratified designs that generalize the principle of matched pairs randomization to arbitary propensity scores $\propfn(x) = P(D=1|X=x)$, with continuous or discrete covariates in general dimension. We also use these new designs for finely stratified sampling, allowing us to implement arbitrary heterogeneous sampling propensities $\propselect(x) = P(T=1|X=x)$, while finely balancing covariates between the sampled and non-sampled units. The basic building block of our method is a matched $k$-tuples design, which uses the baseline covariates to match units into homogeneous groups of $k$, randomly assigning $a$ out of $k$ units in each group to $T=1$ during sampling or $D=1$ during assigment. We formally define the method in the context of finely stratified sampling in the next definition.
{1pt}
{6pt}
{1pt}
{6pt}
Equation (ref) generalizes a similar condition in bai2021inference for the case of matched pairs $k=2$. We discuss matching algorithms and their associated homogeneity rates in Section (ref) below. Consider sampling and assignment propensities $\propselect = a/k$ and $\propfn = a'/k'$. In the rest of the paper, we study treatment effect estimation under a two-stage procedure:
This two-stage procedure is illustrated in Figures (ref) and (ref), using data from an election experiment in Turkey reported in baysan2022. Each color represents a different group of units formed during sampling and assignment. For example, in Figure (ref) we form groups of size $|\group| = 8$, randomly sampling $3$ out of $8$ units from each group to “represent” that group in the experiment. The sampled units are shown in blue in the figure on the right. In figure (ref), we match the sampled units into groups of $k' = 4$, assigning $3$ out of $4$ to $D=1$ in each group.
Section (ref) presents the most general version of our method, allowing different stratification variables $\psi_1$ and $\psi_2$ to be used for sampling and assignment, as well as varying sampling and assignment propensities $\propselect(x)$, $\propfn(x)$. Optimal choice of stratification variables for sampling and assignment is discussed in Section (ref). Our framework enables unified asymptotics and inference for a wide variety of different designs, as shown in the examples below.
One possibility is to treat the left hand side of Equation (ref) as an objective function and minimize it over all partitions of the units into groups of $k$. Denoting $\dimpsi = \dim(\psi)$, Theorem (ref) in the appendix shows that if $E[|\psi(X)|_2^\alpha] < \infty$ for some $\alpha > \dimpsi + 1$ then the optimal groups satisfy
If $\psi(X)$ is bounded this becomes $\Op(n^{-2/(\dimpsi + 1)})$, sharpening the rate achieved under a boundedness assumption in previous work on matched pairs (bai2021inference). For $k=2$, the optimal groups are computable in $O(n^3)$ time using an algorithm due to derigs1988.\footnote{We use the min-weight-matching implementation from the NetworkX $3.1$ module in Python.} Efficient algorithms for computing the optimal groups for general $k$ are not available, and we expect the problem to be NP-hard.\footnote{See karmakar2022 for hardness results in a related problem.}
Iterative Matching - Instead of calculating the optimal partition, for $k > 2$ we iteratively apply Derigs' algorithm to match units into larger groups. For fixed $k$, this procedure can be shown to satisfy the same rate in Equation (ref) above. There are many ways to implement iterative pairing for each $k$. For example, $k = 5$ can be obtained by pairwise matching of $4$-tuples to $1$-tuples, or $3$-tuples to $2$-tuples and so on. This is shown in the figure below, where each level of the tree represents a call to the optimal pairing algorithm.
Before the $j$th algorithm call, we add a certain number of “empty centroids”, represented by the $0$'s in the figure. We also prohibit certain types of matches in order to guarantee the desired sequence of group sizes. For example, at the second level of the tree on the left, we set the distance to $+ \infty$ between groups of size $|\group| = 2$ and $|\group'| = 1$, size $|\group| = 1$ and $|\group'| = 1$, and size $|\group| = 0$ and $|\group'| = 0$. There are many choices of such cardinality trees for each $k$, not all of which can be feasibly implemented using these type of constraints. We provide a canonical way of generating such sequences, as well as the required constraints at each algorithm call, that is guaranteed to implement the desired group cardinality $k$. Technical details are provided in Section (ref) in the appendix.
Large Experiments. This algorithm is highly tractable for small and medium experiment sizes. For example, matching $n=500$ units into $5$-tuples takes 23 seconds on a laptop computer, while $n=2000$ takes about 24 minutes. However, larger experiments quickly become intractable. For example, domurat2021 has $n=87394$, which would take about 3.8 years to match using the algorithm described above. To enable fine stratification in larger experiments, one possibility is to exploit the global shape of the baseline covariate data to rule out matches between distant units. To do so, let $v_1$ be the first principal component of the stratification variables $(\psii)_{i=1}^n$ and consider the following procedure:
Intuitively, we use the first principal component to sort units by their projection along a natural “direction” through the dataset. This exploits the idea that good matches are unlikely between non-adjacent folds: Figure (ref) in the appendix gives a visual representation. By parallelizing over $K=80$ folds, a dataset of size $n=87394$ can be matched about 5 minutes. The original version of our procedure and this “PCA folds” version are asymptotically equivalent for fixed $K$. We focus on the original version in the theory that follows.
This section contains our main asymptotic results, showing nonparametric efficiency gains from both finely stratified sampling and assignment. First, we state our main assumption.
Previous work on fine stratification has required Lipschitz continuity of the outcome function $E[Y(d)|\psi(X)=\psi]$ and variance $\var(Y(d) | \psi(X) = \psi)$, as well as boundedness of the stratification variables $\psi(X)$.\footnote{For example, see bai2021inference, bai2023.} We provide a novel technical analysis that allows all of these assumptions to be removed.
Estimation. Let $\est$ be the regression coefficient on $\Di$ in $Y_i \sim 1 + \Di$, estimated in the sampled units $\{i: \Ti=1\}$. This is just the usual difference-of-means estimator. Before continuing to our asymptotic results, we state a variance decomposition for $\est$ that will be used extensively in what follows. Let $\catefn(\psi) = E[Y(1) - Y(0) | \psi(X)=\psi]$ denote the conditional average treatment effect (CATE) and $\hk_d(\psi) = \var(Y(d) | \psi(X)=\psi)$ the heteroskedasticity function. Define the balance function
Suppose that sampling and assignment are both completely randomized, $\Tn \sim \crdist(\propselect)$ and $\Dn \sim \crdist(\propfn)$, as in Example (ref). Let $\nsampled = \sum_i \Ti$ denote the experiment size. Our work shows that $\sqrt{\nsampled} (\est - \ate) \convwprocess \normal(0, V)$ with variance
We think of the first term as the variance due to treatment effect heterogeneity, in particular, the heterogeneity predictable by the stratification variables. The second term is the variance due to random assignment, which arises from the chance covariate imbalances between treatment and control units created by complete randomization $\Dn \sim \crdist(\propfn)$. The results in the next section show how stratified sampling and assignment nonparametrically dampen the each component of this variance expansion.
In this section we state a central limit theorem for $\ate$ estimation for the simplest case where the sampling and assignment propensities $\propselect = a/k$ and $\propfn = a'/k'$ are constant. This result quantifies the efficiency gains from stratification in each stage of the design. Remarks (ref) and (ref) draw connections between our findings and classical results on semiparametric efficiency in the analysis of observational data. Section (ref) below provides the most general version of the results in this section.
Comparing to Equation (ref) above, the variance component $\var(\balancefn(\psi; \propfn))$ due to covariate imbalance between the treatment arms is now asymptotically negligible. The variance due to treatment effect heterogeneity $\var(\catefn(\psi))$ is now dampened by the sampling proportion $\propselect \in (0,1]$. To see why, observe that the normalization $\sqrt{\nsampled} (\est - \ate)$ effectively holds the experiment size $\nsampled$ constant as we vary $\propselect$. Holding $\nsampled$ constant, the number of eligible units $n \approx \nsampled / \propselect$ grows as $\propselect \to 0$. For small $\propselect$, there are many eligible units to sample from, allowing us to choose a highly representative sample of experimental participants. Our theory shows that this reduces the variance due to treatment effect heterogeneity that is predictable by the stratification variables. Another way to understand this is that under finely stratified sampling, the first variance component scales with the larger size $n$ of eligible units, rather than the true experiment size $\nsampled$. This effectively “boosts” the experiment size for this component of the variance.
Nonparametric Regression by Design. If $\propselect = 1$ or sampling is completely randomized $\Tn \sim \crdist(\propselect)$\footnote{The latter statement follows from the more general results in Section (ref).} then the asymptotic variance in Theorem (ref) is
This is exactly the hahn1998 semiparametric variance bound for $\ate$ with iid observations $(Y, D, \psi(X))$.\footnote{Recent work in armstrong2022 shows that this efficiency bound also holds in settings with covariate-adaptive randomization, including the designs considered here.} We achieve the semiparametric variance bound with a simple difference-of-means estimator, without the need for nonparametric re-estimation of the known propensity (hirano2003) or covariate adjustment with well-specified outcome models (robins95). We can interpret this result as saying that finely stratified treatment assignment $\Dn \sim \localdesigncond(\psi, \propfn)$ does nonparametric covariate adjustment “by design.” See Proposition (ref) below for a more formal equivalence result.
Representative Sampling. If sampling is locally randomized $\Tn \sim \localdesigncond(\psi, \propselect)$, then the variance due to treatment effect heterogeneity decreases from $\var(\catefn(\psi))$ to $\propselect \var(\catefn(\psi))$. In this case, $V$ can be strictly smaller than the classical semiparametric variance bound. Intuitively, by using the additional covariates $(\psi(X_i))_{i=1}^n$ to select a representative sample, we finely represent the distribution of treatment effect levels $(\catefn(\psii))_{i=1}^n$ in the larger eligible population into our experiment. More formally, consider an oracle setting where we observe the treatment effect level $\catefn(\psii)$ for each sampled unit $\Ti = 1$, estimating the $\ate$ by the sampled average $\wh \theta = (1/\nsampled) \sum_i \Ti \catefn(\psii)$. Our analysis shows that if $\Tn \sim \localdesigncond(\psi, \propselect)$ then \[ (1/\nsampled) \sum_{i=1}^n \Ti \catefn(\psii) = \en[\catefn(\psii)] + \op(n^{-1/2}). \] Because of this, the sampled average $(1/\nsampled) \sum_i \Ti \catefn(\psii)$ behaves like the infeasible average $\en[\catefn(\psii)]$ over all eligible units, including those not sampled into the experiment. This nonparametrically dampens the variance due to treatment effect heterogeneity from $\var(\catefn(\psi))$ to $\propselect \var(\catefn(\psi))$ for $\propselect \in (0, 1]$.
Extending bugni2019, the next example proves an equivalence between coarse stratification and linear regression adjustment in the case where sampling and asignment are both coarsely stratified.
This result shows that under coarsely stratified sampling and assignment, the simple difference-of-means estimator $\est$ behaves like the doubly-adjusted estimator $\estadj$ under complete randomization. Proposition (ref) in the next section generalizes this example to the case of fine stratification and nonparametric double adjustment for both sampling and assignment imbalances. The remainder of this subsection provides more intuition for Theorem (ref), connecting our results on fine stratification to the previous literature on semiparametric efficiency.
{1pt}
{6pt}
This section describes the most general version of our method, providing finely stratified designs with heterogeneous sampling and assignment proportions $\propselect(x)$, $\propfn(x)$. The asymptotics developed in this section allows us to formulate and solve the problem of optimal stratification with heterogeneous costs in Section (ref) below.
First, we formally define the procedure. Suppose that $\propselect(x) \in \{a_l / k_l: l \in L\}$ for some finite index set $L$. Similarly, suppose $\propfn(x) \in \{a_l' / k_l': l \in L'\}$ with $|L'| < \infty$. Extending our definition, let $\Tn \sim \localdesigncond(\psi, \propselect(x))$ denote the following double stratification procedure:
Equivalently, we partition each propensity stratum $S_l$ into groups $\group \subseteq S_l$ of size $k_l$ such that $S_l = \bigsqcup_{\group \in \mc G_l} \group$ and the homogeneity condition \[ n \inv \sum_{\group} \sum_{i,j \in \group} |\psii - \psij|_2^2 = n \inv \sum_l \sum_{\group \in \mc G_l} \sum_{i,j \in \group} |\psii - \psij|_2^2 = \op(1) \] In practice, we simply run our matching algorithm separately in each propensity stratum $S_l$, and draw $(\Ti)_{i \in \group} \sim \crdist(a_l/k_l)$ independently for each $\group \in \mc G_l$. Treatment assignment $\Dn \sim \localdesigncond(\psi, \propfn(x))$ is defined identically, partitioning only the units $\{i: \Ti = 1\}$ sampled into the experiment.
Before stating our main result, we extend the definition of our estimator to accommodate varying propensities. Define the double IPW estimator
If $\propfn(x)=\propfn$ and $\propselect(x)=\propselect$ are constant, then $\est_2 = \est + \Op(n\inv)$, where $\est$ is the difference-of-means estimator studied in the previous section.\footnote{This is because $\en[\Di] = \propfn + O(n\inv)$ for stratified designs. It would be false for $\Di \simiid \bern(\propfn)$.} Then abusing notation we denote both estimators by $\est$. The following theorem gives our asymptotic results for fine stratification with varying propensities, extending the fixed propensity results in Theorem (ref) above. We begin with the special case $\psi_1 = \psi_2 = \psi$ and $\propselect = \propselect(\psi)$, $\propfn = \propfn(\psi)$, all non-random.
If $\propselect = 1$ then this is exactly the hahn1998 semiparametric variance bound for $\ate$ with iid observations $(Y, D, \psi(X))$ and propensity $\propfn(\psi)$. This shows that under fine stratification the population IPW estimator is already semiparametrically efficient, with no need to nonparametrically re-estimate the known propensities $\propselect(\psi)$ and $\propfn(\psi)$ as in hirano2003. Next consider the efficiency gain from stratified sampling. The overall sampling proportion $\nsampled / n$ has $\nsampled / n = E[\propselect(\psi)] + \op(1)$. Then defining $\propselectavg = E[\propselect(\psi)]$, local randomization reduces the variance due to treatment effect heterogeneity from $\var(\catefn(\psi))$ to $\propselectavg \var(\catefn(\psi))$ for $\propselectavg \in (0, 1]$, just as in the constant propensity case.
Regression Equivalence. Theorem (ref) shows that if $\propselect = 1$ then difference-of-means is semiparametrically efficient. To the best of our knowledge, no such efficiency bound is available for joint finely stratified sampling and assignment $\Tn \sim \localdesigncond(\psi, \propselect(\psi))$ and $\Dn \sim \localdesigncond(\psi, \propfn(\psi))$ with $\propselect \not = 1$. Instead, we show a direct equivalence between unadjusted estimation under fine stratification and nonparametric regression adjustment under an iid design. In particular, the asymptotic variance $\varlocal$ above is the same as that achieved by a doubly-robust estimator that adjusts for covariate imbalances during both sampling and assignment. To state the result, consider regression estimators $\ceffnest_d(\psi)$ for $\ceffn_d(\psi) = E[Y(d) | \psi]$. Define the doubly-augmented IPW (2-AIPW) estimator
$\estadj$ adjusts for covariate imbalances due to both sampling and assignment. We implement the estimator using cross-fitting, similar to chernozhukov2017dml. See the proof for details. If $\propselect = 1$ this reduces to the familiar AIPW estimator for the $\ate$. The next result provides an equivalence between doubly-robust nonparametric adjustment under an iid design and unadjusted estimation under a finely stratified design.
Intuitively, Proposition (ref) shows that finely stratified sampling and assignment makes the unadjusted estimator $\est$ as efficient as a more complicated estimator that adjusts nonparametrically for covariate imbalances during both sampling and assignment.
Estimated Design Variables. We wish to formally accommodate the case where design variables $\psi, \propselect(\psi), \propfn(\psi)$ are estimated using previously collected data. To do so, define the random element $\proprandfix \indep (X_i, Y_i(1), Y_i(0))_{i=1}^n$ and let $\psi = \psi(X, \proprandfix)$, $\propfn = \propfn(\psi, \proprandfix)$, and $\propselect = \propselect(\psi, \proprandfix)$. For example, we could let $\proprandfix = (\ceffnest_d)_{d=0,1}$ be regression estimates of $\ceffn_d(x) = E[Y(d) | X=x]$ from a pilot experiment and set $\psi(X, \proprandfix) = (\ceffnest_0(X), \ceffnest_1(X))$. Define $\catefn(\psi, \proprandfix) = E[Y(1)-Y(0) | \psi, \proprandfix]$ and $\hk_d(\psi, \proprandfix) = \var(Y(d) | \psi, \proprandfix)$. The proof of Theorem (ref) shows that if $\Tn \sim \localdesigncond(\psi, \propselect(\psi, \proprandfix))$ and $\Dn \sim \localdesigncond(\psi, \propfn(\psi, \proprandfix))$ then $\sqrt{\nsampled} (\est - \ate) | \proprandfix \convwprocess \normal(0, \varlocal(\proprandfix))$ with conditional asymptotic variance
All expectations and variances are conditional on $\proprandfix$. The inference methods provided in Section (ref) are asymptotically exact conditional on $\proprandfix$. Note that marginally over both our experiment and the previous data $\proprandfix$, the estimator is asymptotically mixed normal $\sqrt{\nsampled} (\est - \ate) \convwprocess Z$, where $Z$ has characteristic function $E[\exp(-\frac{1}{2}t^2 V(\proprandfix))]$.\footnote{This mixed normal limit was also observed by cai2023 in a setting with iid treatments.}
Collecting Data After Sampling. In practice, experimenters may want to use different stratification variables $\psi_1$ for sampling and $\psi_2$ for assignment with $\psi_1 \not = \psi_2$. For example, $\psi_1$ may include publicly available administrative data, while $\psi_2$ includes additional survey covariates collected after sampling units into the experiment. To accommodate this, next we state our most general version of Theorem (ref). Suppose Assumption (ref) holds and let sampling and assignment $\Tn \sim \localdesigncond(\psi_1, \propselect(x))$ and $\Dn \sim \localdesigncond(\psi_2, \propfn(x))$. If $\propselect(x) = \propselect(\psi_1)$, require $\psi_1 \sub \psi_2$. Otherwise, require $(\psi_1, q) \sub \psi_2$. Then $\sqrt{\nsampled} (\est - \ate) \convwprocess \normal(0, \propselectavg \varlocal)$ with $\propselectavg = E[\propselect(X)]$ and asymptotic variance $\varlocal = V_1 + V_2$
The variance $V$ does not have a simple form like in the special cases considered above if $\psi_1 \not = \psi_2$. Instead, we expand $V$ relative to the efficient variance $V_1$, which we conjecture is the semiparametric efficiency bound in this setting. If $\psi_1 = \psi_2 = \psi$, and $\propselect = \propselect(\psi)$, $\propfn = \propfn(\psi)$ then $\varlocal$ can be rearranged into the form in Theorem (ref). The requirement that the stratification is increasing\footnote{In fact, we just require that $\psi_1 = h(\psi_2)$ for a measurable function $h$.} $\psi_1 \sub \psi_2$ is subtle. We defer this technical discussion to Remark (ref) in the appendix.
Optimal Stratification Variables. Inspecting the variance $V_2$ in Equation (ref) shows that the minimal dimension efficient stratification variables are $\psi_1^* = \catefn(X)$ and $\psi_2^* = (\catefn(X), \propselect(X), \balancefn(X))$. In this case, $V_2$ is identically zero, and $V = V_1$, the conjectured semiparametric efficiency bound. We include $\propselect(X)$ in $\psi_2^*$ to satisfy the subvector condition $(\psi_1, q) \sub \psi_2$. If $\propselect(x) = \propselect(\psi_1)$, we can just take $\psi_2^* = (\catefn(X), \balancefn(X))$ under the weaker condition $\psi_1 = h(\psi_2)$ for some measurable $h$. In practice, these optimal $\psi_1^*$ and $\psi_2^*$ are not known. This suggests letting $\psi_1(X)$ be a small subvector of the baseline covariates expected to be most predictive of treatment effect heterogeneity, and $\psi_2(X)$ a subvector expected to be predictive of outcomes.\footnote{Since $\balancefn(X; \propfn) \propto E[Y(1)|X] / \propfn + E[Y(0)|X] / (1-\propfn)$, we should prioritize predicting outcomes in the arm assigned with lowest probability.} If $\catefn(X) = F(\psi_1(X))$ for some function $F$, then the choice $\psi_1(X)$ is asymptotically optimal. If none of the baseline covariates predict treatment effect heterogeneity, so that $\catefn(X)$ is constant, then completely randomized sampling with $\psi_1^* = 1$ is efficient. We discuss pilot estimation of $\psi_1^*$ and $\psi_2^*$ in Section (ref).
In this section we formulate and solve the problem of optimal stratification in survey experiments, characterizing the optimal sampling and assignment propensities for budget-constrained experimentation with heterogeneous costs. We show how to implement these optimal propensities using fine stratification and prove that such a design is efficient. For simplicity, in what follows we restrict to the case with stratification variables $\psi_1 = \psi_2 = \psi$ and sampling and assignment propensities $\propselect = \propselect(\psi)$ and $\propfn = \propfn(\psi)$.
In Section (ref), we presented asymptotic results of the form $\sqrt{\nsampled} (\est - \ate) \convwprocess \normal(0, \varlocal)$ normalized by the experiment size $\nsampled = \sum_i \Ti$, which is the typical normalization in the previous literature. However, observe that the experiment size $\nsampled$ varies with the sampling propensity $\propselect(\psi)$ in our setting, making this normalization unsuitable for our current task of optimizing over $\propselect(\psi)$ to find the efficient sampling propensity. Because of this, in what follows we instead normalize by the number of eligible units $n$. Since $\nsampled / n \convp E[\propselect(\psi)]$, this just removes the multiplicative factor $E[\propselect(\psi)]$ from our previous results, so that $\rootn (\est - \ate) \convwprocess \normal(0, \varlocal(\propselect, \propfn))$
With this fixed normalization in hand, consider minimizing Equation (ref) over all sampling propensities $\propselect(\psi)$. Clearly the unconstrained solution is $\propselect^*(\psi) = 1$, making the experiment as large as possible. More generally, we can formalize the problem of experimentation with heterogeneous costs subject to a budget constraint.
Costs. Define $\cost(\psi; \propfn)$ to be the known, potentially heterogeneous cost of including a unit of type $\psi(X)=\psi$ in the experiment. One natural cost specification is
For example, in a development economics context $\cost_s(\psi)$ could be the “sampling cost” of paying volunteers to collect outcome data in a village $\psii = \psi$, while $\cost_1(\psi)$ and $\cost_0(\psi)$ are the marginal costs of assigning treatment and control, respectively.
To finish setting up the problem, define the ex-ante heteroskedasticity function
We interpret $\hkavg(\psi; \propfn)$ as the expected residual variance from sampling a unit with $\psii = \psi$ into the experiment, prior to realization of its random treatment assignment $\Di \in \{0,1\}$, and similarly for the expected costs in Equation (ref) above. For fixed assignment propensity $\propfn(\psi)$, the asymptotic variance objective can be written $V(\propselect) = \var(\catefn(\psi)) + E[\hkavg(\psi; \propfn)/\propselect(\psi)]$. Then the budget-constrained variance minimization problem with budget $\budget$ can be written
The next proposition characterizes the interior solutions to this problem. We assume that $\inf_{\psi} \hk_d(\psi) \geq c > 0$ and costs $\cost(\psi; \propfn) \in [C_l, C_u] \sub (0, \infty)$.
If the feasibility constraint $\sup_{\psi} \propselectopt(\psi) \leq 1$ is violated, the optimal sampling propensity may not have a simple analytical form. Remark (ref) below provides a rounding procedure that can be used to restore feasibility in this case. To build intuition for the form of the solution, consider the following special cases, suppressing dependence on $\propfn$.
Optimal Spending. Under optimal sampling, the total amount spent on units of type $\psi$ is $\cost(\psi)\propselectopt(\psi)dP(\psi) \propto \sqrt{\hkavg(\psi)\cost(\psi)} dP(\psi)$. This shows that we should spend more on units with larger ex-ante variance and larger per unit cost. However, since the optimal propensity undersamples high cost units, spending grows as $\sqrt{\cost(\psi)}$, instead of linearly as it would if $\propselect(\psi) = \propselect$ were constant.
The main result of this section studies implementation of the globally optimal stratified design, subject to the budget constraint. In particular, we show efficiency of a finely stratified implementation of the jointly optimal sampling and assignment propensities $\propselectopt$ and $\propfnopt$. In this section, we restrict to the case with costs $\cost(\psi; \propfn) = \cost(\psi)$ not depending on $\propfn$. See Remark (ref) for discussion of the general case.
First, we characterize the optimal assignment propensity. For any fixed sampling propensity $\propselect(\psi)$, the global minimizer of Equation (ref) is the conditional Neyman allocation $\propfn^*(\psi) = \hksd_1(\psi) / (\hksd_1(\psi) + \hksd_0(\psi))$. In some cases, we may only be interested in implementing a constant assignment propensity $\propfn^* \in (0,1)$. If $\propselect$ is also constant, then
Compare this to the classical Neyman allocation $\hksd_1 / (\hksd_1 + \hksd_0)$ with $\hksd_d = \text{SD}(Y(d))$. In our setting, only the residual variances $\hk_d(\psi) = \var(Y(d) | \psi)$ enter $\propfn^*$, since the fluctuations of $Y(d)$ predictable by $\psi(X)$ do not contribute to first-order asymptotic variance under fine stratification.
The jointly optimal sampling and assignment propensities are obtained by plugging the conditional Neyman allocation $\propfn^*(\psi) = \hksd_1(\psi) / (\hksd_1(\psi) + \hksd_0(\psi))$ into the formula for $\propselectopt(\psi; \propfn)$ above. This gives ex-ante variance $\hkavg(\psi) = (\hksd_1(\psi) + \hksd_0(\psi))^2$ and jointly optimal sampling and assignment propensities
The propensities $\propfn^*(\psi)$ and $\propselect^*(\psi)$ will generally need to be discretized in order to implement them using fine stratification. To do so, we provide novel asymptotics with both the number of distinct propensity levels $\levelsetn$ as well as the number of units in each group $|\group| = k_n$ growing with the sample size $n$.
If $\psi(X)$ is bounded, the final condition simplifies to $\kboundn |\levelsetn| = o(n)$. For example, one way to satisfy Definition (ref) is to round the optimal propensities to the nearest $a/k_n$ for some sequence $k_n \to \infty$, setting $\propselectopt_n(\psi) = \argmin \{|\propselectopt(\psi) - a/k_n| : 1 \leq a \leq k_n-1, k_n = \lfloor n^{1/2 - \epsilon} \rfloor\}$ and similarly for $\propfnopt_n(\psi)$. The main theorem of this section shows that such discretizations are asymptotically efficient.
The design in Theorem (ref) minimizes the asymptotic variance over all sampling and assignment propensities, subject to the budget constraint. If we set $\propselect=1$, then $V^* = \min_{0 \leq \propfn \leq 1} V_{H}(\propfn)$, minimizing the hahn1998 semiparametric efficiency bound for the $\ate$ over all propensities scores. As noted above, armstrong2022 shows that this bound also applies to the designs in this paper for the case $\propselect = 1$.
In this final subsection, we briefly discuss exact optimality in finite samples. Consider the case $\propselect = 1$ and $\propfn(X) = a/k$ fixed and constant. In this setting, bai2020pairs shows that if the balance function $\balancefn(x)$ were known, then matching units into strata of size $k$ according to their sorted $\balancefn(X_i)$ values minimizes $\mse(\est | \Xn)$ over all stratified designs. By contrast, here we show that if $\balancefn(x)$ were known, the class of stratified designs itself would generally be suboptimal. To see this, consider the case $\propfn = 1/2$. Define the complete graph $K_n$ with vertices $\{1, \dots, n\}$ and edge weights $w_{ij} = \balancefn(X_i)\balancefn(X_j)$. The Max-Cut optimization problem asks for a partition of the vertices into disjoint sets $E_1 \cup E_0 = \{1, \dots, n\}$ such that the weight of cut edges between $E_1$ and $E_0$ is maximized
For example, see rendl2008 for an overview. Let $E_0^*, E_1^*$ solve the Max-Cut problem in Equation (ref). Define the optimal treatment allocation $\dn^* = \dn^*(\Xn)$ by $d_i^* = \one(i \in E_1^*)$ for $1 \leq i \leq n$. Define the alternating design $P^*(\Dn | \Xn)$ that alternates between $\dn^*$ and its mirror image $1-\dn^*$ by \[ P^*(\Dn = d_{1:n}^* | \Xn) = P^*(\Dn = 1-d_{1:n}^* | \Xn) = 1/2 \] with $\Dn \indep \Wn | \Xn$ for the full data $\Wn = (X_i, Y_i(1), Y_i(0))_{i=1}^n$. Our next theorem shows that the the alternating design $P^*$ is globally optimal over the set of all covariate-adaptive designs with fixed treatment probability $P(\Di = 1) = 1/2$. We denote this set of designs by $\mc P_{1/2} = \{P: P(\Di = 1) = 1/2, \; \Dn \indep \Wn | \Xn \}$. For simplicity, suppose $n = 2m$ for an integer $m$.
The inequality is strict if Problem (ref) has a unique solution up to permutation of set labels. In particular, note that $P^*$ is not a matched pairs design. Nevertheless, $\balancefn(x)$ is not known, so neither the globally optimal design, nor the optimal stratified design from bai2020pairs are feasible. We also caution against plug-in approaches that use a pilot estimate of $\balancefn(X_i)$. Section (ref) below shows that such approaches are equivalent to regression adjustment with regressions estimated in the pilot instead of the main sample, which may perform poorly if the pilot is small. For these reasons, we do not further pursue finite sample optimal designs in this paper.
In this section, we study a procedure that uses pilot data to estimate and implement the solution to the optimal stratification problem derived in the previous section. We show that this feasible version of the optimal design is asymptotically efficient, achieving the budget-constrained optimal variance of Theorem (ref). In particular, for the case $\propselect = 1$ our procedure minimizes the semiparametric variance bound over all propensity scores, providing the first asymptotically efficient solution to the question of design using a pilot study (hahn2012). The methods in this section are also relevant when observational data from the target population or a closely related previous experiment are available. Small pilot considerations and potential robustifications are discussed in Remark (ref) below. Pilot estimation of the optimal stratification variables $\psi^*$ is discussed in Section (ref) below.
Fix stratification variables $\psi_1 = \psi_2 = \psi$ and consider estimating the optimal design for the budget-constrained problem in Equation (ref). This amounts to using pilot or proxy data to estimate the efficient sampling proportions $\propselect^*(\psi)$ and treatment propensity $\propfn^*(\psi)$. As a proof of concept, we first state our result under large pilot asymptotics, allowing consistent estimation of the heteroskedasticity functions $\hksd_d(\psi)$. We show that the feasible estimated optimal design is asymptotically efficient in the sense of Theorem (ref). Fixed pilot asymptotics and small sample considerations are discussed in Remark (ref) below.
In Section (ref) we derived the optimal propensities \[ \propselect^*(\psi) = \budget \frac{(\hksd_1(\psi) + \hksd_0(\psi))\cost(\psi)^{-1/2}}{E[(\hksd_1(\psi) + \hksd_0(\psi))\cost(\psi) \half]} \quad \quad \propfn^*(\psi) = \frac{\hksd_1(\psi)}{\hksd_1(\psi) + \hksd_0(\psi)}. \] The sampling propensity $\propselect^*(\psi)$ is optimal provided the feasibility condition $\sup_{\psi} \propselect^*(\psi) \leq 1$ is satisfied. Consider pilot heteroskedasticity estimates\footnote{In practice, we use a modification of fan1998 to estimate variance functions. See appendix section (ref) for details.} $\hkest_d(\psi)$ for $d = 0,1$. Define the propensity estimates\footnote{Note in $\wh \propselect(\psi)$ the average is taken over the main experiment covariates, allowing for covariate shift between pilot and main experiment.} \[ \wh \propselect(\psi) = \budget \frac{(\hksdest_1(\psi) + \hksdest_0(\psi))\cost(\psi)^{-1/2}}{\en[(\hksdest_1(\psii) + \hksdest_0(\psii))\cost(\psii) \half]} \quad \quad \wh \propfn(\psi) = \frac{\hksdest_1(\psi)}{\hksdest_1(\psi) + \hksdest_0(\psi)}. \] In practice, we may find $\wh \propselect(\psij) > 1$ for some $j$. This could be because the condition $\sup_{\psi} \propselect^*(\psi) \leq 1$ is violated and the optimal sampling problem does not have an interior solution, or just due to statistical error. However, we can transform $\wh \propselect(\psi)$ into an admissible sampling propensity by an iterative rounding procedure, described in Remark (ref) below. Suppose we have done so and let $\wh \propselect_n(\psi)$ and $\wh \propfn_n(\psi)$ be a sequence of discretizations of $\wh \propselect(\psi)$ and $\wh \propfn(\psi)$, satisfying the conditions in Definition (ref). We require the following technical conditions, including consistency of the pilot heteroskedasticity estimates.
Our main result shows that finely stratified implementation of the optimal propensity estimates is asymptotically fully efficient.
For intuition, it is also helpful to consider certain special cases of Theorem (ref). If we fix $\propselect = 1$, the design $\Dn \sim \localdesigncond(\psi, \wh \propfn_n(\psi))$ asymptotically minimizes the hahn1998 variance bound over all propensity scores
If $\wh \propfn^*$ is a consistent pilot estimate of the optimal constant propensity $\propfn^*$ in Equation (ref), then the design $\Dn \sim \localdesigncond(\psi, \wh \propfn_n)$ has $\rootn(\est - \ate) \convwprocess \normal(0, \varlocal)$ with \[ V^* = \var(\catefn(\psi)) + \min_{\propfn \in (0,1)} E\left [\frac{\hk_1(\psi)}{\propfn}+ \frac{\hk_0(\psi)}{1-\propfn} \right] \]
Section (ref) showed that the stratification variables $\psi_1^* = \catefn(X)$ and $\psi_2^* = (\catefn, \balancefn)(X)$ were asymptotically efficient for both sampling and assignment. This suggests setting $\psi_1(X) = \wh \catefn(X)$ and $\psi_2(X) = (\wh \catefn(X), \wh \balancefn(X))$, using pilot estimates of the various regression functions. In our notation, the design $\Dn \sim \localdesigncond(\wh \balancefn, \propfn)$ was proposed in bai2020pairs for the case with iid sampling.
Regression vs.\ Matching on Estimated Functions. Our first result is negative, suggesting that for small pilots such an approach would be dominated by not using the pilot data at all at the design stage, drawing treatments iid and doing regression adjustment in the main sample. For simplicity, set $\propselect = 1$ and let $\Dn \sim \localdesigncond(\wh \balancefn, \propfn)$ be the design with pilot-estimated stratification variables. Let $\est$ be the difference of means estimator formed using the data $W_{1:n} = (D_i, X_i, Y_i(D_i))_{i=1}^n$. Separately, define iid treatments $\check D_i \sim \bern(\propfn)$ and let $\estadj$ be the cross-fit AIPW estimator in Proposition (ref), estimated using the alternate data $\check W_{1:n} = (\check D_i, X_i, Y_i(\check D_i))_{i=1}^n$ with regression estimators $\check \ceffn_d(\psi)$. Define $\check \balancefn(X)$ by plugging in $\check \ceffn_d(\psi)$ to the formula in Equation (ref). Then with $c_p = (\propfn-\propfn^2)^{1/2}$ the estimators $\est$ and $\estadj$ have identical expansions
The residual terms $\rootn R_n, \rootn \check R_n \convwprocess \normal(0, v)$ for $v > 0$ and are mean-independent of the first term. Denote the imbalance terms $B_n = \en[(\Di - \propfn)(\balancefn - \wh \balancefn)(X_i)]$ and $\check B_n = \en[(\check \Di - \propfn)(\balancefn - \check \balancefn)(X_i)]$. These terms control estimator error due to covariate imbalances between the treatment arms. With pilot regression error $r_n^{pilot} = \max_d \|\ceffnest_d - \ceffn_d\|_{2, \psi}$ and main sample regression error $ r_n^{main} = \max_d \|\check \ceffn_d - \ceffn_d\|_{2, \psi}$, it's easy to show that \[ \rootn B_n = \Op(r_n^{pilot}) \quad \text{and} \quad \rootn \check B_n = \Op(r_n^{main}). \] We expect pilot estimation error to be larger $r_n^{pilot} \gg r_n^{main}$ if the pilot is much smaller than the main sample.
Robustified Approach. The discussion above showed that, for small pilots, the design $\Dn \sim \localdesigncond(\wh \balancefn, \propfn)$ behaves like a noisy version of the AIPW estimator $\estaipw$, with regression adjustments estimated using the pilot instead of the main experiment. However, the bai2020pairs approach could still dominate if e.g.\ $\wh \balancefn$ is estimated consistently from a large observational dataset or a larger previous experiment with closely related covariates and potential outcomes. The large pilot asymptotics in bai2020pairs can be extended to show that the two-stage sampling and assignment design $\Tn \sim \localdesigncond(\wh \catefn, \propselect)$ and $\Dn \sim \localdesigncond((\wh \catefn, \wh \balancefn), \propfn)$ achieves the optimal variance $V_1$ in Equation (ref). Another natural idea is to robustify the bai2020pairs approach, setting $\Tn \sim \localdesigncond((\wh \catefn, \wh \balancefn, \psi'), \propselect)$ and $\Dn \sim \localdesigncond((\wh \catefn, \wh \balancefn, \psi'), \propfn)$ for stratification variables $\psi'$ expected to be predictive of both treatment effects and outcomes ex-ante. This can then be combined with the methods in the previous section, setting $\psi = (\wh \catefn, \wh \balancefn, \psi')$ and proceeding exactly as in Section (ref). The efficiency of such designs under fixed pilot asymptotics is described by Equation (ref). Conditionally asymptotically exact inference, conditional on the pilot data, is available using the methods in Section (ref).
This section provides new methods for asymptotically exact inference on the ATE under two-stage locally randomized designs. To do so, we generalize pairs-of-pairs\footnote{Also known as the method of collapsed strata, as in hansen1953. See abadie2008 and bai2021inference for recent analyses.} type methods to accommodate designs with both finely stratified sampling and assignment, as well as varying propensities $\propselect(\psi), \propfn(\psi)$. Our inference methods enable applied researchers to report smaller confidence intervals that fully reflect the efficiency gains from all of our proposed designs.
For each assignment group $\group \in \groupset_n$, define the centroid $\bar \psi_{\group} = |\group|\inv \sum_{i \in \group} \psii$. Let $\groupmatching: \groupset_n \to \groupset_n$ be a bijective matching between groups satisfying $\groupmatching(\group) \not = \group$, $\groupmatching^2 = \identity$, and the homogeneity condition $\frac{1}{n} \sum_{\group \in \groupset_n} |\bar \psi_{\group} - \bar \psi_{\groupmatching(\group)}|_2^2 = \op(1)$. In practice, $\groupmatching$ is obtained by matching the group centroids $\bar \psi_{\group}$ into pairs using the algorithm in Section (ref). Let $\groupsetnu = \{\group \cup \groupmatching(\group): \group \in \groupset_n\}$ be the unions of paired groups formed by this matching. Define $a(\group) = \sum_{i \in \group} \Di$ and $k(\group) = |\group|$. Define the propensity weights $\weightinferenceone_i = (1-\propfni \propselecti)/(\propfni \propselecti)^2$ and $\weightinferencezero_i = (1-\propselecti(1-\propfni))/(\propselecti(1-\propfni))^2$. Finally, define the variance estimator components
Our inference strategy begins with the sample variance of the double-IPW estimator (Equation (ref)), which is consistent for the true asymptotic variance under an iid design, but too large under stratified designs. We correct this sample variance using the estimators above, which measure how well the stratification variables predict observed outcomes in local regions of the covariate space. Define the variance estimator
Our main result shows that $\varest$ is consistent for the limiting variance of Theorem (ref), enabling asymptotically exact inference.
By Theorem (ref) and the CLT in Section (ref), the confidence interval $\wh C = [\est \pm \varest^{1/2} c_{1-\alpha/2} / \rootn]$ with $c_{\alpha} = \Phi \inv(\alpha)$ is asymptotically exact in the sense that $P (\ate \in \wh C ) = 1-\alpha + o(1)$. Importantly, note that the scaling is by number of eligible units $n$, not the smaller experiment size $\nsampled = \sum_i \Ti \leq n$.
In this section, we quantify the performance of each of our designs on $N=9$ real DGP's from experimental papers covering a range of fields in applied economics. Our theoretical results showed separate variance reductions from each of the following: (a) finely stratified treatment assignment, (b) finely stratified sampling (c) the optimal propensities $\propselectopt(\psi)$ and $\propfnopt(\psi)$ (infeasible), and (d) consistent pilot estimates $\propselectest(\psi)$ and $\propfnest(\psi)$ of the optimal propensities (feasible). To quantify the marginal efficiency gain from each of our proposed methods in finite samples, we simulate unadjusted $\ate$ estimation under the following designs:
We evaluate each of these designs on data from experimental papers published in the AER between May 2021 and November 2022. We exclude papers for which data is unavailable or that do not fit into our framework for various reasons, e.g.\ having multiple interventions on the same unit with a time series structure. The included papers are abebe2021, baysan2022, casey2021, dellavigna2022, domurat2021, hussam2022, and lowe2021. We additionally include data from banerjee2021, a study with significant treatment effect heterogeneity as recently analyzed by chernozhukov2023experiments, as well as data from the Oregon health insurance experiment, reported in finkelstein2012, for a total of $N=9$.
For each paper, we impute missing potential outcomes for all units, defining the imputation $\tilde Y_i(d) = Y_i(d)\one(\Di=d) + \wh Y_i(d) \one(\Di \not = d)$. Following the empirical exercise in bai2020pairs, we use the matching-based imputation $\wh Y_i(d) = Y_{j(i)}(d)$ with $j(i) = \argmin_{j: \Dj=d} |\psii - \psij|_2$.\footnote{Model-based imputation of $Y_i(d) = \ceffn_d(\psii) + \hk_d(\psii) \residuali^d$ with $E[\residuali^d | \psii] = 0$, $\var(\residuali^d | \psii) = 1$, and $\residuali^d \sim \normal(0, 1)$ yields qualitatively similar results.} Let $N_{0}$ denote the size of the original experiment. Using this full panel of imputed potential outcomes, we do the following:
Since the ATE for the imputed DGP $(\tilde Y_i(0), \tilde Y_i(1), \psii)_{i=1}^N$ is known, we can compute the standard deviation (SD),\footnote{All of our designs and estimators have $E[\est - \ate] = 0$, so we do not report MSE.} coverage probabilities, and percent reduction in confidence interval length for each DGP and design. The goal of this exercise is to quantify the marginal variance reduction from each of our methods on the type of DGP's that occur in applied economics research, isolating the separate efficiency gains from finely stratified assignment, finely stratified sampling, as well as implementation of the optimal propensities from Section (ref).
Descriptions of each paper, including the treatment and outcome variables, our choice of stratification variables $\psi$, level of aggregation, and other parameters are provided in Section (ref) in the appendix. Experiment sizes $n$ are as in the original papers, ranging from $n=91$ for casey2021 to $n=1903$ for finkelstein2012. The one exception is domurat2021 ($N_0=87394$), for which we set $n=1000$ for Monte Carlo tractability. Baseline treatment proportions are set to $\propfn = 1/2$, except for lowe2021 with $\propfn=2/3$ and finkelstein2012 with $\propfn=1/3$. We use the large experiment version of our method with folds of size $200$ (Section (ref)) and implement sampling subordinate assignment (Remark (ref)).
Costs and Discretization - Our theory showed that efficiency can be improved by (1) sourcing a large pool of units willing to participate in the experiment and (2) choosing a representative experimental subsample from this pool. If the marginal cost of including a unit is zero, step (2) is trivial, and we just take as many units as possible. The marginal cost $\cost(\psi)$ of including different units is not reported in the papers in our sample. To understand the potential variance reduction from representative sampling with fixed $\propselect$ as well as the varying optimal propensity $\propselect^*(\psi)$ in the types of DGP's that occur in applied economics research, we specify non-zero costs $\cost(\psi) = \one(|\psi|_2 \leq \kappa) + 5 \one(|\psi|_2 > \kappa)$ with $\kappa = \text{Median}_{i=1}^n |\psii|_2$ and $\budget = 1.5$. This results in feasible constant sampling proportions $\propselect \approx 7/10$. For example, if $\psii$ were village location relative to an urban center, this would correspond to higher cost of collecting data in rural villages which, anecdotally, is a common feature of experiments in development economics. We discretize $\propselectopt_{hom}(\psi)$ and $\propselectopt(\psi)$ by choosing $\propselect_k(\cdot)$ to minimize discretization error $\en[(\propselect_k - \propselectopt)^2(\psii)]$ over the set $\{\propselect: \propselect(\psi) \in a/10: a=1, \dots, 10\}$ subject to a constraint on the number of distinct propensity levels $L_n = |\text{Image}(\propselect_k)|$, with $L_n \leq 2$ for $n < 500$, $L_n \leq 3$ for $500 \leq n < 1000$ and $L_n \leq 4$ for $1000 \leq n \leq 2000$.
Results. Our main results are presented in Table (ref), with papers listed by initials of the first author. The largest change in standard deviation (SD) is in the contrast between complete randomization CR and finely stratified assignment CR, Loc, with an average of $-27\%$ across the papers in our sample. This improvement is particulary striking in papers like baysan2022 and banerjee2021 with highly predictive baseline covariates. The average marginal change in SD attributable to finely stratified sampling (from CR, Loc to Loc) is smaller at $-6\%$. Recall that stratified sampling reduces the variance component $\var(\catefn(\psi)) \to \propselect \var(\catefn(\psi))$, with $\propselect = 7/10$ in our simulations. Larger reductions may be expected for smaller $\propselect$ (more eligible units). Using the optimal sampling proportions $\propselectopt_{hom}(\psi)$, which assume homoskedasticity, reduces the variance for some studies but increases it for others, resulting in $+4\%$ change on average. This is not surprising considering that many of these studies have considerable heteroskedasticity. The change in SD between Loc and \textbf{Pilot S} is $+4\%$ on average, while for the case with a large pilot the change in SD from \textbf{Loc} to \textbf{Pilot L} is $-5\%$ on average. This shows that with a large pilot, closely related previous experiment, or observational data from the same population, the feasible estimates of the optimal designs in Section (ref) can be used to increase efficiency. However, our empirical results suggest this may not be appropriate when only a small pilot study is available.
Next we discuss the performance of our inference methods (Section (ref)). Coverage is close to nominal, but somewhat conservative in finite samples. This is due to two different forces. First, match quality between groups $\group$ and $\groupmatching(\group)$ in the collapsed-strata variance estimators $\varestone$ and $\varestzero$ is worse than match quality within groups, which results in $\varestone$ and $\varestzero$ in Section (ref) being conservative. This effect is most severe for designs with highly predictive covariates. Second, match quality at the sampling stage $\Tn \sim \localdesigncond(\psi, \propselect(\psi))$ is generally better than match quality at the assignment stage, since more units are available during sampling. However, our variance estimators can only only use the “thinned out” outcome data available for the units $\Ti=1$ included in the experiment, which underestimates match quality during sampling. This effect will be most severe for small sampling proportions $\propselect \to 0$ and in DGP's with significant treatment effect heterogeneity.
The change in confidence interval length $\% \Delta$CI is slightly conservative but broadly reflects the efficiency gains in the first panel. This shows that our inference methods are able to take advantage of the reduction in variance from both finely stratified sampling and assignment, as well as designs with varying sampling proportions.
Our empirical results show robust variance reductions from fine stratification at both the sampling and assignment stages. When choosing stratification variables $\psi$, we recommend including baseline outcomes, if available, and a small set of other variables suspected to be predictive of outcomes and treatment effect heterogeneity. In particular, if experimenters have pre-registered measuring treatment effect heterogeneity with respect to a certain variable, then it is natural to include this variable in $\psi$. Fine stratification methods increase the value of collecting baseline survey data, insofar as extra investment in the baseline survey process allows us to measure variables expected to be most predictive of outcomes and treatment effect heterogeneity. Our theory in Section (ref) showed that the efficiency gains from stratified sampling are larger the more eligible units we have, since this helps sample more representative experimental units. Because of this, sourcing a large pool of candidate units for the experiment can improve precision, even if the experimental budget constraints do not allow all of these units to ultimately be included in the experiment.
The feasible sampling design $\propselect_{hom}^*(\psi)$ (assuming homoskedasticity) had mixed effects in the empirical application. Relative to the simpler Loc design with constant propensity $\propselect$, the $\propselect_{hom}^*(\psi)$ design reduced variance for some DGP's but increased it for others. Aside from potential misspecification, there is a finite sample tradeoff between (R) better optimization of the residual variance by using varying sampling propensity $\propselect(\psi)$ and (M) worse sampling and assignment matches due to having many different $\propselect(\psi)$ strata. For experiments with highly predictive covariates, effect (M) may dominate, so that a design with constant $\propselect(\psi) = \propselect$ may be preferable. However, if costs are very heterogeneous, then the residual variance effect (R) will dominate, and $\propselect_{hom}^*(\psi)$ can produce significant efficiency gains. The estimated optimal designs from Section (ref) performed well in our empirical application for $n_{pilot}=400$, but were generally too noisy for $n_{pilot} = 100$. In the absence of a very large pilot or related previous experiment, one alternative is to use observational data to calibrate the ex-ante variance function $\hkavg(\psi)$ appearing in the optimal design $\propselectopt(\psi)$. This could improve on the design $\propselect_{hom}^*(\psi)$, which unrealistically assumes perfect homoskedasticity, without the added noise associated with estimating the ex-ante variance $\hkavg(\psi)$ from a small pilot.
Finally, the inference methods in Section (ref) were slightly conservative in finite samples, but still allow researchers to report smaller confidence intervals that reflect the efficiency gains from finely stratified sampling and assignment.
\typeout
If the eligible units $\{1, \dots, n\}$ comprise the entire population of interest, then we may be interested in estimation and inference on the sample average treatment effect $\sate = \en[Y_i(1) - Y_i(0)]$, or the average conditional treatment effect (ACTE) $\en[\catefn(\psii)]$, studied in kolesar2021a. Note that both estimands are defined over the full population of eligible units, not just the smaller set of experiment participants $\{i: \Ti = 1\}$. For the $\sate$, suppose $\Tn \sim \localdesigncond(\psi, \propselect)$ and $\Dn \sim \localdesigncond(\psi, \propfn)$ and define the residual treatment effect variance $\hkte(\psi) = \var(Y(1) - Y(0) | \psi)$. Under the same conditions as Theorem (ref), we have $\sqrt{\nsampled}(\est - \sate) \convwprocess \normal(0, V_{\sate})$ with
The final variance component $E[\hkte(\psi)]$ is not identified, as in the case of $\sate$ estimation under complete randomization. Setting $\psi=1$ and $\propselect = 1$ recovers the classical results in that setting. Observe that $V_{\sate}$ decreases as the residual treatment effect heterogeneity $\hkte(\psi)$ increases. The negative sign reflects a competition between two opposing forces. To see this, let $\bar Y_i = Y_i(1) / \propfn + Y_i(0) / (1-\propfn)$ and consider the error decomposition \[ \est - \sate = \cov_n(\Ti, Y_i(1) - Y_i(0)) / \propselect + \cov_n(\Di, \bar Y_i \, | \Ti=1). \]
The first term parameterizes the errors due to correlation between the sampling variables and treatment effects, and naturally increases with $\hkte(\psi) = \var(Y(1) - Y(0) | \psi)$. The second term is increasing in $\cov(Y(1), Y(0) | \psi)$, or equivalently decreasing in $\hkte(\psi)$, and larger in mean square, resulting in a net negative dependence on $\hkte(\psi)$. As $\propselect \to 0$, the relative variance\footnote{The normalization $\sqrt{\nsampled}(\est - \sate)$ by experiment size $\nsampled = \sum_i \Ti$ holds the number of sampled units fixed.} due to sampling increases, exactly cancelling the $E[\hkte(\psi)]$ factor due to assignment in the limit. Conservative inference for the $\sate$ under stratified sampling and assignment can be based on either of the lower bounds $\hkte(\psi) \geq \hk_1(\psi) + \hk_0(\psi) - 2 \hksd_1(\psi) \hksd_0(\psi) \geq 0$.
The average conditional treatment effect $\acte = \en[\catefn(\psi)]$ is more difficult to motivate from a policy perspective. One potential application is a structural model of the form $Y_{it}(1) - Y_{it}(0) = \catefn(\psii) + \epsilon_{it}$, with systematic component $\catefn(\psii)$ mediated solely through observables $\psi$ and transitory shock component $\epsilon_{it}$.\footnote{Inference on a more general denoised SATE parameter in a model with transitory shocks is studied in deeb2022.} Intuitively, in this model the $\acte$ acts like a “denoised” version of the $\sate$. If the $\epsilon_{it}$ are uncorrelated, the $\acte$ estimated at time $t$ is the best predictor of the $\sate$ for a policy implemented at time $t+1$. The proof of Theorem (ref) shows that $\sqrt{\nsampled} (\est - \en[\catefn(\psii)]) \convwprocess \normal(0, V_{c})$ with identified variance $V_{c} = E[\hk_1(\psi) / \propfn + \hk_0(\psi) / (1-\propfn)] \geq V_{\sate}$.
Write $n = lk + \delta$ for some integer $l$ and $0 \leq \delta < k$. First, match a remainder group of size $\delta$ and set it aside. Then suppose without loss that $n = lk$. Let $J$ be the minimal positive integer such that $2^J \geq k$ and let $\sum_{j=0}^{J-1} a_j 2^j$ with $a_j \in \{0, 1\}$ be the binary representation of $2^J - k \geq 0$. Before the jth call of Derigs' algorithm, add $l$ singleton groups of fake units $\group = \{F_j\}$ of type j to the dataset if and only if $a_j = 1$. Let $R(g)$ denote the real units in a group $R(\group) \sub [n]$ and $F(g)$ the fake units so that $\group = R(g) \cup F(g)$. Before the jth call to Derig's algorithm, set $d(g, g') = +\infty$ if either of the following occur: (1) $F(\group) \cap F(\group') \not = \emptyset$ or (2) $|F(\group) \cup F(\group')| > 0$ but $|F(\group) \cup F(\group')| \not = \sum_{i=0}^j a_i$. Otherwise, set the distance $d(\group, \group') = |\bar \psi_{R(\group)} - \bar \psi_{R(\group')}|_2^2$. Compute the optimal pairing at each step. After $J$ steps, remove all the fake units by setting $\group = R(\group)$. The binary representation trick is inspired by karmakar2022, who studies a related matching problem. However, the algorithm in his paper does not seem to guarantee groups of the correct cardinality for larger $k$.
Figure (ref) shows PCA Folds for the “large experiment” version of our algorithm. We use the data from finkelstein2012, with $K = 4$ folds and $n=1903$ samples.
The theory in Section (ref) requires heteroskedasticity function estimates $\hkest_d(\psi)$. Our simulations and empirical application using a modification of fan1998. In a setting without outcomes $Y = \ceffn(X) + \hk(X) \epsilon$, they propose (1) use local linear regression to estimate $\ceffn(X)$ and (2) use local linear regression to project estimated residuals $(Y - \ceffnest(X))^2$ on $X$. In our setting, we let $\tilde Y_i(1) = Y_i \Di \Ti / \propfni \propselecti$ and (1) project $\tilde Y_i(1) \sim \psii$ to estimate $\ceffn_1(\psii)$. Next, we (2) project weighted residuals $\tilde \epsilon_i(1)^2 = (Y_i - \ceffnest(\psii))^2 \Di \Ti / \propfni \propselecti \sim \psii$ to estimate $\hkest_1(\psii)$, and similarly for $d=0$. We tested linear ridge regression, RBF-kernel ridge regression, and random forests for each regression step, with hyperparameters chosen by cross-validation in all cases. Kernel ridge estimated $\hk_d(\psi)$ the most precisely in dimensions $\dim(\psi)=1,2$, while forests were superior in higher dimensions. Our simulation and empirical results are presented using random forest regression. Pilot data is drawn from a stratified experiment of the specified size with $\propfn = 1/2$.