Finely Stratified Rerandomization Designs

Finely Stratified Rerandomization Designs

\onehalfspacing

Introduction

Stratified randomization is commonly used to increase statistical precision in experimental research.\footnote{For example, cytrynbaum2024adjustment reports a survey of 50 experimental papers in the AER and AEJ from 2018-2023, where 57% used some form of stratified randomization.} Recent theoretical work (e.g.\ bai2021inference) has shown that fine stratification, which randomizes treatment assignments within small groups of tightly matched units, makes unadjusted estimators like difference of means automatically semiparametrically efficient.\footnote{See cytrynbaum2024, armstrong2022, and bai2024efficiency for more detailed discussion.} In finite samples, however, the performance of such designs can deteriorate rapidly with the dimension of the covariates used for stratification, due to a curse of dimensionality in matching.\footnote{Under regularity conditions, the convergence rate of finite sample variance to asymptotic variance is $O(n^{-2/(d+1)})$ for dimension $d$ covariates, see cytrynbaum2024.} This motivates the search for alternative designs that insist upon nonparametric balance for a few important covariates, but only attempt to balance linear functions of the remaining variables. In this paper, we study finely stratified rerandomization designs, which first tightly match the units into groups using a small set of important covariates, then rerandomize within-group treatment assignments until a balance criterion on the remaining covariates is satisfied.

Our first contribution is to derive the asymptotic distribution of generalized method of moments (GMM) estimators under stratified rerandomization, allowing for estimation of generic causal parameters defined by moment equalities. We consider both superpopulation and finite population parameters, the latter of which may be more appropriate for experiments run in a convenience sample (abadie2014fixed), as is the case for the vast majority of experiments in economics (niehaus2017). We introduce novel finite population parameters, studying a finite population local average treatment effect heterogeneity parameter in an application to angrist2013. As in previous work on rerandomization (e.g.\ li2018rerandomization), the asymptotic distribution of GMM estimators is an independent sum of a normal and a truncated normal term. We show that, modulo this truncated term, unadjusted GMM under stratified rerandomization behaves like semiparametrically adjusted GMM (e.g.\ graham2011) under an iid design, with automatic nonparametric control over the stratification covariates and linear control over the rerandomization covariates. Intuitively, stratified rerandomization implements partially linear regression adjustment “by design.”

Our second contribution is to introduce novel forms of rerandomization based on nonlinear balance criteria. For example, we allow acceptance or rejection of an allocation based on the difference of estimated covariate densities between treatment and control units. We also study a design that rerandomizes until an estimated propensity score is approximately constant, forcing the covariates to have no predictive power for treatment assignments in our realized sample. We prove that the designs in a general family of nonlinear rerandomization methods are all asymptotically equivalent to standard rerandomization based on a difference of covariate means, with an implicit choice of covariates and balance criterion, which we characterize.

Our third contribution is to study optimization of the balance criterion itself. We propose a novel minimax scheme that allows the researcher to specify prior information about the relationship between covariates and outcomes, then rerandomizes until the worst case covariate imbalance consistent with this prior is small. We prove that this design minimizes the (asymptotic) computational cost of rerandomization, subject to a strict bound on estimation error over the set of models consistent with the prior. If our prior information set contains the truth, then this design bounds the asymptotic variance of stratified rerandomization within a small additive factor of the optimal semiparametrically adjusted variance. If the information set is instead a Wald region estimated from pilot data, we show that our minimax design bounds the asymptotic variance in the main experiment with high probability.

Our fourth contribution is to provide simple inference methods for generic causal parameters under stratified rerandomization designs. To do this, we first derive the optimal ex-post linear adjustment for GMM estimation, which depends on the stratification.\footnote{This extends recent work on optimal adjustment under pure stratified randomization for ATE estimation, e.g.\ see cytrynbaum2024adjustment, bai2024adjustment, or liu2020.} Optimal adjustment makes the asymptotic distribution of GMM insensitive to the rerandomization acceptance criterion, removing the truncated normal term from the limiting distribution and restoring asymptotic normality. We also show that combining rerandomization with ex-post linear adjustment provides a form of double robustness to covariate imbalances, which helps explain the strong performance of this method in our simulations. For finite population causal parameters, the asymptotic variance is generically not identified (neyman1990). We derive novel identified variance bounds for general finite population causal parameters, enabling asymptotically conservative inference that still exploits the efficiency gains from both stratified rerandomization and optimal adjustment. For superpopulation parameters, we present new inference methods that are asymptotically exact.

Finally, we provide simulations and an empirical application to estimating treatment effect heterogeneity among compliers in angrist2013, which both show the value of adding a rerandomization step to finely stratified designs. This effect can be seen clearly in Figure (ref), which compares stratified rerandomization to stratification plus ex-post adjustment, and Figure (ref), which compares stratified rerandomization to other designs like pure fine stratification. See the relevant sections below for more detailed discussion.

Related Literature

This paper builds on the literature on fine stratification in econometrics as well as the literature on rerandomization in statistics. Stratified randomization has a long history in statistics, see cochran1977 for a survey. Recent work on fine stratification in econometrics includes bai2021inference, bai2020pairs, cytrynbaum2024, armstrong2022, and bai2024efficiency. A sample of some recent work in the statistics literature on rerandomization includes morgan2012 and li2018rerandomization, wang2021, and wang2022. We build on both of these literatures, studying the consequence of rerandomizing treatments within data-adaptive fine strata. We show that finely stratified rerandomization does semiparametric (partially linear) regression adjustment “by design,” providing nonparametric control over a few important variables and linear control over the rest.

For our main asymptotic theory (Section (ref)), the most closely related previous work is wang2021 and bai2024efficiency. wang2021 study estimation of the sample average treatment effect (SATE) under stratified rerandomization, with quadratic imbalance metrics based on the Mahalanobis norm. We study rerandomization within data-adaptive fine strata, providing asymptotic theory for generic superpopulation and finite population causal parameters defined by moment equalities. We also allow for essentially arbitrary rerandomization acceptance criteria, not necessarily based on quadratic forms. bai2024efficiency study estimation of superpopulation parameters defined by moment equalities under pure stratified randomization, without rerandomization. We extend these results to stratified rerandomization as well as generic finite population parameters, providing “SATE-like” versions of the parameters in bai2024efficiency.\footnote{These parameters can be seen as causal versions of the conditional estimand defined in abadie2014fixed.} In concurrent work, li2024 study GMM estimation of univariate superpopulation parameters under stratified rerandomization with fixed, discrete strata. We study significantly more general forms of stratification and rerandomization criteria than considered in their work, allowing for both finite and superpopulation parameters of arbitrary dimension and fine stratification with continuous covariates.

For nonlinear rerandomization (Section (ref)), the closest related results are zhao2024 and li2021. zhao2024 rerandomize based on the p-value of a logistic regression coefficient, while we rerandomize until a general smooth propensity estimate is close to constant in $L_2$ norm. li2021 simulate a density based rerandomization design, but provide limited theoretical results. To the best of our knowledge, we present the first asymptotic theory for rerandomization based on the difference of nonlinear (e.g.\ density) estimates. For acceptance region optimization (Section (ref)), the closest related results are schindl2024, who study the optimal choice of norm for quadratic rerandomization, while liu2023 chooses an optimal quadratic rerandomization design using a Bayesian criterion, in both cases for rerandomization without stratification. We propose a novel minimax approach that accepts or rejects based on the value of a convex penalty function, tailored to prior information provided by the researcher or estimated from a pilot.

Our work on optimal adjustment (Section (ref)) extends recent work on adjustment for stratified designs, e.g.\ liu2020, cytrynbaum2024adjustment, bai2024adjustment, to stratified rerandomization and GMM parameters. Finally our work on inference under data-adaptive fine stratification (Section (ref)) builds on previous work by abadie2008, bai2021inference, and cytrynbaum2024. Other recent work that has considered variance bounds for finite population causal parameters includes aronow2014, fogarty2018b, ding2019, abadie2020, and xu2021.

Framework and Designs

Consider data $W_i = (R_i, S_i(1), S_i(0))$ with $(W_i)_{i=1}^n \simiid F$. The $S_i(d) \in \mr^{d_S}$ denote potential outcome vectors for a binary treatment $d \in \{0, 1\}$, while $R_i$ denote other pre-treatment variables, such as covariates. For treatment assignments $\Di \in \{0, 1\}$, the realized outcome $\Si = S_i(\Di) = \Di \Sione + (1-\Di) \Sizero$. In what follows, for any array $(a_i)_{i=1}^n$ we denote $\en[a_i] = n\inv \sum_{i=1}^n a_i$, with $\bar a_1 = \en[a_i | \Di=1] = \en[a_i \Di] / \en[\Di]$ and similarly $\bar a_0 = \en[a_i | \Di=0]$. Next, we define stratified rerandomization designs.

Intuitively, steps (1) and (2) describe a data-driven “matched k-tuples” design, while step (3) rerandomizes within k-tuples until the balance criterion is satisfied. Equation (ref) is a tight-matching condition, requiring that the groups are clustered locally in $\psi$ space. cytrynbaum2024 provides algorithms to match units into groups that satisfy this condition for any fixed $k$.

Next, we discuss a convenient rerandomization scheme that allows the researcher to select the approximate number of draws until acceptance.

Mahalanobis rerandomization uses a convenient choice of acceptance criterion, but the variance normalization and implicit acceptance region in Equation (ref) are not generally efficient for estimating causal parameters. We provide alternative designs that optimize the shape of acceptance region $A$ in Section (ref) below.

Causal Estimands. Next, we introduce a generic family of causal estimands defined by moment equalities. Let $g(D, R, S, \theta) \in \mr^{\dimmom}$ be a score function for generalized method of moments (GMM) estimation. Recall $W = (R, S(1), S(0))$ and for $D | W \sim \bern(\propfn)$ define $\phi(W, \theta) = E[g(D, R, S, \theta) | W] = \propfn \mom(1, R, S(1), \theta) + (1-\propfn) \mom(0, R, S(0), \theta)$. By construction, we have $E[\phi(W, \theta)] = 0$ $\iff$ $E[g(D, R, S, \theta)] = 0$. The function $\phi(W, \theta)$ provides a convenient parameterization to define our paired finite population and superpopulation causal estimands.

In what follows, we study GMM estimation of both $\thetatrue$ and $\thetan$ under stratified rerandomization designs, showing an asymptotic equivalence between stratified rerandomization and partially linear covariate adjustment. In particular, this framework allows us to introduce several useful finite population estimands $\thetan$ that do not appear to have been previously considered in the literature, such as Example (ref) below. The estimand $\thetan$ may be a more appropriate target for experiments run in a convenience sample, as is the case for the vast majority of experiments reported in the economics literature (niehaus2017). Inference on $\thetan$, provided in Section (ref), is generically more powerful than for $\thetatrue$, since we only have to account for estimation uncertainty due to random assignment, without extra variability from sampling into the experiment.

Note also that GMM estimation of the superpopulation parameter $\thetatrue$ under pure stratification was studied in bai2024efficiency.

Next, consider a setting where the researcher wants to estimate a parametric model of treatment effect heterogeneity in an experiment with noncompliance and randomized binary instrument $Z$. In the next example, we define finite population and superpopulation local average treatment effect (LATE) heterogeneity parameters, applying them in our empirical application to angrist2013 below.

GMM Estimation. For positive-definite weighting matrix $\weightmatest \in \mr^{\dimmom \times \dimmom}$ with $\weightmatest \convp \weightmatpop \succ 0$ and sample moment $\momest(\theta) \equiv \en[g(\Di, R_i, \Si, \theta)]$, the GMM estimator\footnote{In our examples, we will mainly be concerned with the exactly identified case. However, the theory for the over identified case is almost identical, so we include this as well.} is

We are mostly interested in the exactly identified case, where $\est$ solves $\momest(\est) = 0$. In what follows, we study the properties of generalized method of moments (GMM) estimation of the causal parameters $\thetatrue$ and $\thetan$ under stratified rerandomization.

Asymptotics for GMM Estimation

In this section, we characterize the asymptotic distribution of the GMM estimator $\est$ under the stratified rerandomization designs in Definition (ref). We show that the asymptotic variance of $\est$ is proportional to the residuals of a partially linear regression model, up to a remainder term due to slackness in the rerandomization criterion. In this sense, stratified rerandomization does partially linear regression adjustment “by design.” First, we state some technical conditions that are needed for the following results.

Next we state the technical conditions needed for GMM estimation. Define the matrix $G = E[(\partial / \partial \theta') \phi(W, \theta)]|_{\theta = \thetatrue} \in \mr^{\dimmom \times \dimtheta}$ and let $\mom_d(W, \theta) = \mom(d, R, S(d), \theta)$ for $d \in \zerone$. Recall the Frobenius norm $|B|_F^2 = \sum_{ij} B_{ij}^2$ for any matrix $B$.

Compactness could likely be relaxed using concavity assumptions or a VC class condition, but we do not pursue this here. In what follows it will be conceptually useful to reparameterize the score function.

Sampling and Assignment Expansion. Recall $\phi(W, \theta) = E[g(D, R, S, \theta) | W]$ for $W = (R, S(1), S(0))$. Define $\diff(W, \theta) \equiv \var(D)(\mom_1(W, \theta) - \mom_0(W, \theta))$, which we refer to as the “assignment function.” For Horvitz-Thompson weights $H = (D-\propfn)/(\propfn-\propfn^2)$, a calculation shows we can expand

Our results below show that $\diff(W, \theta)$ parameterizes estimator variance due to random assignment, while $\phi(W, \theta)$ parameterizes the variance due to random sampling for the superpopulation estimand $\thetatrue$. We work directly with this expansion in what follows.

Finite Population Estimand

Our first theorem studies GMM estimation of the finite population estimand $\thetan$, which solves $\en[\phi(W_i, \thetan)] = 0$. We extend these results to $\thetatrue$ in Corollary (ref) below. To state the theorem, define the GMM linearization matrix $\matgmm = -(G'\weightmatpop G)\inv G' \weightmatpop \in \mr^{\dimtheta \times \dimmom}$. Note that in the exactly identified case $\dimmom = \dimtheta$, we just have $\matgmm = -G\inv$. For brevity, we also denote the recurring constant $\vard = \var(D) = \propfn - \propfn^2$.

Before stating the main result, we first derive the influence function for GMM estimation of $\thetan$ under stratified rerandomization.

Lemma (ref) generalizes Example (ref) above, showing that \[ \est - \thetan = \matgmm \left (\en[ a(W_i, \thetatrue)| \Di=1] - \en[a(W_i, \thetatrue) | \Di=0] \right ) + \op(\negrootn). \] This implies that the errors in estimating any finite population GMM parameter $\thetan$ are driven by random imbalances in the assignment function $\diff(\Wi, \thetatrue)$ between treatment and control units, at least to first-order. Our main theorem shows that, by balancing $\psi$ and $h$ ex-ante, stratified rerandomization reduces these imbalances, improving precision.

Note the variance $\vtheta$ is a matrix $\vtheta \in \mr^{\dimtheta \times \dimtheta}$, so the minimum should be interpreted in the positive semidefinite sense. In particular, we say $V(\gammaoptmat) = \min_{\gammacoeff} V(\gammacoeff)$ if $V(\gammaoptmat) \preceq V(\gammacoeff)$ for all $\gammacoeff \in \mr^{\dimh \times \dimtheta}$. Theorem (ref) shows that $\rootn(\est - \estn)$ is asymptotically distributed as an independent sum of a normal $\normal(0, \vtheta)$ and truncated normal vector $\residualvar$. Both terms only depend on the “assignment” component of the influence function, $\matgmm \diff(W, \thetatrue)$. The variance is attenuated nonparametrically by the stratification variables $\psi$ and linearly by the rerandomization covariates $h$.

Residual Imbalance. The truncated Gaussian $\residualvar \sim \gammaoptmat'\zh \, | \, \zh \in A$ arises from residual covariate imbalances due to slackness in the acceptance criterion, since $A \not = \{0\}$. If $A$ is symmetric about zero, i.e.\ $x \in A$ iff $-x \in A$, then $E[\residualvar] = 0$, so the GMM estimator $\est$ is first-order unbiased, as usual. In principle, $\residualvar$ can be made negligible relative to $\normal(0, \vtheta)$ in large enough samples by choosing very small $A$. For example, if $A = B(0, \epsilon)$ then $R_{B(0, \epsilon)} \sim \{\gammaoptmat'\zh \, | \, |\zh|_2 \leq \epsilon\} \convp 0$ as $\epsilon \to 0$. However, in finite samples this may be computationally infeasible and could even invalidate our first-order asymptotic approximation.\footnote{See wang2022 for a detailed analysis of complete rerandomization, where $\epsilon_n$ can change with sample size.} We develop a minimax criterion to choose an efficient acceptance region $A$ in finite samples in Section (ref) below.

To isolate the precision gains due to rerandomization, the following corollary specializes Theorem (ref) to the case of stratification without rerandomization ($A = \mr^{\dimh}$), as well as complete randomization, defined in Examples (ref) and (ref).

Corollary (ref) shows that fine stratification reduces the variance of GMM estimation of $\thetan$ to $\vtheta = \vard \inv E[\var(\Pi \diff(W, \thetatrue) | \psi)] \leq \vard \inv \var(\matgmm \diff(W, \thetatrue))$, a nonparametric improvement. Rerandomization as in Definition (ref) provides a further linear variance reduction to $\vtheta = \min_{\gammacoeff \in \mr^{\dimh \times \dimtheta}} E[\var(\Pi \diff(W, \thetatrue) - \gammacoeff'h | \psi)]$, up to the residual imbalance term $\residualvar$.

Superpopulation Estimand

This section extends the asymptotics above to the superpopulation estimand $\thetatrue$ solving $E[\phi(W, \thetatrue)] = 0$. We show that by targeting $\thetatrue$ we incur additional sampling variance that is invariant to the distribution of treatment assignments $\Dn$.

The corollary shows that targeting $\thetatrue$ instead of $\thetan$ adds an extra independent $\normal(0, \vphi)$ term to the asymptotic distribution. The variance $\vphi$ arises due to iid random sampling of the sampling function $\matgmm \phi(W, \thetatrue)$. Notice that stratified rerandomization only reduces the variance due to imbalances in the assignment function $\Pi \diff(W, \thetatrue)$, while the variance due to sampling $\matgmm \phi(W, \thetatrue)$ is irreducible. In this sense, the statistical consequences of different designs and adjustment strategies all happen at the level of the finite population estimand $\thetan$, while targeting the superpopulation $\thetatrue$ just adds extra sampling noise. Note that for pure stratification, bai2024efficiency were the first to derive an analogue of part (b) of Corollary (ref), under different GMM regularity conditions than we use here.\footnote{In particular, bai2024efficiency allow for non-smooth GMM scores and impose a VC dimension condition on $\mom_d(W, \theta)$. We restrict to the smooth case, using compactness of $\Theta$ to avoid entropy conditions.}

Equivalence with Partially Linear Adjustment

Example (ref) showed that, up to the rerandomization imbalance $\residualvar$, the unadjusted estimator $\est = \bar Y_1 - \bar Y_0$ has asymptotic variance $\vtheta = \min_{\gamma \in \mr^{\dimh}} \vard \inv E[\var(\ylevel - \gamma'h | \psi)]$. This can be rewritten in terms of the residuals of a partially linear regression of $\ylevel$ on $\psi$ and $h$:

More generally, Theorem (ref) shows that under stratified rerandomization designs, the unadjusted GMM estimator $\est$ automatically behaves like semiparametrically adjusted GMM in the completely randomized setting. Formally, let $\ltwoproduct(\psi) = L_2^{\dimtheta}(\psi)$ be the $\dimtheta$-fold Cartesian product of $L_2(\psi)$, the space of square-integrable functions. Then the variance due to random assignment $\vtheta$ in Theorem (ref) is can be written in terms of the residuals of the influence function $\Pi \diff(W, \thetatrue)$ in a partially linear regression on $\psi$ and $h$:

Intuitively, stratified rerandomization does partially linear regression adjustment “by design,” providing nonparametric control over $\psi$ and linear control over $h$. For a more explicit equivalence statement, define $\oraclefn(\psi, h) = \gammaoptmat'h + t_0(\psi)$ to be the partially linear function achieving the optimum in Equation (ref). Define the oracle semiparametrically adjusted GMM estimator

For the $\sate$ estimation problem one can show that $\estsemiparam$ is just an oracle version of the usual augmented inverse propensity weighting (AIPW) estimator (robins95), with partially linear regression models in each arm.\footnote{Feasible partially linear adjustment in an iid mean estimation problem with missing data was studied in wang2004. See also the related semiparametric adjustment for GMM parameters in graham2011.}

Under a completely randomized design, we require ex-post semiparametric adjustment to achieve $\vtheta$. Under stratified rerandomization, however, the simple GMM estimator $\est$ automatically achieves $\vtheta$, up to the residual imbalance term $\residualvar$.

Nonlinear Rerandomization

In this section, we introduce several novel “nonlinear” rerandomization criteria, proving that in many cases such designs are first-order equivalent to linear rerandomization (Definition (ref)), with an implicit choice of covariates $h$ and acceptance region $A$. This shows that our asymptotics and inference methods apply to a broad class of asymptotically linear rerandomization schemes, expanding the scope of the results in Section (ref) above.

GMM Rerandomization

First, we generalize the imbalance metric $\imbalance$ in Definition (ref), allowing rejection of $\Dn$ based on potentially nonlinear features of the in-sample distribution of treatments and covariates $(\Di, X_i)_{i=1}^n$. Let $\rerandmom(X_i, \beta)$ be a GMM score function, separate from the score $\mom$ defining the estimands above. We can define a large class of interesting designs by stratifying and rerandomizing until $\rootn (\betaestone - \betaestzero) \approx 0$ for within-arm GMM estimators

Observe that if $\rerandmom(X_i, \beta) = X_i - \beta$, then $\wh \beta_d = \bar X_d$ for $d = 0,1$ and $\imbalancegmm = \imbalance$, so linear rerandomization is a special case. However, Definition (ref) also allows for novel designs, such as rerandomizing until the estimated densities of covariates $X_i | \Di=1$ among treated and $X_i | \Di=0$ among control are similar. To the best of our knowledge, we provide the first formal results for such a design.

Let $\betatrue$ be the unique solution to $E[\rerandmom(X, \betatrue)] = 0$ and define $\rerandjacob = E[(\partial / \partial \beta') \rerandmom(X_i, \betatrue)]$. Our next result shows that GMM rerandomization with acceptance criterion $\imbalancegmm \in A$ is equivalent to linear rerandomization (Definition (ref)) with an implicit choice of rerandomization covariates $\hi = m_i^* \equiv m(X_i, \betatrue)$ and linearly transformed acceptance region.

Theorem (ref) shows that by rerandomizing until $\rootn(\betaestone - \betaestzero) \in A$, we implicitly balance the influence function $-\rerandjacob \inv m(X_i, \betatrue)$ for the difference of GMM estimators above. In particular, this shows that all GMM rerandomization designs are first-order equivalent to linear rerandomization (Definition (ref)) for some choice of $\hi$ and acceptance region $A$.

For completeness, we provide a feasible linear rerandomization that exactly mimics the behavior in Theorem (ref). To do so, let $\wh h_i = \rerandmom(X_i, \wh \beta)$ for $\en[\rerandmom(X_i, \wh \beta)] = 0$ solving the pooled GMM problem, and rerandomize until $\rootn (\en[\wh h_i |\Di=1] - \en[\wh h_i | \Di=0]) \in \rerandjacobest A$ for $\rerandjacobest \convp \rerandjacob$.

One consequence of Corollary (ref) is that density based rerandomization for likelihood in an exponential family with sufficient statistic $r(X_i)$ is asymptotically equivalent to linear rerandomization setting $\hi = r(X_i)$.

Propensity Score Rerandomization

To motivate a propensity score based rerandomization procedure, note that despite $E[\Di | X_i] = p$ for all units, in finite samples the realized propensity $\wh p(B) = \en[\Di | X_i \in B]$ may significantly diverge from $\propfn$ in certain regions $B \sub \mr^{d_X}$ of the covariate space. This implies that covariates $X_i$ are predictive of treatment assignments $D_i$ ex-post, a form of “in-sample confounding,” which vanishes as $n \to \infty$ but affects precision. To prevent this, we could reject allocations where $|\wh p(B) - \propfn| > \epsilon$ for some collection of sets $B$. To make this idea tractable without fully discretizing, consider a parametric propensity model $p(X, \beta) = \link(X'\beta)$ for smooth link function $L$ (e.g.\ Logit) and define the MLE estimator

The average gap between the realized and ex-ante propensity score can be measured by

Intuitively, if $\imbalancesquare$ is large, then the covariates $X$ are predictive of treatment status in some parts of the covariate space. To avoid this, we propose rerandomizing until the imbalance metric $\imbalancesquare$ is below a threshold:

{1pt}

{6pt}

This design is illustrated in Figure (ref). Note that the covariate distribution is approximately balanced between $D=1$ and $D=0$ after acceptance. Our next result shows that propensity rerandomization as in Definition (ref) is equivalent to a simpler linear rerandomization design, with an implicit choice of ellipsoidal acceptance region. We require some extra regularity conditions on the link function $L$, which for brevity we state in Appendix (ref).

Theorem (ref) shows that for any sufficiently regular link function,\footnote{Theorem (ref) uses MLE estimation of $\wh \beta$, though we conjecture the result would be identical for inverse probability tilting (graham2012) or tailored loss function (zhao2019) estimation.} propensity rerandomization is asymptotically equivalent to Mahalanobis rerandomization in Example (ref), with acceptance criterion $n(\hbarone - \hbarzero)'\var_n(\hi)\inv (\hbarone - \hbarzero) \leq \epsilon \vard^{-2}$. Equivalently, propensity rerandomization behaves like linear rerandomization with $\imbalance = \rootn(\hbarone - \hbarzero)$ and ellipsoidal acceptance region $A = \var(h)\half B(0, \epsilon \vard^{-2})$.\footnote{A related result was found by zhao2024, who study rerandomizing until the p-value of a logistic regression coefficient is above a threshold.}

This section introduced a large family of novel rerandomization methods based on nonlinear estimators. Theorems (ref) and (ref) broaden the scope of our asymptotic theory and inference results, showing they also apply to these designs. These equivalence results raise the bar for future methodology improvements, showing that to obtain rerandomization designs with different first-order properties, we may need to consider more stringent imbalance measures, such as nonparametric two-sample test statistics. We leave this extension to future work.

Motivated by the “implicit” acceptance regions chosen by the designs in this section, next we formally study optimal choice of the acceptance region $A$.

Optimizing Acceptance Regions

In this section, we study efficient choice of the acceptance region $A \sub \mr^{\dimh}$. We propose a novel minimax rerandomization scheme and show that it minimizes the computational cost of rerandomization subject to a strict lower bound on statistical efficiency. This can be viewed as a form of dimension reduction, increasing rerandomization acceptance probability by downweighting less important directions in the covariate space $h$.

For simplicity, we first restrict to the case of estimating $\thetan = \sate$. Example (ref) showed that $\rootn(\est - \sate) | \Wn \convwprocess \normal(0, V(\gammaoptmat)) + \gammaoptmat'\zhs$, independent RV's with $\zhs = \zh | \zh \in A$ and variance $V(\gammaoptmat)$ that does not depend on $A$. The term $\zhs$ arises from residual imbalances in $h$ due to slackness in the acceptance region, $A \not = \{0\}$. The coefficient $\gammaoptmat$ comes from the partially linear regression\footnote{This expansion is without loss of generality. We do not impose well-specification $E[e | \psi, h] = 0$.}

All together, the residual imbalance term $\gammaoptmat'\zhs$ is the limiting distribution under rerandomization of $\gammaoptmat' \rootn (\hbarone - \hbarzero)$, the projection of covariate imbalances in $h$ along the direction $\gammaoptmat$. This suggests an oracle acceptance criterion that rerandomizes until the imbalance $|\gammaoptmat' \rootn (\hbarone - \hbarzero)| \leq \epsilon$, with acceptance region $A = \{x: |\gammatrue'x| \leq \epsilon\}$, reducing the problem to one dimension from arbitrary $\dim(h)$. Of course, this oracle design is infeasible since $\gammaoptmat$ is generally unknown when designing the experiment.

Minimax Rerandomization

Since $\gammaoptmat$ is unknown at design-time, we instead take a minimax approach that incorporates prior information about the coefficient $\gammaoptmat$. For belief set $\balancecoeffs \sub \mr^{\dimh}$ specified by the researcher, consider rerandomizing until the worst case imbalance consistent with $\balancecoeffs$ is small enough,

Equivalently, for imbalance $\imbalance = \rootn (\hbarone - \hbarzero)$ we rerandomize until $\fnbalance(\imbalance) \leq \epsilon$ for the convex penalty function $\fnbalance(x) = \sup_{\gamma \in \balancecoeffs} |\gamma' x|$. This significantly generalizes the quadratic imbalance penalty $p(x) = x'\var(h)\inv x$ implicitly used by Mahalanobis rerandomization (Example (ref)). Our next result shows that Equation (ref) is a linear rerandomization design, characterizing the induced acceptance region $A$.

{1pt}

{6pt}

Note that since $\acceptpolar$ is symmetric, the discussion after Theorem (ref) implies that the asymptotic distribution of $\est$ under the design in Equation (ref) is centered at zero. We let $\balancecoeffs$ be totally bounded in what follows. The proposition shows that in this case $\acceptpolar$ is a “nice” set: symmetric, convex, and with non-empty interior, satisfying the conditions of Assumption (ref).

Dimension Reduction. The oracle region $A = \{x: |\gammatrue'x| \leq \epsilon\}$ reduced the rerandomization problem to one dimension for arbitrary $\dim(h)$. Similarly, the minimax acceptance region $\acceptpolar = \epsilon \balancecoeffspolar$ can be viewed as a “soft” form of dimension reduction. To see this, note that the region $\acceptpolar$ is very stringent about imbalances $\rootn (\hbarone - \hbarzero)$ aligned with our belief set $\balancecoeffs$, but can allow large imbalances in directions approximately orthogonal to $\balancecoeffs$, effectively downweighting these directions in the space of covariates $h$. This effect can be seen in the following example, depicted in Figure (ref).

More generally, the following lemma provides a useful characterization of the acceptance region $\acceptpolar = \epsilon \balancecoeffspolar$ from Theorem (ref) for a large family of specifications of the belief set $\balancecoeffs$. To state the lemma, recall that $|x|_p = (\sum_j |x_j|^p)^{1/p}$ for $p \in [1, \infty)$ and $|x|_{\infty} = \max_j |x_j|$. For $p \in [1, \infty]$, denote $\ballp(0, 1) = \{x: |x|_p \leq 1\}$.

Minimizing Computational Cost

Intuitively, by ignoring imbalances $\imbalance = \rootn(\hbarone - \hbarzero)$ approximately orthogonal to our beliefs $\balancecoeffs$, we can “stretch” the acceptance region $\acceptpolar$ in directions unlikely to cause large estimation errors, increasing the probability of acceptance $P(\imbalance \in A)$. Since the expected number of independent randomizations until acceptance is $P(\imbalance \in A)\inv$, we can view this as minimizing the computational cost of rerandomization, subject to a bound on estimation error. This intuition is formalized in Theorem (ref) below. To state the theorem, we first define the family of possible limiting distributions of $\est$ consistent with our beliefs $\gammaoptmat \in \balancecoeffs$ and choice of acceptance region $A \sub \mrh$.

Limiting Distributions. We showed above that $\rootn(\est - \sate) | \Wn \convwprocess \asympdisttrue$ for $\asympdisttrue = \normal(0, V(\gammaoptmat)) + \gammaoptmat'\zhs$. Since $\gammaoptmat$ is unknown, define a family of possible limiting distributions of $\est$ by $\mc L_{\balancecoeffs} = \{\asympdistgamma: \gamma \in \balancecoeffs, A \sub \mrh \}$, with each $\asympdistgamma = \normal(0, V(\gamma)) + \gamma'\zhs$ a sum of independent RV's. For any distribution in this family, the conditional asymptotic bias of $\est$ given realized covariate imbalances $\zhs$ is $\bias(\asympdistgamma | \zhs) \equiv E[\asympdistgamma | \zhs]$. Our main result shows that the polar acceptance region $\acceptpolar = \epsilon \balancecoeffspolar$ minimizes asymptotic computational cost $P(\zh \in A)\inv$, subject to a strict constraint on conditional bias, uniformly over all limiting distributions consistent with our beliefs.

The final statement of the theorem shows that if $\balancecoeffs$ is well-specified ($\gammaoptmat \in \balancecoeffs$), setting $\acceptpolar = \epsilon \balancecoeffspolar$ bounds the magnitude of the conditional asymptotic bias $E[\asympdisttrue | \zhstrue]$ of the GMM estimator $\est$ above by $\epsilon$. By the law of total variance, this implies that the variance $\var(\asympdisttrue)$ of the asymptotic distribution $\rootn(\est - \thetan) \convwprocess \asympdisttrue = \normal(0, \vtheta) + \gammaoptmat'\zhstrue$ is within $\epsilon^2$ of the optimal partially linear variance $\vtheta$ in Equation (ref).

Results closely related to Theorem (ref) can also be found in the previous work of liu2023, who derive optimal Mahalanobis-style completely rerandomized designs under a Bayesian criterion, with Gaussian prior on $\gammaoptmat$.

Beliefs From Pilot Data

Next, we discuss an alternative strategy that uses pilot data to specify the set $B$ in a data-driven way. Suppose we have access to $\datapilot \indep (\Wn, \Dn)$ of size $m$. Suppose $\sqrt{\npilot} (\wh \gamma_{pilot} - \gammatrue) \approx \normal(0, \Sigmapilot)$ for some pilot estimator $\wh \gamma_{pilot}$, discussed below. Consider forming the Wald region $\coeffpilot = \{\gamma: m (\wh \gamma_{pilot} - \gamma)' \Sigmapilot \inv (\wh \gamma_{pilot} - \gamma) \leq c_{\alpha}\}$ using critical value $P(\chi^2_{\dimh} \leq c_{\alpha}) = 1-\alpha$ for $\alpha \in (0, 1)$. Equivalently, one can write this Wald region as

Viewing this $1-\alpha$ confidence region as a belief set, Lemma (ref) above implies that the corresponding minimax acceptance region is

Note that the acceptance region $\wh A_{pilot}$ expands as the pilot size $m$ is larger. This reflects smaller uncertainty about the true parameter $\gammatrue$, and thus less adversarial worst case imbalance $\sup_{\gamma \in \coeffpilot} |\gamma' \rootn(\hbarone - \hbarzero)|$. Conversely, $\wh A_{pilot}$ shrinks as the confidence parameter $\alpha$ and variance estimate $\Sigmapilot$ increase, reflecting greater uncertainty and a more conservative approach to covariate imbalances. Our next result shows that rerandomization with acceptance region $\wh A_{pilot}$ controls the variance of the residual imbalance $\residualvar = \gammatrue'\zh | \zh \in \wh A_{pilot}$ with high probability, marginally over the realizations of the pilot data. The result is an immediate consequence of Theorem (ref) and Theorem (ref).

Formally, the pilot estimate of $\gammatrue$ and Wald region could be constructed as in robinson88. A simpler practical approach suggested by the theory is to let $\gammapilot, \Sigmapilot$ be point and variance estimators from the regression $Y_{T} \sim 1 + h + \psi$, for the “tyranny of the minority” (lin2013) outcomes $Y_T = (1-\propfn)DY / \propfn + p(1-D)Y / (1-p)$, noting that $E[Y_T | W] = (1-p)Y(1) + pY(0) = \ylevel$.

Restoring Normality

In this section, we study optimal linearly adjusted GMM estimation under stratified rerandomization. We show that, to first order, optimal ex-post linear adjustment completely removes the impact of the acceptance region $A$ and imbalance term $\residualvar$, restoring asymptotic normality. This enables standard t-statistic and Wald-test based inference on the parameters $\thetan$ and $\thetatrue$ under stratified rerandomization designs, provided in Section (ref) below. We also describe a novel form of double robustness to covariate imbalances from combining rerandomization with ex-post adjustment.

Let $w$ denote the covariates used for ex-post adjustment and suppose $E[|w|_2^2] < \infty$.

First, we extend Corollary (ref) to provide asymptotics for the adjusted GMM estimator under pure stratification ($A = \mr^{\dimh}$).

A version of this result was given in cytrynbaum2024adjustment for the special case $\thetatrue = \ate$. Motivated by Proposition (ref), we define the optimal linear adjustment coefficient as the minimizer of the asymptotic variance $\vtheta(\alpha)$, in the positive semidefinite sense.

Optimal Adjustment Coefficient. Define the coefficient

Note that if $w = h$ then $\alphaopt = \gammaoptmat$ in Theorem (ref). If $E[\var(w | \psi)] \succ 0$, then the unique minimizer of Equation (ref) is the partially linear regression coefficient matrix $\alphaopt = E[\var(w | \psi)] \inv E[\cov(w, \matgmm \diff(W, \thetatrue) | \psi)]$. Observe that $\alphaopt$ varies with the stratification variables $\psi$, as observed in cytrynbaum2024 and bai2024adjustment for the case of ATE estimation. The main result of this section shows that adjustment by a consistent estimate of $\alphaopt$ restores asymptotic normality.

Two-step Adjustment. For nonlinear models, the optimal coefficient $\alphaopt$ may depend on the unknown parameter $\thetatrue$. This suggests a two-step adjustment strategy:

Similar to two-step efficient GMM, this process can be iterated until convergence to improve finite sample properties. One feasible estimator $\alphaest \convp \alphaopt$ is given in the following theorem. To state the result, define the within-group partialled covariates $\wicheck = \wi - \sum_{j \in \group(i)} w_j$, where $\group(i)$ is the group containing unit $i$ in Definition (ref). Let $\matgmmest \convp \matgmm$ estimate the linearization matrix and denote the score evaluation $\momesti \equiv \mom(D_i, R_i, S_i, \est)$. Define the adjustment coefficient estimator

In some cases, $\alphaopt$ may not depend on $\thetatrue$. For example, if $\diff(W, \theta) = \diff_1(\psi, \theta) + \diff_2(W)$ then $\alphaopt = E[\var(w | \psi)] \inv E[\cov(w, \matgmm \diff_2(W) | \psi)]$. In such cases, one-step optimal adjustment is possible.

One-step adjustment is possible in many linear GMM problems, including the best linear predictor of treatment effect heterogeneity parameter in Example (ref).

Double Robustness from Rerandomization

Theorem (ref) shows that stratified rerandomization has (approximately) the same first-order efficiency as optimal ex-post linear adjustment tailored to both the stratification and GMM problem.\footnote{For the case without stratification, this equivalence was originally shown in li2018rerandomization.} However, our simulations and empirical application show that in finite samples stratified rerandomization can perform significantly better than ex-post adjustment, and further efficiency gains are possible by combining both methods. In this subsection, we provide a brief theoretical justification for this phenomenon, showing that combining rerandomization and adjustment provides a novel form of double robustness to covariate imbalances.

For simplicity, consider the case of $\sate$ estimation with $\psi = 1$. For the difference of means estimator $\est$, by a simple calculation $\est - \sate = \en[\yleveli | \Di=1] - \en[\yleveli | \Di=0]$. Let $\ylevel = c + \gammaoptmat'h + e$ with $e \perp (1, h)$ be the decomposition from Equation (ref). Then we can decompose the estimation error $\est-\sate$ into imbalances in $h$ and $e$:

Rerandomizing until $\rootn (\hbarone - \hbarzero)$ is small shrinks the first imbalance term, ideally making its variance negligible relative to $\var(e) = \var(\ylevel - \gammaoptmat'h)$. Suppse $h = w$ so the adjustment coefficient $\alphaopt = \gammaoptmat$. Writing $\alphaest = \wh \gamma$, we have $\estadj = \est - \wh \gamma'\en[\hti \hi] = \est - \wh \gamma'(\hbarone - \hbarzero)$,

This decomposition shows a novel form of double robustness from combining rerandomization with ex-post adjustment. If the estimation error $\gammaoptmat - \wh \gamma$ is large, then the first imbalance term above may still be negligible as long as we rerandomized until $\rootn(\hbarone - \hbarzero)$ is small enough. For example, in the LHS of Figure (ref) the coefficient $\gammaoptmat$ is not estimated well for small $n$ and large $\dim(w)$, so adjustment without rerandomization performs poorly. This effect is exacerbated by stratification, since the partialling operation $\wicheck = \wi - \sum_{j \in \group(i)} w_j$ tends to decrease the variance of the regressors $w_i$, making estimation of $\alphaopt$ more difficult.\footnote{For example, in our empirical application the condition number of the design matrix $\en[\wicheck \wicheck]$ increases as we stratify more finely.} However, combining adjustment and rerandomization still performs well due to double robustness.

Equation (ref) has a “product of errors” structure, similar to the product of nuisance estimation errors for doubly-robust estimators in the literature on Neyman orthogonal estimating equations (e.g.\ chernozhukov2017dml). This shows that even when both $\wh \gamma - \gammaoptmat$ and $\hbarone - \hbarzero$ are small, we can get an extra benefit from combining the two methods. This is shown in the RHS of Figure (ref) with $n=500$, where adjustment is competitive, but rerandomization still performs better, and stratification + rerandomization is even more efficient due to this product structure. Such double robustness also holds for more general casual parameters and GMM estimators. Let $h = w$ and consider the partially linear decomposition $\matgmm \diff(W, \thetatrue) = \gammaoptmat'h + t(\psi) + e$ with $e \perp h$ and $E[e | \psi] = 0$ and $\bar t_d = \en[t(\psii) | \Di=d]$. Our work shows that

The second equality is a consequence of fine stratification on $\psi$.

Summarizing, this discussion highlights a double robustness property that explains the additional finite-sample precision gains from combining rerandomization with ex-post adjustment. This effect is likely to be especially important in regimes where the optimal adjustment coefficient $\alphaopt$ is poorly estimated, such as for small $n$, large $\dim(w)$, ill-conditioned design matrix $\en[\wicheck \wicheck'] \approx E[\var(w | \psi)]$. A full theory of high-dimensional stratification, rerandomization, and ex-post adjustment is beyond the scope of the current work, but this is an interesting area for future research.\footnote{There are analytical complications from conditioning on $\rootn(\hbarone - \hbarzero) \in A$ e.g.\ with $\dim(h)$ growing. A recent breakthrough on this question was achieved by the careful analysis of wang2022 for the case of complete rererandomization.}

Variance Bounds and Inference Methods

In this section, we provide methods for inference on generic causal parameters under stratified rerandomization designs. We make crucial use of asymptotic normality of the optimally adjusted GMM estimator $\estadj$ developed in the previous section. The asymptotic variance for estimating the finite population parameter $\thetan$ is generally not identified. To enable inference, we provide novel identified upper bounds on the variance, allowing for conservative inference that still reflects the precision gains from stratified rerandomization. The asymptotic variance for estimating the superpopulation parameter $\thetatrue$ is identified, and in this case we provide asymptotically exact inference methods.

Variance Bounds

First, we briefly review the classical variance bounds for $\thetan = \sate$ estimation under completely randomized assignment. In this case, we have $\rootn(\est - \sate) \convwprocess \normal(0, \vtheta)$ with $\vtheta = \var(D)\inv \var(\ylevel)$ for $\ylevel = (1-p)Y(1) + p Y(0)$. The variance $\var(\ylevel) \propto \cov(Y(1), Y(0))$. Since $Y(1)$ and $Y(0)$ are never simultaneously observed, $\vtheta$ is not identified. Let $\hk_d = \var(Y(d))$ and $\tau = Y(1)-Y(0)$. The Cauchy-Schwarz inequality $|\cov(Y(1), Y(0))| \leq \hksd_1 \hksd_0$ and some algebra produces the bounds

Both upper bounds were proposed in neyman1990. Theorem (ref) below extends the sharper bound to generic finite population causal parameters, accounting for both design-time stratified rerandomization and optimal ex-post adjustment.

To develop the bounds, recall from Theorem (ref) that $\rootn(\estadj - \thetan) \convwprocess N(0, \vadj)$ with $\vadj = \vard \inv E[\var(\matgmm \diff(W, \thetatrue) - \alphaopt'w | \psi)]$, where $\alphaopt = E[\var(w | \psi)] \inv E[\cov(w, \matgmm \diff(W, \thetatrue) | \psi)]$ was the optimal adjustment coefficient. By definition, $\matgmm \diff(W, \thetatrue) = \vard \matgmm (\momone(W, \thetatrue) - \momzero(W, \thetatrue))$. Then the adjustment coefficient may be expanded as $\alphaopt = \adjcoeffone - \adjcoeffzero$ for coefficients $\beta_d = E[\var(w | \psi)] \inv E[\cov(w, \vard \matgmm \mom_d(W, \thetatrue) | \psi)]$. Denote $\mom_d = \mom_d(W, \thetatrue)$ and define the “within-arm” influence functions $\momadj_d \equiv \vard \matgmm \mom_d - \adjcoeff_d'w$. Tighter bounds are possible by targeting a fixed scalar contrast $c'\thetan$ for some $c \in \mr^{\dimtheta}$. From Theorem (ref), we have $\rootn(c'\estadj - c'\thetatrue) \convwprocess N(0, \vadj(c))$ for $\vadj(c) = c'\vadj c$. In terms of $\momadj_d$, this is

Similarly to above, $\vadj(c) \propto E[\cov(c'\momadjone, c'\momadjzero | \psi)]$ where the cross-term is generically not identified, since $\momadjone$ and $\momadjzero$ are not simultaneously observed. However, denoting $\hkadj_d(c) = E[\var(c'\momadj_d| \psi)]$ we have the following simple upper bound:

We provide a consistent estimator of the bound $\vadjbound(c)$ in Section (ref) below. The next example shows how Theorem (ref) generalizes the classical Neyman bounds for the simple case of inference on $\thetan = \sate$ under pure stratified randomization ($A = \mrh$) and optimal ex-post adjustment.

Inference on the Finite Population Parameter

Building on the previous section, we construct a consistent estimator of the variance upper bound $\vadjbound(c)$, enabling asymptotically conservative inference on linear contrasts of the finite population parameter $c'\thetan$ under general designs.

We begin with some definitions. Let $\groupset_n$ denote the set of groups (strata) constructed in Definition (ref). For $s \in \groupset_n$, denote number of treated $a(\group) = \sum_{i \in \group} \Di$ and group size $k(\group) = |\group|$. For any $\matgmmest \convp \matgmm$ define estimators of the optimal within-arm adjustment coefficients $\adjcoeff_d$ above by $\adjcoeffest_d = \vard \en[\wicheck \wicheck']\inv \cov_n(\wicheck, \matgmmest \momesti | \Di=d)$. Note that $\adjcoeffoneest - \adjcoeffzeroest = \alphaest$, our estimator of the optimal adjustment coefficient in Section (ref). For $\momesti \equiv \mom(D_i, X_i, S_i, \estadj)$, define $\momestadji \equiv \vard \matgmmest \momesti - \Di \adjcoeffoneest'w_i - (1-\Di) \adjcoeffzeroest'w_i$. First, suppose each group has at least two treated and control units, $2 \leq a(\group) \leq k(\group) - 2$ $\forall s \in \groupset_n$, setting

Collapsed Strata. If number of treated units $a(\group) = 1$ or $a(\group) = k(\group)-1$, as in matched pairs designs, the estimators above do not exist. In this case, we follow\footnote{See abadie2008, bai2021inference, cytrynbaum2024, bai2024efficiency for recent use of this method for inference on superpopulation parameters. In particular, bai2021inference showed asymptotic exactness of the collapsed strata method for matched pairs designs under the matching condition above.} the method of collapsed strata (hansen1953), first agglomerating the original groups $\group \in \groupset_n$ into larger groups satisfying $2 \leq a(\group) \leq k(\group) - 2$. For example, in a matched triples design with $p=1/3$, we agglomerate two triples into a larger group $\group'$ of $6$ units with $a(\group') = 2$. To do so, for each $\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 matching 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 derigs1988 non-bipartite matching algorithm. Define $\groupsetnu_n = \{\group \cup \groupmatching(\group): \group \in \groupset_n\}$ to be the enlarged groups. If $a(\group) = 1$ or $a(\group) = k(\group)-1$, we replace $\groupset_n$ with the larger groups $\groupsetnu_n$ in the definitions of $\varestone$ and $\varestzero$.

Variance Estimator. Finally, define $\varestdiffone = \en[\frac{\Di}{\propfn} \momestadji \momestadji'] - \varestone$ and $\varestdiffzero = \en[\frac{1-\Di}{1-\propfn} \momestadji \momestadji'] - \varestzero $. The proof of Theorem (ref) below shows that $c' \wh u_d c \convp \hkadj_d(c)$ from Theorem (ref), suggesting the variance estimator

To formalize this result we require a slight strengthening of GMM Assumption (ref).

Then the confidence interval $\cifin \equiv [c'\estadj \pm z_{1-\alpha/2} \varestdiffadj(c)\half / \rootn]$ has coverage $P(c'\thetan \in \cifin) \geq 1-\alpha - o(1)$ by Theorem (ref) and Theorem (ref).

The main result is stated for adjusted GMM estimation under stratified rerandomization, with ex-post adjustment to restore normality. For the case of pure stratification (no rerandomization) without adjustment, we can just set $w=0$ in the formulas above, obtaining a specialization $\varestdiff(c)$ of $\varestdiffadj(c)$. We summarize this in a corollary:

Inference on the Superpopulation Parameter

The asymptotic variance $V = \vphi + \vadj$ for adjusted estimation of $\thetatrue$ under stratified rerandomization (Theorem (ref)) is identified. In this case, we can modify the approach above to provide asymptotically exact inference methods. Additionally define

With this extra definition in hand, set $\wh V = \var_n(\momestadji) - \vard (\varestone + \varestzero - \varestcross - \varestcross')$.

By Theorem (ref), $\rootn(\estadj - \thetatrue) \convwprocess N(0, \vphi + \vadj)$, so the result above allows for asymptotically exact joint inference on $\thetatrue$ e.g.\ using standard Wald-test based confidence regions. For example, the interval $\cipop \equiv [c'\estadj \pm z_{1-\alpha/2} (c'\wh V c)\half / \rootn]$ has $P(c'\thetatrue \in \cipop) = 1-\alpha - o(1)$. Similarly to above, this CI can be specialized to pure stratification without adjustment by setting $w = 0$.

Simulations

In this section, we use simulations to study the finite-sample properties of various designs and estimators analyzed above. We consider data generated as $Y(d) = m_d(r) + e_d$ for observables $r$, varying the covariates $\psi$, $h$, and $w$ used for stratification, rerandomization, and adjustment respectively. In models 1-3, we consider quadratic outcome models of the form \[ Y(d) = c_d + r'\linearcoeff_d + r'\quadcoeff_d r + e_d. \] We vary $m = \dim(r)$, setting parameters $\quadcoeff_d$ and $\linearcoeff_d$ as follows:

In Model 1, all covariates have equal importance. In Models 2-4, we think of $r_1$ as a baseline outcome with more importance than $r_{2:m}$. This asymmetric structure arises frequently in practice due to the relatively high predictive power of baseline outcomes for endline outcomes. The covariates are generated $r \sim \normal(0, \Sigma)$. For Tables (ref) and (ref), we let $\Sigma = I_m$. For Table (ref) below, we set $\Sigma_{ii} = 1$ and $\Sigma_{ij} = (1/2)(m-1)\inv$ for $i \not = j$. The residuals $(e_1, e_0) \sim \normal(0, \tilde \Sigma)$ with $\var(e_d) = 4$, $\corr(e_1, e_0) = 0.8$, and $(e_1, e_0) \indep r$. We set $p = 1/2$ in all simulations, corresponding to matched pairs rerandomization for $\psi, h$ non-constant.

In Table (ref), we compare the efficiency and inference properties of various designs for estimating $\thetan = \sate$. The design C refers to complete randomization. Design S is full stratification: for model 1, we set $\psi = r$, while for models 2-4, we let $\psi_1 = \sqrt{2} r_1$ and $\psi_{2:m} = r_{2:m}$ in the matching algorithm, putting more weight on the covariate believed to be important a priori.\footnote{We match using the algorithms in bai2021inference for $p = 1/2$ and cytrynbaum2024 for $p \not = 1/2$.} Design SR is stratified rerandomization, with univariate $\psi = r_1$ and $h = r_{2:m}$. In this first simulation, we use simple Mahalanobis-style rerandomization (Example (ref)), with acceptance probability $\alpha = 1/500$. $\est$ is the unadjusted GMM estimator of Definition (ref), while $\estadj$ is the optimally adjusted GMM estimator of Theorem (ref) with adjustment covariates $w = h$. For each model, we normalize the MSE of $\est$ under complete randomization C to $1$. All inference results are based on the adjusted estimator $\estadj$, comparing performance across different designs. In particular, Cover Fin.\ refers to coverage of $\thetan$ using the (conservative) finite population variance bound estimator $\varestdiff(c)$ in Section (ref) and confidence interval $\cifin$. Cover Pop.\ presents coverage of $\thetatrue$ for $\wh C_{pop}$, using asymptotically exact variance estimator $\wh V$ from Section (ref). CI Width Fin.\ and Pop.\ report the width confidence intervals, normalized so that the width of $\wh C_{pop}$ is $1$ for $\estadj$ and design C.

Next, we summarize a few important findings from Table (ref). Stratified rerandomization SR is the most efficient design across all specifications and for both estimators $\est$ and $\estadj$. While ex-post optimal adjustment and rerandomization have (approximately) the same effect asymptotically (Theorem (ref)), there is an additional finite sample efficiency gain from combining rerandomization and adjustment (SR and $\estadj$), due to the double robustness property discussed in Section (ref). This effect is especially pronounced for small $n$ and large $\dim(r)$, as shown previously in Figure (ref), due to poor estimation of the optimal adjustment coefficient $\gammaoptmat$. For inference, CI Width is slightly larger for S, SR than for C, despite SR being the most efficient. Under design \textbf{C}, the estimators $\wh V$ and $\varestdiff$ tend to be too small, leading to undercoverage.\footnote{This could be fixed by a sample-splitting or jackknife approach for GMM variance estimation under (non-iid) completely randomized treatment assignment, but this is not our focus here.} By contrast, coverage is approximately nominal for designs \textbf{S} and \textbf{SR}. Note that $\wh C_{fin}$ is often much smaller than $\wh C_{pop}$, showing that experimenters only interested in covering $\thetan$ can potentially report smaller confidence intervals. We provide additional results for Model 4 in Figure (ref), letting $n=150$ and varying $\dim(r)$. In the figure, \textbf{CR} refers to complete rerandomization and “quadratic” refers to the Mahalanobis design in Example (ref). Pure fine stratification is competitive for small $\dim(r)$, while stratified rerandomization is preferred for $\dim(r) > 2$.

In Table (ref) we compare different types of stratified rerandomization acceptance criteria. MH is Mahalanobis rerandomization, as in Table (ref). Prop is the propensity-based rerandomization in Definition (ref), using Logit $L(x) = (1 + e^{-x}) \inv$ and $X = (1, w)$. Designs Opt1 and Opt2 refer to the optimal acceptance regions in Section (ref). The belief sets are both well-specified, with either high uncertainty $\balancecoeffs_1 = \{x: |x-\gammaoptmat|_2 \leq 1\}$ or low uncertainty $\balancecoeffs_2 = \{x: |x-\gammaoptmat|_2 \leq 1/10\}$, respectively. In all designs, we set the balance threshold $\epsilon(\alpha)$ so $P(\zh \in A) = 1/500$. Finally, in Best1 and Best2 we rerandomize by implementing the best allocation out of either $k = 500$ or $k=2500$ stratified draws, according to the minimal Mahalanobis imbalance metric. Note that such “best-of-k” stratified rerandomization designs are not formally covered by our theory.\footnote{Recent work by wang2024 provided the first formal results for “best-of-k” designs in the case without stratification.} In addition to $\thetan = \sate$, we also provide efficiency and inference results for the treatment effect heterogeneity parameter from Example (ref). In particular, let $\alpha_n = \argmin_{\alpha} \en[(Y_i(1) - Y_i(0) - \alpha'(1, r_{1i}))^2]$. We define $\thetan$ to be the coefficient on $r_{1}$, denoting $\thetan = \cate$ in the table. Cover Pop.\ and CI Width Pop. refer to inference on the corresponding superpopulation parameter $\thetatrue$.

Next, we summarize a few findings from Table (ref). Theorem (ref) showed that Prop was first-order equivalent to MH, and this is supported by finite-sample evidence in the table. We find that best of $k$ style rerandomization and Mahalnobis rerandomization with acceptance probability $\alpha \approx 1/k$ are indistinguishable in practice. In particular, our inference methods also work well for this design. We don't find major finite sample efficiency improvements from using the optimal acceptance regions in Section (ref). We provide additional results for Model 4 in Figure (ref), showing that Opt1 and Opt2 reduce the curse of dimensionality for rerandomization, since we are able to downweight less important dimensions of $h$. Finally, in Table (ref), we provide additional simulation results for estimating the heterogeneity parameter $\thetan = \cate$. In particular, Example (ref) showed that if the experimenter is interested in treatment effect heterogeneity along dimension $r_1$, then they should balance variables $\psi, h$ and $w$ predictive of the interaction $\ylevel r_1$, not just the outcome level $\ylevel$. The designs in Table (ref) are as above for no interactions (Inter.\ $=$ No). In the “Yes” case, we add interactions so that rerandomization and ex-post adjustment covariates $h = (r, r \cdot r_1)$, and $w = (r, r \cdot r_1)$, keeping $\psi = r_1$. This significantly increases efficiency for $\estadj$ under design C, with smaller efficiency gains for design SR.

Empirical Application

In this section, we apply our methods to data from the “Opportunity Knocks” experiment in angrist2013. The authors randomized eligibility to receive payment for high grades to first and second year students at a large Canadian university. They estimated the effect of the program on future student GPA, successful graduation, and other outcomes. They measured several baseline covariates, including high school GPA, sex, age, native language, and parent's education. Randomization was coarsely stratified on year in college, sex, and quartiles of high school GPA within year-sex cells, with approximately $p = 3/10$ of $n=1203$ students assigned to receive incentives. Some students assigned treatment $Z = 1$ (viewed as an instrument below) did not engage with the program either by checking their earnings or making contact with the program advisor. The authors view this as noncompliance with the instrument $Z$ and estimate both intention-to-treat (ITT) effects and effects on compliers ($\late$). Let $D \in \{0, 1\}$ denote endogeneous decision to engage with the program, with $D(z)$ the potential treatments, $Y(d)$ the potential outcomes, and $T(z) = Y(D(z))$ the ITT potential outcomes with realized outcome $T = Y(D(Z)) = Y$. angrist2013 estimate ITT-style treatment effect heterogeneity along several dimensions, such as gender and student reported financial need.

In what follows, we use this data to study the efficiency and inference properties of various designs and estimators, including complete randomization, fine stratification on different variable sets, and coarse stratification as in the original study, including both rerandomized and standard versions of each. To do so, we follow the common approach (e.g.\ li2018rerandomization, bai2020pairs) of imputing the missing potential outcomes, which allows us to simulate the MSE, coverage properties, and CI width under various counterfactual designs. In particular, we set $\wh T(z) = T = Y$ if $Z=z$ in the observed data, and impute $\wh T(z) = \wh m_z^T(X) + \hksdest_z^T(X) \epsilon_z$ if $Z=1-z$, where $\wh m_z(X)$, $\hksdest_z(X)$ are estimated using cross-validated LASSO and random forests applied to $11$ baseline covariates their full pairwise interactions. The residual $\epsilon_z \sim \normal(0, 1)$. We similarly impute missing potential treatments $\wh D(z)$ for all units with $\wh D(z) = D$ if $Z=z$. See Section (ref) for more details on this procedure.

Given imputed data $(X_i, \wh T_i(z), \wh D_i(z))$ for units $i=1, \dots, 1203$, we simulate an experiment of size $n$ as follows: (1) sample $(X_i, \wh T_i(z), \wh D_i(z))_{i=1}^n$ with replacement, (2) draw instrument assignments $\tilde Z_{1:n}$ e.g.\ by stratified rerandomization with covariates $\psii, \hi \sub X_i$. Then we (3) observe realized treatments $\tilde D_i = \wh D_i(\tilde Z_i)$ and outcomes $\tilde Y_i = \tilde T_i = \wh T_i(\tilde Z_i)$ and (4) form estimators $\est$ and $\estadj$ and confidence intervals $\wh C_{fin}$ and $\wh C_{pop}$ for the causal parameters $\sate$, $\late$, $\cate$, and $\clate$ described below.

We let rerandomization and adjustment sets $h$, $w$ include all $11$ covariates above, as well as the pairwise interactions of HS GPA, sex, year, and mother and father's education with both financial need $F \in \{0, 1\}$ and HS GPA $G \in \mr$, for a total of $21$ adjustment covariates. The interactions are motivated by our desire to estimate treatment effect heterogeneity along the dimensions $F$ and $G$, as discussed in Example (ref). We simulate the following designs: C is complete randomization, and CR is rerandomization. S is the original study design (coarse stratification), and SR is its rerandomized version using covariates $h$ above. F is fine stratification on HS GPA, and FR is finely stratified rerandomization. \textbf{F+} is fine stratification on HS GPA, sex, and year and similarly for the rerandomized version \textbf{FR+}.\footnote{For the last four designs \textbf{F}-\textbf{FR+}, we remove covariates included in $\psi$ from $w$ and $h$, to ensure that $E[\var(w | \psi)] \succ 0$, as discussed in Section (ref). This does not affect first-order efficiency.} We let $p=3/10$ and $n=1200$ for all.

We present empirical results for several causal estimands. Table (ref) presents results on the ITT estimands $\sate = \en[T_i(1) - T_i(0)]$ and “$\cate$,” the coefficient on $x_i$ in \[ \thetan = \argmin_{\theta} \en[(T_i(1) - T_i(0) - \theta'(1, x_i))^2]. \] We consider both $x_i = F_i \in \{0, 1\}$, an indicator for student financial stress, and $x_i = G_i \in \mr$, the student's HS GPA. For $x_i = F_i$, this has a simple interpretation as the difference in ITT effects between students with and without financial stress: \[ \cate = \en[T_i(1)-T_i(0)|F_i=1] - \en[T_i(1)-T_i(0)|F_i=0]. \] Table (ref) presents efficiency and inference results for LATE-style treatment effects on compliers. In particular, if $C_i = \one(D_i(1) - D_i(0) > 0)$ is a compliance indicator then $\late = \en[Y_i(1)-Y_i(0) | C_i=1]$ and CLATE (Example (ref)) is the coefficient on $x_i$ in \[ \thetan = \argmin_{\theta} \en[(Y_i(1) - Y_i(0) - \theta'(1, x_i))^2 | C_i=1]. \] In both tables, Cover Pop.\ and CI Width Pop.\ refer to inference on the corresponding superpopulation estimands $\thetatrue$, e.g.\ $\thetatrue = \argmin_{\theta} E[(Y(1) - Y(0) - \theta'(1, x))^2 | C=1]$ for $\thetan = $ CLATE and $\thetatrue = E[T(1)-T(0)] = \ate$ for $\thetan = \sate$. The MSE of $\estadj$ and the CI width of $\wh C_{pop}$ are normalized to $1$ under design C.

We briefly summarize our main findings from the tables. The efficiency differences between designs are more pronounced for the heterogeneity variables CATE and CLATE than for average effects SATE and LATE. Finely stratified rerandomization FR is the efficient for the majority of estimands, while SR is slightly more efficient for estimating treatment effect heterogeneity CATE and CLATE along the financial need variable $F \in \{0, 1\}$. Confidence intervals broadly have correct coverage. The width of $\wh C_{fin}$ for inference on $\thetan$ is slightly smaller than $\wh C_{pop}$ for inference on $\thetatrue$ on average, with the largest improvements for estimating CATE (GPA) and CLATE (GPA).

Discussion and Recommendations for Practice

In general, we recommend experimenters finely stratify on a few variables expected to be highly predictive of outcomes, as well as their interactions with covariates for which treatment effect heterogeneity is of interest.\footnote{More generally, “highly predictive” is defined in the estimand-specific sense of Equation (ref).} Based on our work, we strongly recommend adding a rerandomization step to this procedure to balance the remaining baseline covariates, e.g.\ using the stratified Mahalanobis design in Example (ref) or the optimized designs in Section (ref), if the researcher has a strong prior. Separating the baseline covariates into stratification and rerandomization tiers is an easy way to balance linear functions of the less important covariates, without degrading match quality when finely stratifying on the most important covariates like baseline outcomes.

Our work in Section (ref) showed that combining stratified rerandomization with optimal ex-post adjustment provides a form of double robustness to covariate imbalances between treatment groups. This effect seemed to matter in our simulations and empirical application, where rerandomization plus adjustment sometimes performed much better than either method alone. We recommend experimenters adopt this doubly-robust approach, using the stratification-tailored adjustment coefficients in Section (ref). In Section (ref), we provide the first valid methods for inference on both finite population and superpopulation GMM parameters under stratified rerandomization, enabling inference in new settings. Our work also provides new tools even for some settings with existing inference methods. For example, experimenters can use the finite population methods in Section (ref) for more powerful inference than is currently available in the setting of stratification without rerandomization, e.g.\ in experiments in a convenience sample where we only require coverage of the finite population parameter.

This discussion also touches on several practical questions for which the theory does not give concrete guidance. For example, exactly which and how many covariates should we finely stratify on and which should we rerandomize in a given experiment to maximize finite sample efficiency? It may be possible to formally develop a high dimensional theory of stratified rerandomization in future work. However, even with such new technical results, optimizing the partition of covariates into stratification vs.\ rerandomization sets would likely require knowledge of DGP-specific constants that are not estimable at design-time before outcomes are observed, and may be difficult to specify beliefs over.\footnote{For example, optimizing which variables to stratify on vs.\ rerandomize would likely require researchers to specify a prior on objects like the Lipschitz coefficient of the function $\psi \to E[a(W, \thetatrue) | \psi]$.} Providing practically useful and implementable theoretical guidance for such design issues remains a difficult open question for future work.