On flexible covariate adjustment under covariate-constrained randomization

Abstract

Covariate-constrained randomization is an effective treatment allocation procedure for controlling imbalance across multiple baseline covariates in randomized trials. Motivated by the GroupPMPlus cluster randomized trial, we introduce the asymptotic theory for a broad class of estimators, known as M-estimators, under covariate-constrained randomization. Here, M-estimators refer to estimators obtained by optimizing an objective function, such as a log-likelihood function, and include commonly used methods such as analysis of covariance and linear mixed models. We show that M-estimators remain consistent in this setting but can exhibit non-Gaussian asymptotic distributions depending on the specification. Using examples of common M-estimators, we delineate conditions under which covariate-constrained randomization can be safely ignored in statistical analysis. Our results extend to stratified covariate-constrained randomization and semiparametric efficient estimators based on data-adaptive machine learning methods. We illustrate these theoretical findings using the GroupPMPlus study to evaluate the causal effect of a psychological treatment on mental health outcomes following a disaster.

Introduction

Covariate-constrained randomization refers to a restricted randomization procedure in which treatment assignments are repeatedly generated until a predefined balance criterion on baseline covariates is met. Compared with simple randomization or stratified randomization,¹ this approach can simultaneously balance several covariates of different types and allows continuous control over the randomization space through a design parameter. Compared with minimization, which balances multiple categorical covariates, covariate-constrained randomization can accommodate both continuous and categorical covariates, offering additional flexibility. These advantages have contributed to its increasing use in both economics,² where it is often referred to as rerandomization, and biomedical research,^3,4 particularly in the context of cluster randomized trials to improve baseline balance.^5–7

Our work is motivated by the GroupPMPlus trial,⁸ a cluster randomized study in Nepal designed to improve the mental health of individuals affected by humanitarian crises such as pandemics, armed conflict, and environmental disasters. The intervention, Group Problem Management Plus, consisted of five weekly sessions and was compared with standard care. A total of 72 wards, the smallest administrative units in Nepal, were enrolled, comprising 609 participants. The intervention and standard care were both implemented at the ward level, and wards were equally randomized to receive intervention or standard care. The trial employed stratified covariate-constrained randomization, with stratification based on gender (uniform within wards) and covariate-constrained randomization based on three binary cluster-level covariates: access to mental health services (high or low), disaster risk (high or low), and rural or urban geographical location. The primary outcome was the GHQ-12 score, a continuous measure of psychological distress, assessed 3 months after treatment initiation.

Several methodological questions arose as we approached the analysis of the GroupPMPlus trial. First, although mixed-effects models are routinely used to adjust for baseline covariates while accounting for the within-cluster correlations,⁹ their asymptotic validity under covariate-constrained randomization has not been fully established; see Wang et al.¹⁰ for the asymptotic theory for linear mixed model analysis of cluster randomized trials under simple randomization. Second, incorporating covariate-constrained randomization into the statistical analysis typically requires knowledge of the balancing threshold, which is not directly reported for this study from their primary publication. The implications of the balancing threshold on the asymptotic property of the treatment effect estimator remain elusive. Third, while data-adaptive machine learning methods have shown potential for maximizing efficiency gains in randomized trials,^11,12 it is unclear whether this property carries to covariate-constrained randomization. Addressing these several methodological questions can not only strengthen the statistical analysis of the GroupPMPlus trial but also provide a sound theoretical foundation for future individually randomized and cluster randomized studies that adopt covariate-constrained randomization.

In this paper, we expand the asymptotic theory of M-estimators and efficient estimators—in which nuisance functions are estimated via data-adaptive machine learning algorithms—to the settings of covariate-constrained randomization and its stratified extension.¹³ These large-sample results generalize earlier work that focuses on linearly adjusted estimators^14–16 to a much wider class of treatment effect estimators commonly used in analyzing clinical trials.

Methods

Definitions and assumptions

As a general setup, we consider a randomized clinical trial with $n$ individuals. For each individual $i(i=1,\dots ,n)$, we observe a real-valued outcome $Y_i$, a treatment allocation variable $A_i\left(A_i=1\right)$ for treatment and 0 for control), and a vector of baseline covariates $X_i$. We define $Y_i(a)$ as the potential outcome if individual $i$ were assigned to treatment $a,a\in \{0,1\}$ and assume causal consistency such that $Y_i=Y_i\left(A_i\right)$. Furthermore, we assume that the vectors $\left(Y_i(1),Y_i(0),X_i\right)$ for $i=1,\dots ,n$ are independent and identically distributed. Under the super-population causal inference framework, the target estimand is the average treatment effect, defined as ${\mathrm{\Delta }}^{\mathrm{*}}=E[Y(1)-Y(0)]$, where the expectation is taken with respect to a notional super-population that generates the full data vector.¹⁷

Covariate-constrained randomization controls for imbalance on a pre-specified set of baseline covariates, which we denote as $X_i^r$ with $X_i^r$ being a subset of $X_i^r$, and involves the following three steps. First, we independently generate $A_1^{\mathrm{*}},\dots ,A_n^{\mathrm{*}}$ from a Bernoulli distribution with $P(A_i^{\mathrm{*}}=1)=\pi \in (0,1)$ as in simple randomization. Next, we compute the imbalance statistic and its variance estimator as

$$I=\frac{1}{N_1}\sum _{i=1}^n{A}_i^{\mathrm{*}}{X}_i^r-\frac{1}{N_0}\sum _{i=1}^n\left(1-A_i^{\mathrm{*}}\right)X_i^r,$$

$$\widehat{\mathrm{V}\mathrm{a}\mathrm{r}}(I)=\frac{1}{N_1{N}_0}\sum _{i=1}^n\left(X_i^r-\overline{X^r}\right){\left(X_i^r-\overline{X^r}\right)}^{\mathrm{\top }}$$

where $N_1=\sum _{i=1}^n{A}_i^{\mathrm{*}},N_0=n-N_1^{\mathrm{*}}$ and $\overline{X^r}=n^{-1}\sum _{i=1}^n{X}_i^r$. Finally, given a pre-specified balance threshold $t>0$, we check a Mahalanobis distance type condition of whether $I^{\mathrm{\top }}\{\widehat{\mathrm{V}\mathrm{a}\mathrm{r}}(I){\}}^{-1}I<t$. If true, the final treatment assignment ($A_1,\dots ,A_n$) is set to be $(A_1^{\mathrm{*}},\dots ,A_n^{\mathrm{*}})$; otherwise, we return to the first step and regenerate the treatment assignment until the condition is met.¹⁸ As a result, the output will be a set of treatment assignments satisfying the balancing criterion.

Under covariate-constrained randomization, a smaller $t$ corresponds to a higher rejection rate for the generated allocation and a stronger control over chance imbalance. When covariate-constrained randomization is combined with stratified randomization (e.g. the stratified covariate-constrained randomization), the first step (generating treatment allocation under simple randomization) will be replaced by stratified randomization among pre-specified discrete baseline strata; the second and third steps remain unchanged. Thus, the strata variables will typically be perfectly balanced as covariate-constrained randomization is implemented within each stratum.

Asymptotic results for M-estimators

M-estimators¹⁹ broadly refer to the class of estimators obtained by solving estimating equations (or equivalently optimizing an objective function). Denoting $O_i=\left(Y_i,A_i,X_i\right)$, we first specify an estimation function $\psi \left(O_i\in \theta \right)$, where $\theta$ is a vector of parameters, and then estimate $\theta$ by solving $\sum _{i=1}^n\psi \left(O_i;\theta \right)=0$. The resulting estimator $\hat{\theta }$ is called an M-estimator for $\theta$. In our setting, the average treatment effect $\mathrm{\Delta }$ will be a parameter in $\theta$ (or as a function of $\theta$) such that this procedure can output an M-estimator $\widehat{\mathrm{\Delta }}$ that converges to a probability limit ${\mathrm{\Delta }}^{\mathrm{*}}$.

This framework encompasses many estimation methods. For instance, maximum likelihood estimators are M-estimators where $\psi \left(O_i;\theta \right)$ is the score function. To contextualize the development, we describe two additional examples of M-estimators that will be applied to the GroupPMPlus trial. In the context of cluster randomized trials with simple or stratified randomization, these two M-estimators are previously discussed by Wang et al.¹⁰ and are shown to provide model-robust estimators for the average treatment effect under arbitrary misspecification.

Example 1 (The ANCOVA estimator). Consider fitting a working model $E[Y\mid A,X]={\beta }_0+{\beta }_{A}A+{\beta }_X^{\mathrm{\top }}X$ with ordinary least squares to obtain estimators $({\widehat{\beta }}_0,{\widehat{\beta }}_A,{\widehat{\beta }}_X)$ for parameters (${\beta }_0,{\beta }_A,{\beta }_X$). The analysis of covariance (ANCOVA) estimator for the average treatment effect is $\widehat{\Delta }={\widehat{\beta }}_A$. In the context of the cluster randomization, because the ordinary least squares approach is used for point estimation, $\widehat{\mathrm{\Delta }}$ can also be considered as the generalized estimating equations estimator under working independence.

Example 2 (The Mixed-ANCOVA estimator). In cluster randomized trials, the outcomes for each cluster become an $N_i$-dimensional vector $Y_i=\left(Y_{i1},\dots ,Y_{i{N}_i}\right)$, where $N_i$ is the size of cluster $i$. In addition, the covariate $X_i$ is a collection of individual covariates $\left\{X_{1i},\dots ,X_{i{N}_i}\right\}$. Then, covariate-constrained randomization is based on $X_i^r$ being summary function of $\left(X_{i1}^r,\dots ,X_{i{N}_i}^r\right)$, such as cluster averages. In this context, we fit the linear mixed model $Y_{\mathrm{i}\mathrm{j}}={\beta }_0+{\beta }_A{A}_i+{\beta }_X{X}_{\mathrm{i}\mathrm{j}}+{\delta }_i+{\epsilon }_{\mathrm{i}\mathrm{j}}$, where ${\delta }_i~N\left(0,{\tau }^2\right)$ is the random intercept and ${\epsilon }_{\mathrm{i}\mathrm{j}}~N\left(0,{\sigma }^2\right)$ is the independent error. Under maximum likelihood estimation, the mixed-model ANCOVA estimator for the average treatment effect is $\widehat{\mathrm{\Delta }}={\widehat{\beta }}_A.$

Under mild regularity conditions, M-estimators have the following properties. First, an M-estimator has the same consistency (asymptotically unbiasedness) under simple and covariate-constrained randomization. Second, its asymptotic distribution is typically non-Gaussian, but a weighted average of a Gaussian distribution and a truncated normal distribution. Third, it has no larger asymptotic variance under covariate-constrained randomization than under simple randomization, with more variance reduction given a smaller $t$ (tighter balance threshold or equivalently stronger balance control) or higher prognostic values in $X^r$. These statistical properties are illustrated Figure 1. When stratified randomization is combined with covariate-constrained randomization, the same results hold with stratified randomization replacing simple randomization (the original version of covariate-constrained randomization). A more rigorous presentation of these technical results can be found in Wang and Li.¹³

Figure 1.

Illustration of the asymptotic distribution under different randomization schemes and parameter settings. The blue solid curve shows the distribution under simple randomization (standard normal). The red dashed curve corresponds to covariate-constrained randomization with $t=1$ and a moderately prognostic covariate $X_i^r$ (with $R^2=0.5$ in the linear regression of $Y_i$ on $X_i^r$). The green dotted curve corresponds to covariate-constrained randomization with $t=1$ and a strongly prognostic covariate $X_i^r$ (with $R^2=0.75$ in the linear regression of $Y_i$ on $X_i^r$).

According to these theoretical results, statistical analysis ignoring the covariate-constrained randomization, for example, using normal approximation for hypothesis testing, will generally be valid but conservative.¹³ This is because covariate-constrained randomization brings additional precision gain by balancing baseline covariates. To avoid such power loss, one can make statistical inferences with the exact asymptotic non-Gaussian distribution. Alternatively, appropriate covariate adjustment via outcome modeling is also an effective approach to restore the nominal rejection rate of the hypothesis testing procedure. Specifically, in Examples 1 and 2, since we have included $X^r$ in the regression model, under equal randomization (i.e. $\pi =0.5$), these estimators will have the same asymptotic distribution under simple or covariate-constrained randomization.¹³ Therefore, valid yet non-conservative statistical inference for these estimators can be performed as if simple randomization were used in the design stage. This theoretical finding provides a foundation to support previously observed patterns in simulation studies, see, for example, Li et al.,^5,6 when analyzing continuous and binary outcomes in cluster randomized trials. Furthermore, this result holds regardless of the choice of $t$, which obviates the need to know $t$ for valid inference in large samples. For other estimators, such as ANCOVA2 (ANCOVA with treatment–covariate interaction terms²⁰), G-computation with a working logistic regression model,²¹ and doubly robust estimators to address missing outcomes,²² adjusting for in the outcome regression can similarly capture the precision gains induced by covariate-constrained randomization. Formal justification is provided in Wang et al.,¹³ which extends results established under stratified permuted block randomization in Ye et al.²³ and Wang et al.²⁴ to covariate-constrained randomization.

Extension to efficient estimation with debiased machine learning

In the analysis of randomized trials, statistical precision through covariate adjustment can be maximized using the debiased machine learning (DML) approach with cross-fitting.¹¹ For example, Wang et al.¹² have developed the DML approach for the analysis of cluster randomized trials with the additional complication of within-cluster subsampling. Compared to traditional regression methods, DML has been proven to achieve semiparametric efficiency in randomized trials with simple randomization. The procedure is as follows. First, we randomly partition the data into $K$ folds of approximately equal size. Next, for each fold $k=1,\dots ,K$, we use the other $K-1$ fold to train a machine learning algorithm for $E[Y\mid A,X]$, which we denote as ${\hat{\eta }}^{(k)}(A,X)$. Then, we compute ${\hat{\eta }}^{(k)}(A,X)$ for individuals in the $k$ the fold. Going through all $K$ folds, we can obtain the machine learning prediction for each individual $i=1,\dots ,n$, which we denote as $\hat{\eta }\left(A_i,X_i\right)$. Finally, the DML estimator for the average treatment effect is constructed based on the efficient influence function as

$${\widehat{\mathrm{\Delta }}}^{\mathrm{d}\mathrm{m}\mathrm{l}}=\frac{1}{n}\sum _{i=1}^n\left[\frac{A_i-\pi }{\pi (1-\pi )}\left\{Y_i-\hat{\eta }\left(A_i,X_i\right)\right\}+\hat{\eta }\left(1,X_i\right)-\hat{\eta }\left(0,X_i\right)\right].$$ (1)

For deriving the asymptotic results of this estimator, we assume ${\hat{\eta }}^{(b)}(A,X)$ is consistent in $L_2$ norm, which can be achieved by many existing machine learning methods, including random forests,²⁵ deep neural network,²⁶ and highly adaptive lasso.²⁷ When stratified covariate-constrained randomization is used, the above sample splitting needs to be stratified based on randomization strata to preserve treatment balance; the detailed procedure is described in Rafi²⁸ and Wang and Li.¹³

Under covariate-constrained randomization or stratified covariate-constrained randomization, we can show that ${\widehat{\mathrm{\Delta }}}^{\mathrm{d}\mathrm{m}\mathrm{l}}$ is consistent, asymptotically normal, and as efficient as under simple randomization. This property holds regardless of the values of $t$ or choices of $X^r$, provided that $X^r$ is included in the machine learning algorithms as predictors. Thus, inference can proceed as if the trial used simple randomization, and statistical precision remains optimal with regard to the semiparametric variance lower bound. The detailed presentation of this result can be found in Section 6 of Wang et al.¹³ An important implication of this result is that it justifies the use of flexible, data-adaptive covariate adjustment methods under (stratified) covariate-constrained randomization, and offers a recipe for valid statistical inference as long as $E\left[Y\right|A,X]$ is consistently estimated via machine learning methods.

Illustrate analysis of the GroupPMPlus trial

We demonstrate the theoretical findings by reanalyzing the GroupPMPlus cluster randomized trial. Here, we implemented the unadjusted, ANCOVA, and DML estimators at the cluster level. Specifically, we first take cluster averages of the outcomes and covariates and then conduct the analysis using these cluster-level data. The ANCOVA estimator is defined as in Example 1, and the unadjusted estimator is defined as ANCOVA without covariates. The DML estimator is defined in equation (1), and we use an ensemble learner of generalized linear models, regression trees, and neural networks to construct $\hat{\eta }\left(A_i,X_i\right)$. We also fit the mixed-ANCOVA estimator of Example 2 using individual-level data. Baseline covariates included the baseline GHQ-12 score and all variables used in stratified covariate-constrained randomization.

Since we are unable to obtain information about the balancing threshold $t$ used in the randomization procedure from the published report, we carry out our analysis assuming simple randomization; that is, we compute the standard errors of each estimator and their confidence intervals under normal approximation. This choice will lead to conservative confidence intervals for an unadjusted analysis but has no impact on the covariate-adjusted estimators, as supported by our theory. That is, as long as the randomization variables are controlled by the adjusted estimators, the covariate-constrained randomization procedure should be safely ignorable asymptotically, and statistical inference procedures established under simple randomization can be used without modification.

The results are summarized in Figure 2. While all estimators have similar point estimates, their standard error estimates differ. Since we are unaware of the design parameter $t$, the unadjusted estimator has a conservative standard error estimate of 0.78, leading to failure to reject the null at the 5% level. In contrast, all three covariate-adjusted estimators have fully accounted for the precision gain from stratified covariate-constrained randomization, thus leading to valid estimation of the variance. This is achieved even without knowing the exact randomization parameters, which further demonstrates the benefit of adjusting for randomization variables. Among the three covariate-adjusted estimators, the cluster-level ANCOVA estimator has the highest precision, while the DML estimator leads to the least variance reduction in this particular data analysis. This may be because either the sample size is relatively limited for machine learning methods to fully realize their asymptotic efficiency gain or the true data-generating distribution is almost linear in the baseline covariates (hence the cluster-level ANCOVA reasonably well approximated the data-generating process).

Figure 2.

Summary of data application results. The left panel presents point estimates with 95% confidence intervals, where negative values indicate a treatment benefit. The right panel shows relative efficiency compared with the unadjusted estimator, calculated as the variance of the unadjusted estimator divided by the variance of each alternative estimator. Larger values correspond to greater precision.

Conclusion and recommendation

Covariate adjustment in randomized clinical trials can be incorporated at both the design stage, through approaches such as covariate-constrained randomization, and at the analysis stage, via outcome modeling and other established covariate-adjustment techniques. Each stage offers distinct but complementary opportunities to improve statistical precision while preserving validity, as mentioned in recent regulatory guidance.²⁹ Our theoretical results demonstrate that, similar to conventional covariate-adaptive designs such as stratified randomization,²⁴ covariate-constrained randomization never increases asymptotic variances of the M-estimator for $\mathrm{\Delta }$ and, under appropriate modeling, can yield substantial precision gains.

From a methodological standpoint, our asymptotic theory for M-estimators under covariate-constrained randomization clarifies when and how the randomization scheme affects the large-sample distributional properties of treatment effect estimators under the potential outcomes framework. While consistency is retained under a broad class of outcome models due to randomization, the asymptotic distribution may deviate from normality, particularly when insufficient randomization variables are adjusted for in the analysis. This has implications for statistical inference under covariate-constrained randomization; that is, direct reliance on standard errors that ignore the randomization constraints can lead to mis-calibrated tests and confidence intervals that are typically conservative. This has been previously shown in the simulation studies for cluster randomized trials,^5,6 where the unadjusted $F$-test under a linear mixed model can grow substantially conservative when only a few covariates are used in the covariate-constrained randomization procedure.

From a practical standpoint, linear adjustment for the randomization variables in ANCOVA offers substantial efficiency gains, is simple to carry out, and has strong theoretical support^30,31 in many common trial settings. For example, under equal allocation in two-arm trials, the ANCOVA estimator has been shown to asymptotically equivalent to the semiparametric efficient estimator (when the true outcome surface is linear),^32,33 and this can fully capture the precision gain from covariate-constrained randomization as long as the randomization variables are sufficiently adjusted. While ANCOVA is attractive for its simplicity and robustness, it may not fully capture the nuanced data-generating process (especially under unequal allocation or when the true outcome surface is nonlinear), making more flexible, data-adaptive covariate adjustment methods such as DML appealing. These methods can approximate complex relationships and, in large samples, achieve semiparametric efficiency under weak assumptions. However, there is a trade-off between model complexity and theoretical optimality. That is, DML requires a sufficiently large number of independent units to accurately estimate complex nuisance functions and realize its asymptotic benefits. In our motivating example, with 72 clusters as independent units, the advantage of DML over ANCOVA may be limited. In contrast, in individually randomized trials or cluster randomized trials with over 100 clusters, DML may have greater potential to outperform ANCOVA. A similar discussion of this balance between complexity and efficiency under simple randomization in cluster randomized trials can be found in Wang et al.,²⁴ and the same considerations apply in the context of covariate-constrained randomization designs.

For implementation, covariate-constrained randomization encompasses many variations, differing in balance metrics, weighting of covariates, and acceptance thresholds, each of which may influence both balance properties and downstream analytic requirements. Our theoretical results provide formal justification for these variants and guidance for aligning the analytic model with the specific form of constrained randomization used. In practice, we recommend that trial statisticians document the chosen balancing algorithm in sufficient detail to permit replication and to inform the choice of analytic model. Specifically, our results reinforce the recommendation of aligning the design and analysis in clinical trials (e.g. those that were discussed in Li et al.^5,6 based on their extensive simulation findings). That is, one is recommended to ensure that the analysis model incorporates randomization variables used in covariate-constrained randomization, to the extent possible, such that inference procedures developed under simple randomization can be directly applicable.

Covariate-constrained randomization is primarily used in settings where trial participants are predetermined, such as cluster randomized trials where clusters are often recruited prior to randomization or individual randomized trials involving a closed cohort identified prior to randomization. In contrast, many clinical trials enroll participants sequentially, so covariate information is only observed at the time of enrollment. To address this setting, sequential rerandomization³⁴ has been developed, in which participants are enrolled in batches and covariate-constrained randomization is applied within each batch. A formal theoretical investigation of this design is left for future research.

Finally, our work can be extended to several future directions involving high-dimensional covariate adjustment, adaptive treatment allocation, or complex missing data mechanisms. Such work will further bridge the gap between rigorous statistical theory and the increasingly rich data structures encountered in modern clinical trials. By integrating careful design choices with robust and efficient covariate-adjusted analyses, investigators can maximize both the precision and credibility of their trial findings.