Asymptotic inference with flexible covariate adjustment under rerandomization and stratified rerandomization

We consider a randomized experiment with $n$ individuals. Each individual $i\ (i=1,\dots,n)$ has an outcome $Y_{i}$, a non-missing indicator of outcome $R_{i}$ ($R_{i}=1$ if $Y_{i}$ observed and $0$ if missing), a treatment assignment indicator $A_{i}$ ($A_{i}=1$ if assigned treatment and $0$ otherwise), and a vector of baseline covariates $\boldsymbol{X}_{i}$. The observed data are $(\boldsymbol{O}_{1},\dots,\boldsymbol{O}_{n})$ with $\boldsymbol{O}_{i}=(R_{i}Y_{i},R_{i},A_{i},\boldsymbol{X}_{i})$. We adopt the potential outcomes framework and define $Y_{i}(a)$ as the potential outcome if individual $i$ were assigned to treatment group $a,a\in\{0,1\}$. Similarly, we denote $R_{i}(a)$ as the potential non-missing indicator given treatment $a$. We assume causal consistency such that $Y_{i}=A_{i}Y_{i}(1)+(1-A_{i})Y_{i}(0)$ and $R_{i}=A_{i}R_{i}(1)+(1-A_{i})R_{i}(0)$. Denoting the complete data vector for each individual as $\boldsymbol{W}_{i}=(Y_{i}(1),Y_{i}(0),R_{i}(1),R_{i}(0),\boldsymbol{X}_{i})$, we make the following assumptions on $(\boldsymbol{W}_{1},\dots,\boldsymbol{W}_{n})$.

Assumption 1 is a standard condition for making inference under a sampling-based framework. It postulates a notional super-population of units, from which the observed sample is randomly drawn and for which the target estimands may be defined. We invoke Assumption 2 as a standard condition to accommodate missing outcomes in randomized experiments. This assumption is required for the doubly-robust estimators in Theorems 2, 4, 5, and 6, but is not required for other results without missing data.

Our goal is to estimate the average treatment effect, defined in a general form as

where $f(x,y)$ is a pre-specified function defining the scale of effect measure. For example, $f(x,y)=x-y$ corresponds to the treatment effect on the difference scale, whereas $f(x,y)=x/y$ corresponds to the risk ratio scale, which is a standard choice when the outcome is binary. Finally, by Assumption 1, the expectation operator in the estimand definition (and throughout the paper) is defined with respect to the super-population distribution $\mathcal{P}^{W}$.

Simple randomization assigns treatment via independent coin flips, that is, $A_{i}$ is independently determined by a Bernoulli distribution $\mathcal{P}^{A}$ with $P(A_{i}=1)=\pi\in(0,1)$. By design, each $A_{i}$ is independent of $(\boldsymbol{W}_{1},\dots,\boldsymbol{W}_{n})$, and the observed data $(\boldsymbol{O}_{1},\dots,\boldsymbol{O}_{n})$ are independent and identically distributed given Assumption 1.

Rerandomization improves simple randomization by controlling imbalance on a subset of measured baseline covariates, which we denote as $\boldsymbol{X}^{r}$ with $\boldsymbol{X}^{r}\subset\boldsymbol{X}$ (Morgan and Rubin,, 2012); we refer to $\boldsymbol{X}^{r}$ as rerandomization variables. Rerandomization involves the following three steps. First, we independently generate $A_{1}^{*},\dots,A_{n}^{*}$ from a Bernoulli distribution with $P(A_{i}^{*}=1)=\pi\in(0,1)$ as in simple randomization. In the second step, we compute the imbalance statistic on $\boldsymbol{X}^{r}_{i}$ and its variance estimator as

where $N_{1}=\sum_{i=1}^{n}A_{i}^{*}$, $N_{0}=n-N_{1}^{*}$ and $\overline{\boldsymbol{X}^{r}}=n^{-1}\sum_{i=1}^{n}\boldsymbol{X}^{r}_{i}$. Lastly, given a pre-specified balance threshold $t>0$, we check whether $\boldsymbol{I}^{\top}\{\widehat{Var}(\boldsymbol{I})\}^{-1}\boldsymbol{I}<t$. If true, the final treatment assignment $(A_{1},\dots,A_{n})$ is set to be $(A_{1}^{*},\dots,A_{n}^{*})$; otherwise, we repeat the steps until we obtain the first randomization scheme that satisfies the balance criterion.

The above rerandomization procedure considers the Mahalanobis distance to define the balance criterion, and $\boldsymbol{I}^{\top}\{\widehat{Var}(\boldsymbol{I})\}^{-1}\boldsymbol{I}$ follows a chi-squared distribution asymptotically. The balance threshold $t$ can then be selected to adjust the rejection rate. For example, letting $q$ denote the dimension of $\boldsymbol{X}^{r}$, then $t$ being the $\alpha$-quantile of $\mathcal{X}^{2}_{q}$ leads to an approximate rejection rate $1-\alpha$. In general, smaller $t$ corresponds to stronger control over imbalance; if $t\rightarrow+\infty$, then rerandomization reduces to simple randomization. From a statistical perspective, rerandomization introduces correlation among $(A_{1},\dots,A_{n})$ and hence correlation among observed data $(\boldsymbol{O}_{1},\dots,\boldsymbol{O}_{n})$, leading to a key complication in studying large-sample results of treatment effect estimators.

Beyond the Mahalanobis distance, rerandomization can also be more generally based on $\boldsymbol{I}^{\top}\widehat{\mathbf{H}}^{-1}\boldsymbol{I}<t$, where $\widehat{\mathbf{H}}$ is a positive definite matrix (e.g., a diagonal matrix collecting the sample variance for each component of $\boldsymbol{I}$). For our main theorems, we focus on rerandomization using the Mahalanobis distance, which leads to the most interpretable results. Nevertheless, we show how our results can be extended to accommodate such a more general distance function in Remarks 1 and 4.

M-estimator (van der Vaart,, 1998, Section 5) refers to a wide class of estimators that are solutions to estimating equations. Specifically, let $\boldsymbol{\theta}=(\Delta,\boldsymbol{\beta})\in\mathbb{R}^{l+1}$ be an $(l+1)$-dimensional vector of parameters, where $\Delta\in\mathbb{R}$ is the parameter of interest and $\boldsymbol{\beta}\in\mathbb{R}^{l}$ is an $l$-dimensional vector of nuisance parameters. An M-estimator $\widehat{\boldsymbol{\theta}}=(\widehat{\Delta},\widehat{\boldsymbol{\beta}})$ for $\boldsymbol{\theta}$ is the solution to

where $\boldsymbol{\psi}$ is a pre-specified $(l+1)$-dimensional estimating function. For example, $\boldsymbol{\psi}$ is the score function if $\widehat{\boldsymbol{\theta}}$ is obtained by maximum likelihood estimation. The M-estimator for $\Delta^{*}$ is $\widehat{\Delta}$. For $a=0,1$, we denote $\boldsymbol{O}(a)=(R(a)Y(a),R(a),a,\boldsymbol{X})$ and $||\cdot||_{2}$ as the $L_{2}$ matrix norm. Throughout, we assume the following regularity conditions for M-estimation.

Assumption 3 are largely moment and continuity assumptions to rule out irregular or degenerate M-estimators, which are similar to those invoked in Section 5.3 of van der Vaart, (1998). It is important to note that when the M-estimator is defined based on a working regression model, Assumption 3 does not necessarily imply the working model is correctly specified. For example, if the analysis of covariance (ANCOVA, Example 1 in Section 4) is used, Assumption 3 (2)-(6) simply reduce to assuming $\boldsymbol{W}$ has finite fourth moments and invertible covariance matrices, rather than assuming the ANCOVA model is the true data generating process.

Under simple randomization, it has been established that Assumptions 1 and 3 yield $\widehat{\Delta}\xrightarrow{p}\underline{\Delta}$, i.e., $\widehat{\Delta}$ converge in probability to a limit quantity $\underline{\Delta}$, and $\sqrt{n}(\widehat{\Delta}-\underline{\Delta})\xrightarrow{d}N(0,V)$, i.e., $\sqrt{n}$-rate asymptotic normality of $\widehat{\Delta}$ (van der Vaart,, 1998). To achieve consistent estimation, i.e., $\widehat{\Delta}\xrightarrow{p}\Delta^{*}$, one needs to choose appropriate working models (e.g., ANCOVA in randomized experiments) or make additional assumptions (e.g., Assumption 2 in the presence of missing outcomes) such that $\Delta^{*}=\underline{\Delta}$; we refer to Section 4 for detailed examples. In the next section, we formally describe the asymptotic behavior of $\widehat{\Delta}$ under rerandomization.

In practice, we may want to control the treatment imbalance on $\boldsymbol{X}^{r}$ at level $t$ and, meanwhile, eliminate the treatment imbalance on certain categorical variables. We refer to this design as stratified rerandomization (Wang et al., 2023c, ). Compared to rerandomization defined in Section 2.2, this design further ensures a perfectly balanced assignment within each stratum. To formally define this scheme, we first introduce stratified randomization without rerandomization. Stratified (block) randomization refers to a randomization scheme that achieves exact balance within each pre-specified stratum. Let $S\in\boldsymbol{X}$ be a categorical baseline variable encoding the randomization strata, e.g., $S$ taking values in $\mathcal{S}=\{\textup{male smoker},\ \textup{female smoker},\ \textup{male non-smoker},\ \textup{female non-smoker}\}$ if the randomization strata are defined by sex and smoking status. Within each randomization stratum, a size-$k$ permutation block with $\pi k$ ones and $(1-\pi)k$ zeros in random order is independently sampled to assign the treatment for the first $k$ participants (1 for treatment and 0 for control). When a permutation block is exhausted, a new one is sampled to assign the treatment for the next $k$ participants. The block size $k$ is chosen to make $\pi k$ an integer. Unlike simple randomization where treatment assignment is independent, stratified randomization directly introduces correlation among treatment assignment and observed data $(\boldsymbol{O}_{1},\dots,\boldsymbol{O}_{n})$.

Stratified rerandomization combines stratified randomization and rerandomization through the following steps. First, we generate $\widetilde{A}_{1},\dots,\widetilde{A}_{n}$ by stratified randomization based on randomization strata $\{S_{1},\dots,S_{n}\}$ and parameter $\pi\in(0,1)$. Second, we compute

where $\widetilde{N}_{1}=\sum_{i=1}^{n}\widetilde{A}_{i}$, $\widetilde{N}_{0}=n-\widetilde{N}_{1}$, and $\overline{\boldsymbol{X}_{s}^{r}}=(\sum_{i=1}^{n}I\{S_{i}=s\})^{-1}\sum_{i=1}^{n}\boldsymbol{X}_{i}^{r}I\{S_{i}=s\}$. In the last step, if $\widetilde{\boldsymbol{I}}^{\top}\{\widehat{Var}(\widetilde{\boldsymbol{I}})\}^{-1}\widetilde{\boldsymbol{I}}<t$ for a pre-specified balance threshold $t$, we set the final treatment assignment $(A_{1},\dots,A_{n})$ to be $(\widetilde{A}_{1},\dots,\widetilde{A}_{n})$; otherwise, we return to the first step to obtain the first randomization scheme that satisfies the balance criterion.

Compared to rerandomization defined in Section 2.2, stratified rerandomization involves two changes: the first step replaces simple randomization by stratified randomization, and the variance estimator $\widehat{Var}(\widetilde{\boldsymbol{I}})$ is updated to reflect the true variance of the imbalance statistic under stratification, which is smaller than under simple randomization. This design is asymptotically equivalent to the first stratified rerandomization criterion (overall Mahalanobis distance) of Wang et al., 2023c since we assume $\pi$ is constant across randomization strata—the most common setting in randomized experiments. Without loss of generality, we assume $S\notin\boldsymbol{X}^{r}$; otherwise, rerandomization is effectively controlling for imbalance on $\boldsymbol{X}^{r}\setminus S$. Under stratified rerandomization, we next develop the large-sample properties for M-estimators regarding the general estimand $\Delta^{*}$.

Theorem 3 parallels Theorem 1 in terms of convergence in probability, asymptotic linearity, and weak convergence results. The major difference lies in the constant factors of the asymptotic distribution. In particular, $\widetilde{V}$ substitutes $V$ in Equation (3) to account for the variance reduction from stratification (Wang et al., 2023b, ). Likewise, compared to Theorem 1, $\widetilde{R}^{2}$ is also updated by replacing marginal covariance with the expectation of conditional covariance. Despite these changes, the properties of the asymptotic distribution remain similar to those under rerandomization, and statistical inference considerations under rerandomization discussed in Section 3 still apply to stratified rerandomization. Finally, if we consider the unadjusted estimator, the difference estimand $\Delta^{*}=E[Y(1)-Y(0)]$, and fixed strata categories, Theorem 3 becomes the counterpart of Theorem 3 of Wang et al., 2023c under a sampling-based framework. Theorem 3, however, includes a wider class of M-estimators and addresses a more general class of estimands $\Delta^{*}$ (e.g. to allow for ratio estimands).

For Examples 1-4 in Section 4, their model-robustness or double-robustness property remains the same under stratified rerandomization. We thus extend Theorem 3 for stratified rerandomization given below.

Theorem 4 clarifies that adjusting for both $\boldsymbol{X}^{r}$ and $S$ in the working models offers asymptotic normality of the considered M-estimators under stratified rerandomization. In this case, the asymptotic distribution for each estimator is no different from that under simple randomization, and asymptotic results developed for those estimators under simple randomization can directly apply even though the actual assignment comes from a stratified rerandomization procedure. Among the three conditions, conditions (1) and (3) remain the same as Theorem 3, but condition (2) is modified to further include additional treatment-by-stratum interaction terms (as intuitively, stratum variable is part of the randomization procedure). Similar to Theorem 3, the threshold $t$ and weighting matrix $\widehat{\mathbf{H}}$ do not contribute to the asymptotic distribution, suggesting that the choice of these design parameters in general does not impact the asymptotic inference of the considered treatment effect estimators. Due to the appeal of applying normal approximation for inference, we maintain our recommendation that all rerandomization variables (including rerandomization variables and dummy stratum variables) be adjusted for in the analysis.

Under rerandomization, we study the efficient estimator that uses double machine learning (DML) for estimating nuisance functions with cross-fitting (Chernozhukov et al.,, 2018). Let $\widehat{\eta}_{a}(\boldsymbol{X};\mathcal{D}),\widehat{\kappa}_{a}(\boldsymbol{X};\mathcal{D})$ denote the pre-specified estimators for $E[Y(a)|\boldsymbol{X}],E[R(a)|\boldsymbol{X}]$ trained on data $\mathcal{D}$ and evaluated at $\boldsymbol{X}$. For cross-fitting, we randomly partition the index set $\{1,\dots,n\}$ into $K$ folds with approximately equal sizes. Specifically, let $\mathcal{I}_{k}$ denote the index set for the $k$-th fold, we have $\{1,\dots,n\}=\cup_{k=1}^{K}\mathcal{I}_{k}$ and $|\mathcal{I}_{k}|=K^{-1}n$. If $K$ does not divide $n$, we require $||\mathcal{I}_{k}|-K^{-1}n|\leq 1$ instead. Defining $\mathcal{O}_{a,-k}=\{(R_{i}Y_{i},R_{i},\boldsymbol{X}_{i}):i=1,\dots,n;i\notin\mathcal{I}_{k};A_{i}=a\}$ as the training data for fold $k$ within treatment group $a$, we specify the nuisance function estimators as $\widehat{\eta}_{a}(\boldsymbol{X}_{i})=\sum_{k=1}^{K}I\{i\in\mathcal{I}_{k}\}\widehat{\eta}_{a}(\boldsymbol{X}_{i};\mathcal{O}_{a,-k})$ and $\widehat{\kappa}_{a}(\boldsymbol{X}_{i})=\sum_{k=1}^{K}I\{i\in\mathcal{I}_{k}\}\widehat{\kappa}_{a}(\boldsymbol{X}_{i};\mathcal{O}_{a,-k})$. Next, the efficient estimator for $E[Y(a)]$ is constructed based on the efficient influence function as

and output $\widehat{\Delta}^{\textup{dml}}=f(\widehat{\mu}_{1}^{\textup{dml}},\widehat{\mu}_{0}^{\textup{dml}})$. For this estimator, we assume the following conditions.

Assumption 4(1) requires that the nuisance function estimators converge to their target in $L_{2}$-norm with a rate faster than $n^{1/4}$, if the training data consist of independent and identically distributed data. This rate can be achieved by existing machine learning methods, including random forests (Wager and Athey,, 2018), deep neural network (Farrell et al.,, 2021), and highly-adpative lasso (Benkeser and Van Der Laan,, 2016), under mild regularity conditions. Of note, for the outcome model, the training data $\mathcal{D}_{n}$ are used to learn $E[R(a)Y(a)|R(a),\boldsymbol{X}]$, which yields $E[Y(a)|R(a)=1,\boldsymbol{X}]$ and thus $E[Y(a)|\boldsymbol{X}]$ under Assumption 2. In Assumption 4(2), we assume uniform boundedness on several components of $\mathcal{P}^{W}$ and the estimators so that the remainder terms can be appropriately controlled. Overall, Assumption 4 resembles the standard assumptions for efficient inference made in Chernozhukov et al., (2018), and importantly, no extra condition is specifically assumed for addressing rerandomization. Theorem 5 below provides the asymptotic results for $\widehat{\Delta}^{\textup{dml}}$.

Theorem 5 justifies the validity of efficient estimation of the average treatment effect estimand under rerandomization. Additionally, the overall structure of its asymptotic expansion resembles that for M-estimators in Theorem 1, with the only difference being that the influence function is given in Equation (5), which is the efficient influence function under simple randomization. If covariate adjustment includes the rerandomization variables, then $\widehat{\Delta}^{\textup{dml}}$ is asymptotically normal with asymptotic variance $V$—the usual efficiency lower bound given by the variance of $EIF(\boldsymbol{O}_{i})$ under simple randomization. On the contrary, if $\boldsymbol{X}$ does not include all rerandomization variables, then $V$ does not fully account for the precision gain from rerandomization, which leads to a positive $R^{2}$ and non-normal asymptotic distribution.

Comparing efficient estimators with parametric M-estimators, efficient estimators have the advantage of optimal efficiency, while machine learners often require a large sample size to avoid over-fitting. We recommend using efficient estimators if the sample size per arm is no less than 200, and caution against data-adaptive fitting with small sample sizes, e.g., $n\leq 50$.