A unified framework for covariate adjustment under stratified randomization

We start by considering the oracle case where the space $\Omega$ spanned by functions of $X$ with finite variance is given. Conditional on $B_{i}=k$, let $h_{[k]}(X,1)$ and $h_{[k]}(X,0)$ denote the projections of $Y(1)$ and $Y(0)$ onto $\Omega$, respectively. Following the idea of the regression-adjusted estimator, our proposed oracle estimator is:

Let $r_{i}(a)=Y_{i}(a)-[(1-\pi)h_{[k]}(X_{i},1)+\pi h_{[k]}(X_{i},0)],i\in[k]$, we establish the following proposition for the oracle estimator. The detailed proof is given in Section 2 of the Supplementary Materials.

Plugging in the estimates of $h_{[k]}(X,a)$, our empirical treatment effect estimator adjusting for baseline covariates is defined as (Liu et al.,, 2023):

where $\hat{h}_{[k]}(X_{i},a),\ a=0,1$, are projection functions estimated by different methods.

The discussion presented by Tsiatis et al., (2008) implied that estimates of $h_{[k]}(X_{i},a),\ a=0,1$ are desirable if they can estimate $h_{[k]}(X_{i},a)$ well, but no formal conditions were given on the estimation error. To study the asymptotic properties of $\hat{\tau}_{\text{emp}}$ under stratified randomization, the following conditions were outlined by Liu et al., (2023).

The first equation in Assumption 4, i.e. Equation (2), implies that the difference between the oracle estimator and the empirical estimator is negligible at the rate of $o_{P}(n^{-1/2})$, the assumption of the similar form has also been used for the regression-adjusted estimation of quantile treatment effects (Assumption (i) in Jiang et al., (2023)). Equation (3) allows control of the estimation error of the projection function and serves as a guarantee of consistent variance estimation. Based on this assumption, we can draw valid inference for the average treatment effect. According to the low or high dimensionality of the covariates, a wide range of methods can be applied to estimate $h_{[k]}(X_{i},a),\ a=0,1$. If we consider the space spanned by arbitrary measurable functions of $X$ with finite variance, then nonparametric methods can be applied. Additionally, if we have prior information on the format of $h_{[k]}(X_{i},a)$, we can adopt parametric regressions. For linear adjustments, these conditions have been verified for lasso and linear regressions. For nonlinear adjustments, Wang et al., (2023) considered the M-estimation of the parameter of interest and proved relevant asymptotic properties. However, theoretical properties of covariate adjusted estimators using more general nonparametric or machine learning methods are not clear and yet to be explored, see discussion in Section 8. Meanwhile, Assumption 4 is a high-level assumption that, regardless of whether the setting is low-dimensional or not, as long as the adjusting method satisfies this assumption, we can obtain a consistent treatment effect estimator and make valid inference. One of the main contributions of this paper is the verification that this assumption is satisfied for local linear kernel in low-dimensional cases. In high-dimensional cases, this assumption may be hard to meet, and we thus propose a sample splitting algorithm for machine learning methods, which is another contribution of this paper. In brief, this algorithm separates the samples used for the estimation of $\hat{h}_{[k]}(\cdot,a)$ and the evaluation sample $X_{i}$ in $\hat{h}_{[k]}(X_{i},a)$, hence simplifies the required assumption. Detailed information about the sample splitting process, simplified assumption can be found in Section 6.3.

Denote $\hat{r}_{i}(a)=Y_{i}(a)-[(1-\pi)\hat{h}_{[k]}(X_{i},1)+\pi\hat{h}_{[k]}(X_{i},0)]$ as the estimated value of $r_{i}(a)$. The asymptotic normality and a consistent asymptotic variance estimator of $\hat{\tau}_{\text{emp}}$ is provided as in the following theorem by Liu et al., (2023), thus justifying the Wald-type inference of the average treatment effect.

There are different choices of the subtractor in transformed outcomes, but not all of them improve the efficiency of the treatment effect estimator. How to achieve the optimal efficiency of the treatment effect estimator has been established as in the following theorem by Liu et al., (2023), and we give a more detailed explanation here.

To conduct valid and efficient inference, we need to minimize the distance between the estimation function and true function, i.e., control both the estimation error and approximation error. Assumption 4 imposes detailed requirements on the convergence rate of the estimating methods, which controls the estimation error to obtain valid estimates. To optimize efficiency, we need to further eliminate the approximation error, which is determined by the function space spanned by estimating methods. If the function space contains the true function (i.e., $E\{Y(a)|X,B=k\},\ a=0,1$ in Theorem 2), then the corresponding treatment effect estimator has optimal efficiency. However, it is often the case that we cannot know exactly which function space contains the true function, especially in high-dimensional cases. Therefore, we can only consider the rationality of function spaces for different estimating methods, such as the linear function space for linear regressions, the sieve space for artificial neural networks (ANNs), and the spline space for smoothing methods, in the context of current data, and compare their practical performances. In the following sections, we introduce commonly used estimating methods and demonstrate their empirical efficiencies through numerical simulations.

Although Tsiatis et al., (2008) suggested a general strategy for estimating the projection function with the flexible use of modeling methods for covariate-adjusted estimators under simple randomization, they did not give theoretical details. Several nonparametric statistical tools including kernel, spline, and orthogonal series have been proposed and widely used for estimating $h_{[k]}(X_{i},a)$. Among these tools, local linear kernel has better asymptotic behavior (Fan,, 1992). In this paper, we give a formal justification for the use of the local linear kernel smoother of $h_{[k]}(X_{i},a)$ under stratified randomization. Consider the problem that for $i\in[k]$,

where $H$ is a $d\times d$ symmetric positive definite matrix depending on $n$, $K_{H}(u)=|H|^{-1/2}K(H^{-1/2}u)$ with $K$ being a $d$-dimensional kernel such that $\int K(u)d_{u}=1$, and $|\cdot|$ denotes the determinant of a matrix. $\mathds{1}_{A_{j}=a}$ is an indicator function that equals $1$ if $A_{j}=a$ and $0$ otherwise. $H^{1/2}$ is called the bandwidth matrix. Then, $\hat{h}_{[k]}(X_{i},a)=\hat{\alpha}$ is the local linear kernel smoother of $h_{[k]}(X_{i},a)$. We make the following general assumptions about local linear kernel.

Then based on Theorem 1, we establish the following theorem.

As a nonparametric method competing with kernel methods, spline smoothing has commonly been investigated in the literature. Unlike the local linear weighted regression idea of the kernel, spline smoothing considers adding penalties when minimizing the sum of squares, to avoid over-fitting. One of the most frequently adopted penalties is the integral over the second-order derivative of the estimated function, which corresponds to the cubic spline. Additionally, despite the different forms and concepts of kernel and spline, Silverman, (1984) showed the asymptotic equivalence of spline smoothing and a kernel method with a bandwidth depending on the local density of design points. Our simulation results suggest the similar performance of spline smoothing and kernel methods under certain scenarios and the superiority of local linear kernel in certain situations. Thus, we do not explore the theoretical justification of spline smoothing-adjusted estimators in this paper.

In the presence of high-dimensional covariates, linear regression is prone to problems of multi-collinearity or overfitting. In practice, penalized regressions are usually adopted to solve these problems. When the exact or weak sparsity assumption of the population projection coefficients is reasonable, lasso can be adopted, and the satisfiability of the lasso-adjusted estimators for Assumption 4 was proved by Liu et al., (2023). Fundamental work on the concentration inequality and restricted eigenvalue under stratified randomization in the work of Liu et al., (2023) can be used for other penalized regressions, such as adaptive-lasso (Zou,, 2006; Huang et al.,, 2008), Ridge regression (Hoerl and Kennard,, 1970), and elastic net regression (Zou and Hastie,, 2005). Moreover, although the debiased-lasso (Zhang and Zhang,, 2014) usually has a performance similar to that of lasso, it has a higher computational cost, and it is thus not recommended here. When there are many weak predictors, Ridge regression can be chosen.

With the abundance of available data and the increase of model complexity, we need the help of computers to perform data-intensive and complex operations in addition to model training and estimation, i.e., machine learning. Nowadays, there are many machine learning methods available to researchers, among which the two types of method commonly used in medical research are tree-based methods and neural networks (Garg and Mago,, 2021). Both types of method can handle data with high dimensionality and complex interactions through stepwise deconstruction. However, as has been pointed out, no single method can solve problems in a one-size-fits-all manner (e.g., Peel,, 2010; Pedersen et al.,, 2020; De Cristofaro,, 2021). In practice, we should consider the type of problem to be solved, the structure of covariates, and the parameters of interest in selecting the most appropriate method. The good approximability of machine learning methods effectively reduces the approximation error.

Tree-based methods hierarchically partition the covariate space, recursively dividing the entire space into small regions. They can handle categorical and ordinal covariates in a simple way and automatically select covariates in steps and reduce model complexity. However, the single-tree structure tends to have insufficient prediction accuracy. The prediction performance has thus been improved using ensemble trees, commonly through bagging and boosting. Random forest, proposed by Breiman, (2001), is a substantial modification of bagging (Breiman,, 1996) in ensemble learning that constructs de-correlated trees on different bootstrap samples of the data and averages the results. Random forest is widely used in randomized experiments. For example, Wu and Gagnon-Bartsch, (2018) proposed a “leave-one-out potential outcomes” estimator by imputing potential outcomes using random forest, and Wager and Athey, (2018) estimated and inferred heterogeneous treatment effects using random forest. Boosting is a sequential process that continuously trains weak classifiers or regressors and adjusts the weights of samples and classifiers or regressors in each iteration, to reduce the prediction error. The tree-based methods that are widely used in practice in combination with boosting are the gradient boosting decision tree (Friedman,, 2002), Adaboost (Freund and Schapire,, 1997) and XGboost (Chen and Guestrin,, 2016).

Benefitting from the rapid development of computer technology, ANNs are appealing machine learning methods that can approach a wide variety of (nonlinear) functions, especially in the case of high-dimensional data (e.g., White,, 1992; Yarotsky,, 2018; Schmidt-Hieber,, 2020). The willingness to use ANNs has been demonstrated in different areas of research, including pattern recognition, decision-making, and pharmaceutical research. In recent years, ANNs have also been introduced to improve the efficiency of estimating treatment effects. Farrell et al., (2021) used ANNs to estimate nuisance functions in the estimation procedure of treatment effects, based on efficient influence functions. Chen et al., (2024) considered a more general framework of treatment effects with propensity score functions estimated using ANNs.

Machine learning methods can well estimate projection functions and make better predictions than traditional methods, but the fitted functions may have plenty of parameters and be complex in form. Additionally, substantive biases induced by machine learning methods are inevitable, and the restrictions on the estimation error in Assumption 4 are thus difficult to verify and may not be satisfied for various machine learning methods. In this case, we can use sample splitting for the separate estimation of the projection function and treatment effect, which relaxes the consistency conditions and thus makes most machine learning methods feasible (Chernozhukov et al.,, 2018).

Sample splitting is a common technique in statistics. This technique usually involves dividing data into two parts, one part for the inference of parameters or functions of interest and the other part for validation or estimating nuisance parameters (Picard and Berk,, 1990). Through sample splitting, we gain accuracy and robustness of the inference at the cost of prediction efficiency owing to the reduction of the sample size (Rinaldo et al.,, 2019). To regain full efficiency, Chernozhukov et al., (2018) proposed a “cross-fitting” process, which includes dividing data evenly into $M$ parts, estimating the parameter of interest and its variance on each inference-estimation data pair, and averaging them to obtain the final estimator. As the sample size increases, the estimators obtained from inference-estimation data pairs become asymptotically independent, and their variances can thus be aggregated for inference. The adoption of additional orthogonalization to reduce the regularization bias is known as the double/debiased machine learning method, which was proposed by Chernozhukov et al., (2018). This method has gained prevalence in research on high-dimensional statistical inference (e.g., Kallus et al.,, 2019; Bodory et al.,, 2022).

Moreover, the asymptotic properties of the methods described in Section 6.2 and the double/debiased machine learning method are deduced under the assumption that the treatment assignments and thus the outcomes are independent. However, under stratified randomization, there may be correlations between outcomes and between treatment assignment indicators within each stratum because of the treatment assignment procedure, which may lead to the invalidity of the aforementioned statistical inference. This motivates us to combine the form of the variance estimator in Theorem 1 with the sample splitting technique, where we expect to obtain a valid statistical inference using machine learning methods under stratified randomization. The proposed algorithm is described in Algorithm 1.

Indeed, this algorithm incorporates the sample splitting technique into the general framework of inference under stratified randomization. Using this algorithm, Assumption 4 on the projection function can reduce to the following second moment convergence assumption.

The above assumption can be satisfied by many machine learning methods under certain assumptions (Chernozhukov et al.,, 2018), and we establish the following theorem.

From the subsequent simulation results (Table 4), we can see that all the treatment effect estimators obtained using Algorithm 1 enjoy unbiasedness and validity.

Model 1:

with $\mu_{0}=1,\ \mu_{1}=4,\ \beta_{0}^{\textnormal{T}}=(75,35,125,80)$, and $\beta_{1}^{\textnormal{T}}=(100,80,60,40)$. $X_{i}$ is a four-dimensional vector, $X_{i1}\sim\textup{Beta}(3,4)$, $X_{i2}\sim\textup{Unif}[-2,2]$, $X_{i3}$ takes values in $\{-1,1\}$ with equal probability, $X_{i4}$ takes values in $\{3,5\}$ with probability $0.6,0.4$, and they are independent of each other. The variable used for randomization is an additional variable that takes a value in $\{1,2,3,4\}$ with probability $0.2,0.3,0.3,0.2$ and is independent of $X_{ij}$. Model 1 considers the regular linear model as a baseline for other nonlinear models.

Model 2:

with $\mu_{0}=-3,\ \mu_{1}=0,\ \beta_{0}^{\textnormal{T}}=(10,24,15,20)$, and $\beta_{1}^{\textnormal{T}}=(20,17,10)$. $X_{i}$ is a two-dimensional vector, $X_{i1}\sim\textup{Beta}(3,4)$, $X_{i2}\sim\textup{Unif}[-2,2]$, and they are independent of each other. The variable used for randomization is an additional variable that takes a value in $\{1,2,3,4\}$ with probability $0.2,0.3,0.3,0.2$ and is independent of $X_{ij}$. Model 2 is an additive but nonlinear model of covariates.

Model 3:

with $\mu_{0}=5,\ \mu_{1}=2,\ \beta_{0}^{\textnormal{T}}=(42,83)$, and $\beta_{1}^{\textnormal{T}}=(30,75)$. $X_{i}$ is a four-dimensional vector, $X_{i1}\sim\textup{Beta}(3,4)$, $X_{i2}\sim\textup{Unif}[-2,2]$, $X_{i3}\sim\mathcal{N}(0,1)$, $X_{i4}\sim\textup{Unif}[0,2]$, and they are independent of each other. The variable used for randomization is an additional variable that takes a value in $\{1,2\}$ with probability $0.4,0.6$ and is independent of $X_{ij}$. Model 3 is a nonlinear model including interaction terms of covariates, hence is more complex.

Model 4:

with $\mu_{0}=5,\ \mu_{1}=5,\ \beta_{0}^{\textnormal{T}}=(20,30,50)$, and $\beta_{1}^{\textnormal{T}}=(20,30,65)$. $X_{i}$ is a two-dimensional vector, $X_{i1}\sim\textup{Beta}(3,4)$, $X_{i2}\sim\textup{Unif}[-2,2]$, and they are independent of each other. $\mathds{1}_{S_{i}=1}$ is an indicator function that equals $1$ if $S_{i}=1$ and $0$ otherwise. $\mathds{1}_{S_{i}=-1}$ is defined likewise. $S_{i}$ is the variable used for randomization, it is an additional variable that takes a value in $\{1,-1\}$ with equal probability and is independent of $X_{ij}$. Model 4 further takes the interaction between covariates and stratum into consideration.

For low-dimensional data-generating models, we apply linear regression, local linear kernel (kernel), natural spline (nspline), neural networks (nnet), and random forest (rf) to approach $\hat{h}_{[k]}(X_{i},a)$. In the result tables, $\hat{\tau}$ denotes the stratum-common estimators ($\hat{h}_{[k]}(X_{i},a)$ are the same for all strata $k,\ k=1,\dots,K$) and $\tilde{\tau}$ denotes the stratum-specific estimators ($\hat{h}_{[k]}(X_{i},a)$ can be different in each stratum).