High-dimensional regression adjustments in randomized experiments

Significance

As datasets get larger and more complex, there is a growing interest in using machine-learning methods to enhance scientific analysis. In many settings, considerable work is required to make standard machine-learning methods useful for specific scientific applications. We find, however, that in the case of treatment effect estimation with randomized experiments, regression adjustments via machine-learning methods designed to minimize test set error directly induce efficient estimates of the average treatment effect. Thus, machine-learning methods can be used out of the box for this task, without any special-case adjustments.

Abstract

We study the problem of treatment effect estimation in randomized experiments with high-dimensional covariate information and show that essentially any risk-consistent regression adjustment can be used to obtain efficient estimates of the average treatment effect. Our results considerably extend the range of settings where high-dimensional regression adjustments are guaranteed to provide valid inference about the population average treatment effect. We then propose cross-estimation, a simple method for obtaining finite-sample–unbiased treatment effect estimates that leverages high-dimensional regression adjustments. Our method can be used when the regression model is estimated using the lasso, the elastic net, subset selection, etc. Finally, we extend our analysis to allow for adaptive specification search via cross-validation and flexible nonparametric regression adjustments with machine-learning methods such as random forests or neural networks.

Randomized controlled trials are often considered the gold standard for estimating the effect of an intervention, as they allow for simple model-free inference about the average treatment effect on the sampled population. Under mild conditions, the mean observed outcome in the treated sample minus the mean observed outcome in the control sample is a consistent and unbiased estimator for the population average treatment effect.

However, the fact that model-free inference is possible in randomized controlled trials does not mean that it is always optimal: As argued by Fisher (1), if we have access to auxiliary features that are related to our outcome of interest via a linear model, then controlling for these features using ordinary least squares will reduce the variance of the estimated average treatment effect without inducing any bias. This line of research has been thoroughly explored: Under low-dimensional asymptotics where the problem specification remains fixed while the number of samples grows to infinity, it is now well established that regression adjustments are always asymptotically helpful—even in misspecified models—provided we add full treatment-by-covariate interactions to the regression design and use robust standard errors (2–10).

The characteristics of high-dimensional regression adjustments are less well understood. In a recent advance, Bloniarz et al. (11) show that regression adjustments are at least sometimes helpful in high dimensions: Given an “ultrasparsity” assumption from the high-dimensional inference literature, they establish that regression adjustments using the lasso (12, 13) are more efficient than model-free inference. This result, however, leaves a substantial gap between the low-dimensional regime—where regression adjustments are always asymptotically helpful—and the high-dimensional regime where we have only special-case results.

In this paper, we show that high-dimensional regression adjustments to randomized controlled trials work under much greater generality than previously known. We find that any regression adjustment with a free intercept yields unbiased estimates of the treatment effect. This result is agnostic as to whether the regression model was obtained using the lasso, the elastic net (14), subset selection, or any other method that satisfies this criterion. We also propose a simple procedure for building practical confidence intervals for the average treatment effect.

Furthermore, we show that the precision of the treatment effect estimates obtained by such regression adjustments depends only on the prediction risk of the fitted regression adjustment. In particular, any risk-consistent regression adjustment can be made to yield efficient estimates of the average treatment effect in the sense of refs. 15–18. Thus, when choosing which regression adjustment to use, practitioners are justified in using standard model selection tools that aim to control prediction error, e.g., Mallow’s Cp or cross-validation.

This finding presents a striking contrast to the theory of high-dimensional regression adjustments in observational studies. In a setting where treatment propensity may depend on covariates, simply fitting low-risk regression models to the treatment and control samples via cross-validation is not advised, as there exist regression adjustments that have low predictive error but yield severely biased estimates of the average treatment effect (19–22). Instead, special-case procedures are needed: For example, Belloni et al. (21) advocate a form of augmented model selection that protects against bias at the cost of worsening the predictive performance of the regression model. The tasks of fitting good high-dimensional regression adjustments to randomized vs. observational data thus present qualitatively different challenges.

The first half of this paper develops a theory of regularized regression adjustments with high-dimensional Gaussian designs. This analysis enables us to highlight the connection between the predictive accuracy of the regression adjustment and the precision of the resulting treatment effect estimate and also to considerably improve on theoretical guarantees available in prior work. In the second half of the paper, we build on these insights to develop cross-estimation, a practical method for inference about average treatment effects that can be paired with either high-dimensional regularized regression or nonparametric machine-learning methods.

1. Setting and Notation

We frame our analysis in terms of the Neyman–Rubin potential outcomes model (23, 24). Given n i.i.d. observations $\left(X_i,Y_i,W_i\right),$ $i=1,\hspace{0.17em}2,\hspace{0.17em}\dots ,\hspace{0.17em}n,$ we posit potential outcomes $Y_i^{\left(1\right)}$ and $Y_i^{\left(0\right)};$ then, the outcome that we actually observe is $Y_i=Y_i^{\left(W_i\right)}.$ We focus on randomized controlled trials, where $W_i$ is independent of all pretreatment characteristics,

$$\left\{X_i,Y_i^{\left(0\right)},Y_i^{\left(1\right)}\right\}\perp W_i.$$ [1]

We take the predictors to be generated as $X_i\sim F(\cdot )\in {\mathrm{ℝ}}^p$ and assume a homoskedastic linear model in each arm,

$$Y_i=c^{\left(W_i\right)}+X_i\cdot {\beta }^{\left(W_i\right)}+{\epsilon }_i^{\left(W_i\right)},c^{\left(w\right)}\in \mathrm{ℝ},{\beta }^{\left(w\right)}\in {\mathrm{ℝ}}^p,$$ [2]

for $w=0,\hspace{0.17em}1,$ where ${\epsilon }_i^{\left(W_i\right)}$ is mean-zero noise with variance ${\sigma }^2;$ more general models are considered later. We use the notation $n_0=\left|\left\{i:W_i=0\right\}\right|$ and $n_1=\left|\left\{i:W_i=1\right\}\right|.$ We study inference about the average treatment effect $\tau =\mathbb{E}\left[Y\left(1\right)-Y\left(0\right)\right].$ In our analysis, it is sometimes also convenient to study estimation of the conditional average treatment effect

$$\begin{array}{c}\overline{\tau }=\frac{1}{n}{\displaystyle \sum _{i=1}^n\mathbb{E}\left[Y_i^{\left(1\right)}-Y_i^{\left(0\right)}|X_i\right]}\\ =\overline{X}\cdot \left({\beta }^{\left(1\right)}-{\beta }^{\left(0\right)}\right)+c^{\left(1\right)}-c^{\left(0\right)}.\end{array}$$ [3]

As discussed by ref. 17, good estimators for $\overline{\tau }$ are generally good estimators for τ and vice versa. In the homogeneous treatment effects model $Y_i=c+X_i\cdot \beta +W_i\tau +{\epsilon }_i,$ τ and $\overline{\tau }$ coincide.

All technical derivations can be found in the Supporting Information, A. Proofs.

2. Regression Adjustments with Gaussian Designs

Suppose that we have obtained parameter estimates ${\widehat{c}}^{\left(w\right)},$ ${\widehat{\beta }}^{\left(w\right)},$ $w\in \left\{\mathrm{0,1}\right\}$ for the linear model Eq. 2 via the lasso, the elastic net, or any other method. We then get a natural estimator for the average treatment effect:

$$\widehat{\tau }=\overline{X}\cdot \left({\widehat{\beta }}^{\left(1\right)}-{\widehat{\beta }}^{\left(0\right)}\right)+{\widehat{c}}^{\left(1\right)}-{\widehat{c}}^{\left(0\right)}.$$ [4]

In the case where ${\widehat{\beta }}^{\left(w\right)}$ is the ordinary least-squares estimator for $\widehat{\tau },$ the behavior of this estimator has been carefully studied by refs. 8 and 10. Our goal is to characterize its behavior for generic regression adjustments ${\widehat{\beta }}^{\left(w\right)},$ all while allowing the number of predictors p to be much larger than the sample size n.

The only assumption that we make on the estimation scheme is that it be centered: For $w\in \left\{\mathrm{0,1}\right\},$

$${\overline{Y}}_w={\overline{X}}_w\cdot {\widehat{\beta }}^{\left(w\right)}+{\widehat{c}}^{\left(w\right)};$$ [5]

i.e., the mean of the predicted outcomes matches that of the observed outcomes; and ${\widehat{\beta }}^{\left(w\right)}$ is translation invariant and depends only on

$${\mathrm{ℱ}}_{\beta }={\left\{X_i-{\overline{X}}_{W_i},Y_i-{\overline{Y}}_{W_i},W_i\right\}}_{i=1}^n.$$ [6]

Here, ${\overline{X}}_w$ and ${\overline{Y}}_w$ denote the mean of the outcomes $Y_i$ and features $X_i$ over all observations with $W_i=w.$ Algorithmically, a simple way to enforce this constraint is to first center the training samples $X_i\to X_i-{\overline{X}}_{W_i},$ $Y_i\to Y_i-{\overline{Y}}_{W_i},$ run any regression method on these centered data, and then set the intercept using Eq. 5; this is done by default in standard software for regularized regression, such as glmnet (25). We also note that ordinary least-squares regression is always centered in this sense, even after common forms of model selection.

Now, if our regression adjustment has a well-calibrated intercept as in Eq. 5, then we can write Eq. 4 as

$$\begin{array}{c}\widehat{\tau }=\overline{X}\cdot \left({\widehat{\beta }}^{\left(1\right)}-{\widehat{\beta }}^{\left(0\right)}\right)+\left({\widehat{c}}^{\left(1\right)}-{\widehat{c}}^{\left(0\right)}\right)\\ ={\overline{Y}}_1-{\overline{Y}}_0+\left(\overline{X}-{\overline{X}}_1\right)\cdot {\widehat{\beta }}^{\left(1\right)}-\left(\overline{X}-{\overline{X}}_0\right)\cdot {\widehat{\beta }}^{\left(0\right)}.\end{array}$$ [7]

To move forward, we focus on the case where the data-generating model for $\left(X_i,\hspace{0.17em}{Y}_i\right)$ is Gaussian; i.e., $X_i\sim \mathcal{N}\left(m,\hspace{0.17em}\mathrm{\Sigma }\right)$ for some $m\in {\mathrm{ℝ}}^p$ and positive-semidefinite matrix $\mathrm{\Sigma }\in {\mathrm{ℝ}}^{p\times p},$ and $Y_i-\mathbb{E}\left[Y_i|X_i,W_i\right]\sim \mathcal{N}\left(0,{\sigma }^2\right).$ For our purpose, the key fact about Gaussian data is that the mean of independent samples is independent of the within-sample spread; i.e.,

$${\left\{X_i-{\overline{X}}_{W_i},Y_i-{\overline{Y}}_{W_i}\right\}}_{i=1}^n\perp \left\{{\overline{X}}_0,{\overline{X}}_1,{\overline{Y}}_0,{\overline{Y}}_1\right\},$$ [8]

conditionally on the treatment assignments $W_1,\hspace{0.17em}\dots ,\hspace{0.17em}{W}_n.$ Thus, because ${\widehat{\beta }}^{\left(w\right)}$ depends only on the centered data $X_i-{\overline{X}}_{W_i}$ and $Y_i-{\overline{Y}}_{W_i},$ we can derive a simple expression for the distribution of $\widehat{\tau }.$ The following is an exact finite sample result and holds no matter how large p is relative to n; a key observation is that $\overline{X}-{\overline{X}}_w$ is mean zero by randomization of the treatment assignment, for $w=0,\hspace{0.17em}1.$

Proposition 1.

Suppose that our regression scheme for ${\widehat{c}}^{\left(w\right)}$ and ${\widehat{\beta }}^{\left(w\right)}$ is centered and that our data-generating model is Gaussian as above. Then, writing ${\Vert v\Vert }_{\mathrm{\Sigma }}^2≔v^{\mathrm{\top }}\mathrm{\Sigma }\hspace{0.17em}v$ for $v\in {\mathrm{ℝ}}^p,$

$$\begin{array}{l}\widehat{\tau }-\overline{\tau }|n_0,n_1,{\widehat{\beta }}^{\left(0\right)},{\widehat{\beta }}^{\left(1\right)}\stackrel{d}{=}\mathcal{N}\left(0,A\right),\\ A=\left(\frac{1}{n_0}+\frac{1}{n_1}\right)\left({\sigma }^2+{\Vert \widehat{\overline{\beta }}-\overline{\beta }\Vert }_{\mathrm{\Sigma }}^2\right),\\ \overline{\beta }=\frac{n_1\hspace{0.17em}{\beta }^{\left(0\right)}+n_0\hspace{0.17em}{\beta }^{\left(1\right)}}{n},\hspace{1em}\widehat{\overline{\beta }}=\frac{n_1\hspace{0.17em}{\widehat{\beta }}^{\left(0\right)}+n_0\hspace{0.17em}{\widehat{\beta }}^{\left(1\right)}}{n}.\end{array}$$ [9]

If the errors in ${\widehat{\beta }}^{\left(0\right)}$ and ${\widehat{\beta }}^{\left(1\right)}$ are roughly orthogonal, then

$${\Vert \widehat{\overline{\beta }}-\overline{\beta }\Vert }_{\mathrm{\Sigma }}^2\approx \frac{n_1^2}{n^2}{\Vert {\widehat{\beta }}^{\left(0\right)}-{\beta }^{\left(0\right)}\Vert }_{\mathrm{\Sigma }}^2+\frac{n_0^2}{n^2}{\Vert {\widehat{\beta }}^{\left(1\right)}-{\beta }^{\left(1\right)}\Vert }_{\mathrm{\Sigma }}^2$$ [10]

and, in any case, twice the right-hand side is always an upper bound for the left-hand side. Thus, the distribution of $\widehat{\tau }$ effectively depends on the regression adjustments ${\widehat{\beta }}^{\left(w\right)}$ only through the excess predictive error

$${\Vert {\widehat{\beta }}^{\left(w\right)}-{\beta }^{\left(w\right)}\Vert }_{\mathrm{\Sigma }}^2=\mathbb{E}\left[\left(\left(X-m\right)\cdot {\left({\widehat{\beta }}^{\left(w\right)}-{\beta }^{\left(w\right)}\right)}^2|{\widehat{\beta }}^{\left(w\right)}\right)\right]\hspace{0.17em},$$ [11]

where the above expectation is taken over a test set example X. This implies that, in the setting of Proposition 1, the main practical concern in choosing which regression adjustment to use is to ensure that ${\widehat{\beta }}^{\left(w\right)}$ has low predictive error.

The above result is conceptually related to recent work by Berk et al. (3) (also ref. 26), who showed that the accuracy of low-dimensional covariate adjustments using ordinary least-squares regression depends on the mean-squared error of the regression fit; they also advocate using this connection to provide simple asymptotic inference about τ. Here, we showed that a similar result holds for any regression adjustment on Gaussian designs, even in high dimensions; and in the second half of this paper we discuss how to move beyond the Gaussian case.

Risk Consistency and the Lasso.

As stated, Proposition 1 provides the distribution of $\widehat{\tau }$ conditionally on ${\widehat{\beta }}^{\left(w\right)}$ and so is not directly comparable to related results in the literature. However, whenever ${\widehat{\beta }}^{\left(w\right)}$ is risk consistent in the sense that

$$R\left({\widehat{\beta }}^{\left(w\right)}\right)≔{\Vert {\widehat{\beta }}^{\left(w\right)}-{\beta }^{\left(w\right)}\Vert }_{\mathrm{\Sigma }}^2{\to }_{p}0,$$ [12]

for $w=0,\hspace{0.17em}1,$ we can asymptotically omit the conditioning.

Theorem 2.

Suppose that, under the conditions of Proposition 1, we have a sequence of problems where ${\widehat{\beta }}^{\left(w\right)}$ is risk consistent (Eq. 12), and $\mathrm{\mathbb{P}}\left[W=1\right]\to \pi .$ Then,

$$\sqrt{n}\left(\widehat{\tau }-\overline{\tau }\right)\Rightarrow \mathcal{N}\left(0,\frac{{\sigma }^2}{\pi \left(1-\pi \right)}\right),$$ [13]

or, in other words, $\widehat{\tau }$ is efficient for estimating $\overline{\tau }$ (15–18).

In the case of the lasso, Theorem 2 lets us substantially improve over the best existing guarantees in the literature (11). The lasso estimates ${\widehat{\beta }}^{\left(w\right)}$ as the minimizer over β of

$${\displaystyle \sum _{\left\{i:W_i=w\right\}}\frac{1}{2}{\left(Y_i-{\overline{Y}}_w-\left(X_i-{\overline{X}}_w\right)\cdot \beta \right)}^2+n_w\lambda {\Vert \beta \Vert }_1},$$ [14]

for some penalty parameter $\lambda >0.$ Typically, the lasso is used when we believe a sparse regression adjustment to be appropriate. In our setting, it is well known that the lasso satisfies $R\left({\widehat{\beta }}^{\left(w\right)}\right)={\mathcal{O}}_P\left({\Vert \mathrm{\Sigma }\Vert }_{\text{op}}^2{\Vert {\beta }^{\left(w\right)}\Vert }_0\log \left(p\right)/n_w\right),$ provided the penalty parameter λ is well chosen and $\mathrm{\Sigma }$ does not allow for too much correlation between features (27, 28).

Thus, whenever we have a sequence of problems as in Theorem 2 where ${\beta }^{\left(w\right)}$ is k-sparse, i.e., ${\beta }^{\left(w\right)}$ has at most k nonzero entries, and $k\log \left(p\right)/n\to 0,$ we find that $\widehat{\tau }$ is efficient in the sense of Eq. 13. Note that this result is much stronger than the related result of ref. 11, which shows that lasso regression adjustments yield efficient estimators $\widehat{\tau }$ in an ultrasparse regime with $k\ll \sqrt{n}/\log \left(p\right).$

To illustrate the difference between these two results, it is well known that if $k\ll \sqrt{n}/\log \left(p\right),$ then it is possible to do efficient inference about the coefficients of the underlying parameter vector β (29–31), and so the result of ref. 11 is roughly in line with the rest of the literature on high-dimensional inference. Conversely, if we have only $k\ll n/\log \left(p\right),$ accurate inference about the coefficients of β is in general impossible without further conditions on the covariance of X (32, 33). However, we have shown that we can still carry out efficient inference about τ. In other words, the special structure present in randomized trials means that much more is possible than in the generic high-dimensional regression setting.

Inconsistent Regression Adjustments.

Even if our regression adjustment ${\widehat{\beta }}^{\left(w\right)}$ is not risk consistent, we can still use Proposition 1 to derive unconditional results about $\widehat{\tau }$ whenever

$$R\left(\widehat{\overline{\beta }}\right)≔{\Vert \widehat{\overline{\beta }}-\overline{\beta }\Vert }_{\mathrm{\Sigma }}^2{\to }_p{R}_{\infty }.$$ [15]

We illustrate this phenomenon in the case of ridge regression, where regression adjustments generally reduce—but do not eliminate—excess test-set risk. Recall that ridge regression estimates ${\widehat{\beta }}^{\left(w\right)}$ as the minimizer over β of

$${\displaystyle \sum _{\left\{i:W_i=w\right\}}\frac{1}{2}{\left(Y_i-{\overline{Y}}_w-\left(X_i-{\overline{X}}_w\right)\cdot \beta \right)}^2+n_w\lambda {\Vert \beta \Vert }_2^2}.$$ [16]

The following result relies on random-matrix theoretic tools for analyzing the predictive risk of ridge regression (34).

Theorem 3.

Suppose we have a sequence of problems in the setting of Proposition 1 with $n,\hspace{0.17em}p\to \infty$ and $p/n\to \gamma \in \left(0,\infty \right),$ such that the spectrum of the covariance $\mathrm{\Sigma }$ has a weak limit. Following ref. 34, suppose moreover that the true parameters ${\beta }^{\left(0\right)}$ and ${\beta }^{\left(1\right)}$ are independently and randomly drawn from a random-effects model with

$$\mathbb{E}\left[{\beta }^{\left(w\right)}\right]=0\hspace{1em}and\hspace{1em}\text{Var}\left[{\beta }^{\left(w\right)}\right]=\frac{{\alpha }^2}{p}{I}_{p\times p},\hspace{0.17em}\hspace{1em}with\hspace{0.17em}\alpha >0.$$ [17]

Then, selecting ${\widehat{\beta }}^{\left(w\right)}$ in Eq. 4 via ridge regression tuned to minimize prediction error, and with $\mathrm{\mathbb{P}}\left[W=1\right]\to \pi ,$ we get $\sqrt{n}\left(\widehat{\tau }-\overline{\tau }\right)\Rightarrow \mathcal{N}\left(0,S\right),$

$$S=2{\sigma }^2+\frac{{\alpha }^2}{\gamma }\left(\frac{\pi }{v_0\left(-\frac{\gamma {\sigma }^2}{{\alpha }^2\left(1-\pi \right)}\right)}+\frac{1-\pi }{v_1\left(-\frac{\gamma {\sigma }^2}{{\alpha }^2\pi }\right)}\right),$$ [18]

where the $v_w\left(-\lambda \right)$ are the companion Stieltjes transforms of the limiting empirical spectral distributions for the treated and control samples, as defined in the proof.

To interpret this result, we note that the quantity $v_w\left(-\lambda \right)$ can also be induced via the limit (35, 36)

$$\frac{1}{n_w}tr\left({\left(\frac{1}{n_w}{\displaystyle \sum }_{\left\{i:W_i=w\right\}}{X}_i{X}_i^{\mathrm{\top }}+\lambda I_{n_w\times n_w}\right)}^{-1}\right){\to }_p{v}_w\left(-\lambda \right),\hspace{0.17em}\hspace{0.5em}\text{for}\hspace{0.17em}\lambda >0.$$

Finally, we note that the limiting variance of $\widehat{\tau }-\overline{\tau }$ obtained via ridge regression above is strictly smaller than the corresponding variance of the unadjusted estimator, which converges to $\left({\sigma }^2+\left({\pi }^2+\left(1-{\pi }^2\right)\right){\alpha }^2\mathrm{tr}\left(\mathrm{\Sigma }\right)/p\right)/\left(\pi \left(1-\pi \right)\right);$ this is because optimally tuned ridge regression strictly improves over the “null” model ${\widehat{\beta }}^{\left(w\right)}=0$ in terms of its predictive accuracy.

3. Practical Inference with Cross-Estimation

In the previous section, we found that—given Gaussianity assumptions—generic regression adjustments yield unbiased estimates of the average treatment effect and also that low-risk regression adjustments lead to high-precision estimators. Here, we seek to build on this insight and to develop simple inferential procedures about τ and $\overline{\tau }$ that attain the above efficiency guarantees, all while remaining robust to deviations from Gaussianity or homoskedasticity.

Our approach is built around cross-estimation, a procedure inspired by data splitting and the work of refs. 37 and 38. We first split our data into K equally sized folds (e.g., $K=5\hspace{0.17em}\text{or}\hspace{0.17em}10$) and then, for each fold $k=1,\hspace{0.17em}\dots ,\hspace{0.17em}K,$ we compute

$$\begin{array}{l}{\widehat{\tau }}^{\left(k\right)}={\overline{Y}}_1^{\left(k\right)}-{\overline{Y}}_0^{\left(k\right)}+\left({\overline{X}}^{\left(k\right)}-{\overline{X}}_1^{\left(k\right)}\right)\cdot {\widehat{\beta }}^{\left(1,-k\right)}\\ -\left({\overline{X}}^{\left(k\right)}-{\overline{X}}_0^{\left(k\right)}\right)\cdot {\widehat{\beta }}^{\left(0,-k\right)}.\end{array}$$ [19]

Here, ${\overline{Y}}_1^{\left(k\right)},$ ${\overline{Y}}_0^{\left(k\right)},$ etc., are moments taken over the kth fold, whereas ${\widehat{\beta }}^{\left(1,\hspace{0.17em}-k\right)}$ and ${\widehat{\beta }}^{\left(0,\hspace{0.17em}-k\right)}$ are centered regression estimators computed over the $K-1$ other folds. We then obtain an aggregate estimate $\widehat{\tau }={\displaystyle {\sum }_{k=1}^K{\widehat{\tau }}^{\left(k\right)}\hspace{0.17em}{n}^{\left(k\right)}/n},$ where $n^{\left(k\right)}$ is the number of observations in the kth fold. An advantage of this construction is that an analog to the relation Eq. 8 now automatically holds, and thus our treatment effect estimator $\widehat{\tau }$ is unbiased without assumptions. Note that the result below references both the average treatment effect τ and the conditional average treatment effect $\overline{\tau }.$

Theorem 4.

Suppose that we have n independent and identically distributed samples satisfying Eq. 1, drawn from a linear model Eq. 2, where $X_i$ has finite first moments and the conditional variance of $Y_i^{\left(w\right)}$ given $X_i$ may vary. Then, $\mathbb{E}\left[\widehat{\tau }|X_1,\hspace{0.17em}\dots ,X_n\right]=\overline{\tau }.$ If, moreover, the ${\widehat{\beta }}^{\left(w,\hspace{0.17em}-k\right)}$ are all risk consistent in the sense of Eq. 12 for $k=1,\hspace{0.17em}\dots ,\hspace{0.17em}K,$ and both the signals $X_i\cdot {\beta }^{\left(w\right)}$ and residuals $Y_i-\mathbb{E}\left[Y_i|X_i,W_i=w\right]$ are asymptotically Gaussian when averaged, then writing ${\overline{\sigma }}_w^2=\mathbb{E}\left[\text{Var}\left[Y_i^{\left(w\right)}|X_i\right]\right],$ we have

$$\sqrt{n}\left(\widehat{\tau }-\tau \right)\Rightarrow \mathcal{N}\left(0,\frac{{\overline{\sigma }}_0^2}{1-\pi }+\frac{{\overline{\sigma }}_1^2}{\pi }+{\Vert {\beta }^{\left(1\right)}-{\beta }^{\left(0\right)}\Vert }_{\mathrm{\Sigma }}^2\right)\hspace{0.17em}.$$ [20]

In the homoskedastic case, i.e., when the variance of $Y^{\left(w\right)}$ conditionally on X does not depend on X, then the above is efficient. With heteroskedasticity, the above is no longer efficient because we are in a linear setting and so inverse-variance weighting could improve precision; however, Eq. 20 can still be used as the basis for valid inference about τ.

Confidence Intervals via Cross-Estimation.

Another advantage of cross-estimation is that it allows for moment-based variance estimates for $\widehat{\tau }.$ Here, we discuss practical methods for building confidence intervals that cover the average treatment effect τ. We can verify that the variance of ${\widehat{\tau }}^{\left(k\right)}$ is $V_k$ after conditioning on the ${\widehat{\beta }}^{\left(w,\hspace{0.17em}-k\right)}$ and $n_w^{\left(k\right)},$ with

$$V_k={\displaystyle \sum _{w\in \left\{0,\hspace{0.17em}1\right\}}\frac{1}{n_w^{\left(k\right)}}\text{Var}\left[Y^{\left(w\right)}-X\cdot {\widehat{\overline{\beta }}}^{\left(-k\right)}|{\widehat{\overline{\beta }}}^{\left(-k\right)}\right]}.$$ [21]

Now, the above moments correspond to observable quantities on the kth data fold, so we immediately obtain a moment-based plug-in estimator ${\widehat{V}}_k$ for $V_k.$ Finally, we build α-level confidence intervals for τ as

$$\tau \in \widehat{\tau }\pm z_{1-\alpha /2}\widehat{V},\hspace{0.17em}\widehat{V}={\displaystyle \sum _{k=1}^K{\left(\frac{n^{\left(k\right)}}{n}\right)}^2{\widehat{V}}_k},$$ [22]

where $z_{1-\alpha /2}$ is the appropriate standard Gaussian quantile. In the setting of Theorem 4, i.e., with risk consistency and bounded second moments, we can verify that the ${\widehat{\tau }}^{\left(k\right)}$ are asymptotically uncorrelated and so the above confidence intervals are asymptotically exact.

Cross-Validated Cross-Estimation.

High-dimensional regression adjustments usually rely on a tuning parameter that controls the amount of regularization, e.g., the parameter λ for the lasso and ridge regression. Although theory provides some guidance on how to select λ, practitioners often prefer to use computationally intensive methods such as cross-validation.

Now, our procedure in principle already allows for cross-validation: If we estimate ${\widehat{\beta }}^{\left(0,\hspace{0.17em}-k\right)}$ in Eq. 19 via any cross-validated regression adjustment that relies only on all but the kth data folds, then ${\widehat{\tau }}^{\left(k\right)}$ will be unbiased for τ. However, this requires running the full cross-validated algorithm K times, which can be very expensive computationally.

Here, we show how to obtain good estimates of $\widehat{\tau }$ using only a single round of cross-validation. First, we specify K regression folds, and for each $k\in \left\{1,\hspace{0.17em}\dots ,\hspace{0.17em}K\right\}$ and $w\in \left\{\mathrm{0,1}\right\}$ we compute ${\overline{X}}_{k,\hspace{0.17em}w}$ and ${\overline{Y}}_{k,\hspace{0.17em}w}$ as the mean of all observations in the kth fold with $W_i=w.$ Next, we center the data such that ${\stackrel{\sim }{X}}_i=X_i-{\overline{X}}_{k,\hspace{0.17em}{W}_i}$ and ${\stackrel{\sim }{Y}}_i=Y_i-{\overline{Y}}_{k,\hspace{0.17em}{W}_i}$ for all observations in the kth fold. Finally, we estimate ${\widehat{\beta }}^{\left(w,\hspace{0.17em}-k\right)}$ by running a standard out-of-the-box cross-validated algorithm (e.g., cv. glmnet for R) on the $\left({\stackrel{\sim }{X}}_i,\hspace{0.17em}{\stackrel{\sim }{Y}}_i,\hspace{0.17em}{W}_i\right)$ triples with the same K folds as specified before and then use Eq. 19 to compute $\widehat{\tau }.$

The actual estimator that we use to estimate ${\widehat{\beta }}^{\left(0\right)}$ and ${\widehat{\beta }}^{\left(1\right)}$ in our experiments is inspired by the procedure of Imai and Ratkovic (39). Our goal is to let the lasso learn shared “main effects” for the treatment and control groups. To accomplish this, we first run a $2p$-dimensional lasso problem,

$$\widehat{\beta },\hspace{0.17em}\widehat{\gamma }={\text{argmin}}_{\beta ,\hspace{0.17em}\gamma }\left\{\lambda \left({\Vert \beta \Vert }_1+{\Vert \gamma \Vert }_1\right)+{\displaystyle \sum {\left({\stackrel{\sim }{Y}}_i-\left({\stackrel{\sim }{X}}_i\cdot \beta +\left(2W_i-1\right){\stackrel{\sim }{X}}_i\cdot \gamma \right)\right)}^2}\right\},$$ [23]

and then set ${\widehat{\beta }}^{\left(0\right)}=\widehat{\beta }-\widehat{\gamma }$ and ${\widehat{\beta }}^{\left(1\right)}=\widehat{\beta }+\widehat{\gamma }.$ We simultaneously tune λ and estimate τ by cross-validated cross-estimation as discussed above. When all our data are Gaussian, this procedure is exactly unbiased by the same argument as used in Proposition 1; and even when X is not Gaussian, it appears to work well in our experiments.

4. Nonparametric Machine-Learning Methods

In our discussion so far, we have focused on treatment effect estimation using high-dimensional, linear regression adjustments and showed how to provide unbiased inference about τ under general conditions. Here, we show how to extend our results about cross-estimation to general nonparametric regression adjustments obtained using, e.g., neural networks or random forests (40). We assume a setting where

$$\mathbb{E}\left[Y\left(w\right)|X=x\right]={\mu }^{\left(w\right)}\left(x\right)$$

for some unknown regression functions ${\mu }^{\left(w\right)}\left(x\right),$ and our goal is to leverage estimates ${\widehat{\mu }}^{\left(w\right)}\left(x\right)$ obtained using any machine-learning method to improve the precision of $\widehat{\tau },$ as*^,†

$$\widehat{\tau }=\frac{1}{n}{\displaystyle \sum _{i=1}^n\left({\widehat{\mu }}^{\left(1,\hspace{0.17em}-i\right)}\left(X_i\right)-{\widehat{\mu }}^{\left(0,\hspace{0.17em}-i\right)}\left(X_i\right)\right)}+{\displaystyle \sum _{\left\{i:W_i=1\right\}}\frac{Y_i-{\widehat{\mu }}^{\left(1,\hspace{0.17em}-i\right)}\left(X_i\right)}{n_1}}-{\displaystyle \sum _{\left\{i:W_i=0\right\}}\frac{Y_i-{\widehat{\mu }}^{\left(0,\hspace{0.17em}-i\right)}\left(X_i\right)}{n_0}},$$ [24]

where ${\widehat{\mu }}^{\left(w,\hspace{0.17em}-i\right)}$ is any estimator that does not depend on the ith training example; for random forests, we set ${\widehat{\mu }}^{\left(w,\hspace{0.17em}-i\right)}\left(X_i\right)$ to be the “out-of-bag” prediction at $X_i.$ To motivate Eq. 24, we start from Eq. 7 and expand out terms using the relation

$$\left(\overline{X}-{\overline{X}}_1\right)\cdot {\widehat{\beta }}^{\left(1\right)}=\frac{1}{n}{\displaystyle \sum _{i=1}^n{\widehat{\mu }}^{\left(1\right)}\left(X_i\right)}-\frac{1}{n_1}{\displaystyle \sum _{\left\{i:W_i=1\right\}}{\widehat{\mu }}^{\left(1\right)}\left(X_i\right)},$$ [25]

where ${\widehat{\mu }}^{\left(1\right)}\left(x\right)=x\cdot {\widehat{\beta }}^{\left(1\right)}+{\widehat{c}}^{\left(1\right)},$ etc. The remaining differences between Eq. 24 and Eq. 7 are due to the use of out-of-bag estimation to preserve randomization of the treatment assignment $W_i$ conditionally on the corresponding regression adjustment. We estimate the variance of $\widehat{\tau }$ using the formula

$$\widehat{V}={\displaystyle \sum _{w\in \left\{0,\hspace{0.17em}1\right\}}{\displaystyle \sum _{\left\{i:W_i=w\right\}}\frac{{\left(Y_i-\frac{n_0}{n}{\widehat{\mu }}^{\left(1,\hspace{0.17em}-i\right)}\left(X_i\right)-\frac{n_1}{n}{\widehat{\mu }}^{\left(0,\hspace{0.17em}-i\right)}\left(X_i\right)\right)}^2}{n_w\left(n_w-1\right)}}}.$$ [26]

Formal Results.

The following result characterizes the behavior of this estimator, under the assumption that the estimator is “jackknife compatible,” meaning that the expected jackknife estimate of variance for ${\widehat{\mu }}^{\left(w\right)}$ converges to 0. We define this condition in Supporting Information, A. Proofs and verify that it holds for random forests.

Theorem 5.

Suppose that $\widehat{\mu }$ is jackknife compatible. Then, the estimator $\widehat{\tau }$ in Eq. 24 is asymptotically unbiased, $\mathbb{E}\left[\widehat{\tau }|X_1,\hspace{0.17em}\dots ,X_n\right]=\overline{\tau }+o\left(1/\sqrt{n}\right).$ Moreover, if the regression adjustments ${\widehat{\mu }}_w$ are risk consistent in the sense that^‡ $1/n{\displaystyle {\sum }_{i=1}^n{\left({\widehat{\mu }}^{\left(w,\hspace{0.17em}-i\right)}\left(X_i\right)-{\mu }^{\left(w\right)}\left(X_i\right)\right)}^2}{\to }_{p}0,$ and the potential outcomes $Y_i^{\left(w\right)}$ have finite second moments, then $\widehat{\tau }$ is efficient and $\left(\widehat{\tau }-\tau \right)/{\left(\widehat{V}\right)}^{1/2}$ is asymptotically standard Gaussian.

We note that there has been considerable recent interest in using machine-learning methods to estimate heterogeneous treatment effects (42–45). In relation to this literature, our present goal is more modest: We simply seek to use machine learning to reduce the variance of treatment effect estimates in randomized experiments. This is why we obtain more general results than the papers on treatment heterogeneity.

5. Experiments

In our experiments, we focus on two specific variants of treatment effect estimation via cross-estimation. For high-dimensional linear estimation, we use the lasso-based method Eq. 23 tuned by cross-validated cross-estimation. For nonparametric estimation, we use Eq. 24 with random forest adjustments. We implement our method as an open-source R package, crossEstimation, built on top of glmnet (25) and randomForest (46) for R. Supporting Information, B. Additional Simulation Results and Table S1 have additional simulation results.

Table S1.

Coverage rates for 95% nominal confidence intervals obtained by cross-validated cross-estimation with the joint lasso procedure, Eq. 23

				p
				60		500
				n
Signal	Noise	Correlation	Design	80	200	80	200
Dense	$\sigma =0.1$	$\rho =0$	Gauss. X	0.94	0.96	0.95	0.95
	$\sigma =0.1$	$\rho =0$	Bern. X	0.95	0.95	0.94	0.95
	$\sigma =0.1$	$\rho =0.9$	Gauss. X	0.94	0.94	0.94	0.95
	$\sigma =0.1$	$\rho =0.9$	Bern. X	0.93	0.95	0.93	0.92
	$\sigma =1$	$\rho =0$	Gauss. X	0.95	0.95	0.92	0.94
	$\sigma =1$	$\rho =0$	Bern. X	0.93	0.95	0.94	0.93
	$\sigma =1$	$\rho =0.9$	Gauss. X	0.93	0.94	0.93	0.94
	$\sigma =1$	$\rho =0.9$	Bern. X	0.95	0.96	0.92	0.94
Geometric	$\sigma =0.1$	$\rho =0$	Gauss. X	0.94	0.95	0.93	0.94
	$\sigma =0.1$	$\rho =0$	Bern. X	0.92	0.95	0.94	0.96
	$\sigma =0.1$	$\rho =0.9$	Gauss. X	0.95	0.94	0.95	0.95
	$\sigma =0.1$	$\rho =0.9$	Bern. X	0.94	0.93	0.95	0.94
	$\sigma =1$	$\rho =0$	Gauss. X	0.94	0.94	0.94	0.94
	$\sigma =1$	$\rho =0$	Bern. X	0.95	0.94	0.95	0.95
	$\sigma =1$	$\rho =0.9$	Gauss. X	0.91	0.96	0.94	0.95
	$\sigma =1$	$\rho =0.9$	Bern. X	0.95	0.97	0.94	0.95
Sparse	$\sigma =0.1$	$\rho =0$	Gauss. X	0.94	0.96	0.94	0.95
	$\sigma =0.1$	$\rho =0$	Bern. X	0.95	0.95	0.92	0.94
	$\sigma =0.1$	$\rho =0.9$	Gauss. X	0.92	0.95	0.93	0.96
	$\sigma =0.1$	$\rho =0.9$	Bern. X	0.93	0.95	0.92	0.97
	$\sigma =1$	$\rho =0$	Gauss. X	0.94	0.93	0.94	0.95
	$\sigma =1$	$\rho =0$	Bern. X	0.94	0.93	0.93	0.95
	$\sigma =1$	$\rho =0.9$	Gauss. X	0.95	0.94	0.93	0.95
	$\sigma =1$	$\rho =0.9$	Bern. X	0.94	0.94	0.94	0.94

All numbers are aggregated over 500 simulation runs; the above numbers thus have a standard sampling error of roughly 0.01. Bern., Bernoulli; Gauss., Gaussian.

Simulations.

We begin by validating our method in a simple simulation setting with $Y=X\beta +W\tau +\epsilon ,$ where $\epsilon \sim \mathcal{N}\left(0,\hspace{0.17em}1\right).$ In all simulations, we set the features X to be Gaussian with autoregressive AR-ρ covariance. We compare our lasso-based cross-estimation with both the simple difference-in-means estimate $\widehat{\tau }={\overline{Y}}_1-{\overline{Y}}_0$ and the proposal of Bloniarz et al. (11) that uses lasso regression adjustments tuned by cross-validation. Our method differs from that of Bloniarz et al. (11) in that we use a different algorithm for confidence intervals and also that we use the joint lasso algorithm Eq. 23 instead of computing separate lassos in both treatment arms.

Figs. 1 and 2 display results for different choices of β, ρ, etc., while varying n. In both cases, we see that the confidence intervals produced by our cross-estimation algorithm and the method of Bloniarz et al. (11) are substantially shorter than those produced by the difference-in-means estimator. Moreover, our confidence intervals accurately represent the variance of our estimator (compare solid and dashed-dotted lines in Figs. 1 and 2, Left) and achieve nominal coverage at both the 95% and 99% levels. Conversely, especially in small samples, the method of Bloniarz et al. (11) underestimates the variance of the method and does not achieve target coverage.

Fig. 1.

Simulation results with $\beta =\left(1,\hspace{0.17em}0,\hspace{0.17em}0,\hspace{0.17em}\dots ,\hspace{0.17em}0\right),$ $\mathrm{\mathbb{P}}\left[W=1\right]=0.2,$ $\rho =0,$ and $p=500.$ All numbers are based on 1,000 simulation replications. Left panel shows both the average variance estimate $\widehat{V}$ produced by each estimator (solid lines) and the actual variance Var[$\widehat{\tau }$] of the estimator (dashed-dotted lines); note that $\widehat{V}$ is directly proportional to the squared length of the confidence interval. Right panel depicts realized coverage for both 95% confidence intervals (solid lines) and 99% confidence intervals (dashed lines).

Fig. 2.

Simulation results with β proportional to a permutation of $\left(1,\hspace{0.17em}{2}^{-1},\hspace{0.17em}{3}^{-1},\hspace{0.17em}\dots ,\hspace{0.17em}{p}^{-1}\right),$ ${\Vert \beta \Vert }_2=2,$ $\mathrm{\mathbb{P}}\left[W=1\right]=0.5,$ $\rho =0.8,$ and $p=500.$ All numbers are based on 1,000 simulation replications. The plots are produced the same way as in Fig. 1.

Understanding Attitudes Toward Welfare.

We also consider an experimental dataset collected as a part of the General Social Survey.^§ The dataset is large ($N=\mathrm{28,646}$ after preprocessing), so we know the true treatment effect essentially without error: The fraction of respondents who say we spend too much on assistance to the poor is smaller than the fraction of respondents who say we spend too much on welfare by 0.35. To test our method, we repeatedly drew subsamples of size $n=\mathrm{2,000}$ from the full dataset and examined the ability of both lasso-based and random-forest–based cross-estimation to recover the correct answer. We had $p=12$ regressors.

First, we note that both variants of cross-estimation achieved excellent coverage. Given a nominal coverage rate of 95%, the simple difference-in-means estimator, lasso-based cross-estimation, and random-forest cross-estimation had realized coverage rates of 96.3%, 96.5%, and 95.3%, respectively, over 1,000 replications. Meanwhile, given a nominal target of 99%, the realized numbers became 99.0%, 99.0%, and 99.3%. We note that this dataset has non-Gaussian features and exhibits considerable treatment effect heterogeneity.

Second, Fig. 3 depicts the reduction in squared confidence interval length for individual realizations of each method. More formally, we show boxplots of ${\widehat{V}}_{\text{lasso}/\text{rf}}/{\widehat{V}}_{\text{simple}},$ where $\widehat{V}$ is the variance estimate used to build confidence intervals. Here, we see that although cross-estimation may not improve the precision of the simple method by a large amount, it consistently improves performance by a small amount. Moreover, in this example, random forests result in a larger improvement in precision than lasso-based cross-estimation.

Fig. 3.

Reduction in squared confidence interval length achieved by random forests and a lasso-based method, relative to the simple difference-in-means estimator. Confidence intervals rely on cross-estimation. The plot was generated using 1,000 simulation replications.

6. Discussion

In many applications of machine-learning methods to causal inference, there is a concern that the risk of specification search, i.e., trying out many candidate methods and choosing the one that gives us a significant result, may reduce the credibility of empirical findings. This has led to considerable interest in methodologies that allow for complex model-fitting strategies that do not compromise statistical inference.

One prominent example is the design-based paradigm to causal inference in observational studies, whereby we first seek to build an “observational design” by looking only at the features $X_i$ and the treatment assignments $W_i$ and reveal the outcomes $Y_i$ only once the observational design has been set (48, 49). The observational design may rely on matching, inverse-propensity weighting, or other techniques. As the observational design is fixed before the outcomes $Y_i$ are revealed, practitioners can devote considerable time and creativity to fine-tuning the design without compromising their analysis.

From this perspective, we have shown that regression adjustments to high-dimensional randomized controlled trials exhibit a similar opportunity for safe specification search. Concretely, imagine that once we have collected data from a randomized experiment, we provide our analyst only with class-wise centered data: $W_i,$ $X_i-{\overline{X}}_{W_i},$ and $Y_i-{\overline{Y}}_{W_i}.$ The analyst can then use these data to obtain any regression adjustment they want, which we will then plug into Eq. 4. Our results guarantee that—at least with a random Gaussian design—the resulting treatment effect estimates will be unbiased regardless of the specification search the analyst may have done using only the class-wise centered data. Cross-estimation enables us to mimic this phenomenon with non-Gaussian data.

A. Proofs

Proof of Proposition 1.

Beginning from Eq. 3, we see that

$$\begin{array}{c}\widehat{\tau }-\overline{\tau }={\overline{Y}}^{\left(1\right)}-{\overline{Y}}^{\left(0\right)}+\left(\overline{X}-{\overline{X}}_1\right)\cdot {\widehat{\beta }}^{\left(1\right)}-\left(\overline{X}-{\overline{X}}_0\right)\cdot {\widehat{\beta }}^{\left(0\right)}\\ \hspace{1em}-\left(c^{\left(1\right)}-c^{\left(0\right)}\right)-\overline{X}\cdot \left({\beta }^{\left(1\right)}-{\beta }^{\left(0\right)}\right)\\ ={\overline{\epsilon }}^{\left(1\right)}-{\overline{\epsilon }}^{\left(0\right)}+\left(\overline{X}-{\overline{X}}_1\right)\cdot \left({\widehat{\beta }}^{\left(1\right)}-{\beta }^{\left(1\right)}\right)\\ \hspace{1em}-\left(\overline{X}-{\overline{X}}_0\right)\cdot \left({\widehat{\beta }}^{\left(0\right)}-{\beta }^{\left(0\right)}\right)\\ ={\overline{\epsilon }}^{\left(1\right)}-{\overline{\epsilon }}^{\left(0\right)}-\left({\overline{X}}_1-{\overline{X}}_0\right)\left(\widehat{\overline{\beta }}-\overline{\beta }\right),\end{array}$$

where $\overline{\beta }$ and $\widehat{\overline{\beta }}$ are as defined in the statement of Proposition 1, and ${\overline{\epsilon }}^{\left(w\right)}=n_w^{-1}{\displaystyle {\sum }_{\left\{i:W_i=w\right\}}{\epsilon }_i^{\left(w\right)}}.$ Now, because of Eq. 8, we can verify that that the three summands above are independent conditionally on $n_0,$ $n_1,$ and $\widehat{\overline{\beta }},$ with

$${\overline{\epsilon }}^{\left(w\right)}\sim \mathcal{N}\left(0,{\sigma }^2/n_w\right)\hspace{0.5em}\text{for}\hspace{0.5em}w\in \left\{\mathrm{0,1}\right\}$$

and

$$\left({\overline{X}}_1-{\overline{X}}_0\right)\left(\widehat{\overline{\beta }}-\overline{\beta }\right)\sim \mathcal{N}\left(0,\left(\frac{1}{n_0}+\frac{1}{n_1}\right){\Vert \widehat{\overline{\beta }}-\overline{\beta }\Vert }_{\mathrm{\Sigma }}^2\right).$$

Proof of Theorem 2.

Given our hypotheses, we immediately see that $n_0/n{\to }_{p}1-\pi ,$ $n_1/n{\to }_p\pi ,$ and ${\Vert \widehat{\overline{\beta }}-\overline{\beta }\Vert }_{\mathrm{\Sigma }}^2{\to }_{p}0.$ The conclusion follows from Proposition 1 via Slutsky’s theorem.

Proof of Theorem 3.

For a covariance matrix $\mathrm{\Sigma },$ we define its spectral distribution $F\left(\mathrm{\Sigma }\right)$ as the empirical distribution of its eigenvalues. We assume that, in our sequence of problems, $F\left(\mathrm{\Sigma }\right)$ converges weakly to some limiting population spectral distribution $F_P.$ Given this assumption, it is well known that the spectra of the sample covariance matrices $\widehat{\mathrm{\Sigma }}$ also converge weakly to a limiting empirical spectral distribution $F_E,$ with probability 1 (35, 36). In this notation, the companion Stieltjes transform $v\left(-\lambda \right)$ is defined as

$$v\left(-\lambda \right)=\frac{1-\gamma }{\lambda }+\gamma {\displaystyle {\int }_0^{\infty }\frac{1}{z+\lambda }d{F}_E\left(z\right)}.$$

Given these preliminaries and under the listed hypotheses, ref. 34 shows that the risk of optimally tuned ridge regression converges in probability,

$$\frac{1}{{\sigma }^2}R\left({\widehat{\beta }}_{\text{ridge}}\right){\to }_p\frac{1}{\lambda \ast v\left(-\lambda \ast \right)}-1,$$

where $\lambda \ast =\gamma {\sigma }^2/{\alpha }^2$ is the asymptotically optimal choice for λ.

Now, in our setting, we need to apply this result to the treatment and control samples separately. The asymptotically optimal regularization parameters for ${\widehat{\beta }}^{\left(0\right)}$ and ${\widehat{\beta }}^{\left(1\right)}$ are $\gamma {\sigma }^2{\alpha }^{-2}{\left(1-\pi \right)}^{-1}$ and $\gamma {\sigma }^2{\alpha }^{-2}{\pi }^{-1},$ respectively. Moreover, by spherical symmetry,

$${\left({\widehat{\beta }}^{\left(1\right)}-{\beta }^{\left(1\right)}\right)}^{\mathrm{\top }}\mathrm{\Sigma }\left({\widehat{\beta }}^{\left(0\right)}-{\beta }^{\left(0\right)}\right){\to }_{p}0.$$

Thus, together with the above risk bounds, we find that

$$\frac{1}{{\sigma }^2}{\Vert \widehat{\overline{\beta }}-\overline{\beta }\Vert }_{\mathrm{\Sigma }}^2{\to }_p{\pi }^2\left(\frac{1}{\frac{\gamma {\sigma }^2}{{\alpha }^2\left(1-\pi \right)}{v}_0\left(-\frac{\gamma {\sigma }^2}{{\alpha }^2\left(1-\pi \right)}\right)}-1\right)+{\left(1-\pi \right)}^2\left(\frac{1}{\frac{\gamma {\sigma }^2}{{\alpha }^2\pi }{v}_1\left(-\frac{\gamma {\sigma }^2}{{\alpha }^2\pi }\right)}-1\right),$$

where $v_0$ and $v_1$ are the companion Stieltjes transforms for the control and treatment samples, respectively. The desired conclusion then follows from Proposition 1.

Proof of Theorem 4.

By randomization of the treatment assignment $W_i$ we have, within the kth fold and for $w\in \left\{\mathrm{0,1}\right\},$

$$\mathbb{E}\left[{\overline{X}}_w^{\left(k\right)}-{\overline{X}}^{\left(k\right)}|\left\{X_i\right\},{\widehat{\beta }}^{\left(0,\hspace{0.17em}-k\right)},{\widehat{\beta }}^{\left(1,\hspace{0.17em}-k\right)}\right]=0,$$

and

$$\begin{array}{l}\mathbb{E}\left[{\overline{Y}}_1^{\left(k\right)}-{\overline{Y}}_0^{\left(k\right)}|\left\{X_i\right\},{\widehat{\beta }}^{\left(0,\hspace{0.17em}-k\right)},{\widehat{\beta }}^{\left(1,\hspace{0.17em}-k\right)}\right]\\ \hspace{1em}=c^{\left(1\right)}-c^{\left(0\right)}+{\overline{X}}^{\left(k\right)}\cdot \left({\beta }^{\left(1\right)}-{\beta }^{\left(0\right)}\right).\end{array}$$

Thus, we see that

$$\begin{array}{c}\mathbb{E}\left[\widehat{\tau }|\left\{X_i\right\}\right]=c^{\left(1\right)}-c^{\left(0\right)}+\left({\displaystyle \sum _{k=1}^K{\overline{X}}^{\left(k\right)}}\frac{n^{\left(k\right)}}{n}\right)\cdot \left({\beta }^{\left(1\right)}-{\beta }^{\left(0\right)}\right)\\ =c^{\left(1\right)}-c^{\left(0\right)}+\overline{X}\cdot \left({\beta }^{\left(1\right)}-{\beta }^{\left(0\right)}\right)=\overline{\tau }.\end{array}$$

Meanwhile, writing ${\mathrm{ℱ}}_k\in \left\{1,\hspace{0.17em}\dots ,\hspace{0.17em}n\right\}$ for the set of observations in the kth fold and $n_w^{\left(k\right)}$ for the number of those observations with $W_i=k,$ we can write

$$\begin{array}{c}{\widehat{\tau }}^{\left(k\right)}={\displaystyle \sum _{\left\{i\in {\mathrm{ℱ}}_k:W_i=1\right\}}\frac{Y_i-\frac{n_1^{\left(k\right)}}{n^{\left(k\right)}}{\beta }^{\left(0\right)}\cdot X_i-\frac{n_0^{\left(k\right)}}{n^{\left(k\right)}}{\beta }^{\left(1\right)}\cdot X_i}{n_1^{\left(k\right)}}}\\ -{\displaystyle \sum _{\left\{i\in {\mathrm{ℱ}}_k:W_i=0\right\}}\frac{Y_i-\frac{n_1^{\left(k\right)}}{n^{\left(k\right)}}{\beta }^{\left(0\right)}\cdot X_i-\frac{n_0^{\left(k\right)}}{n^{\left(k\right)}}{\beta }^{\left(1\right)}\cdot X_i}{n_0^{\left(k\right)}}}\\ +{\displaystyle \sum _{\left\{i\in {\mathrm{ℱ}}_k\right\}}\frac{{\left(-1\right)}^{W_i}}{n_{W_i}^{\left(k\right)}}(\frac{n_1^{\left(k\right)}}{n^{\left(k\right)}}\left({\widehat{\beta }}^{\left(0,\hspace{0.17em}-k\right)}-{\beta }^{\left(0\right)}\right)\cdot X_i}\\ +\frac{n_0^{\left(k\right)}}{n^{\left(k\right)}}\left({\widehat{\beta }}^{\left(1,\hspace{0.17em}-k\right)}-{\beta }^{\left(1\right)}\right)\cdot X_i).\end{array}$$

Now, by consistency of the regression adjustment, the third summand decays faster than $1/\sqrt{n}$ and so can asymptotically be ignored. Rearranging the first two summands, we get

$${\widehat{\tau }}^{\left(k\right)}=\frac{1}{n^{\left(k\right)}}{\displaystyle \sum _{\left\{i\in {\mathrm{ℱ}}_k\right\}}\left({\beta }^{\left(1\right)}-{\beta }^{\left(0\right)}\right)\cdot X_i}-{\displaystyle \sum _{\left\{i\in {\mathrm{ℱ}}_k\right\}}{\left(-1\right)}^{W_i}\frac{Y_i-X_i\cdot {\beta }^{\left(W_i\right)}}{n_{W_i}^{\left(k\right)}}}+o_P\left(\frac{1}{\sqrt{n}}\right),$$

which has the desired asymptotic variance and is asymptotically Gaussian under the stated regularity conditions.

Proof of Theorem 5.

We begin with a definition. The estimator $\widehat{\mu }$ is jackknife compatible if, for any $w\in \left\{\mathrm{0,1}\right\}$ and a new independently drawn test point X,

$$\mathbb{E}\left[{\displaystyle \sum _{\left\{i:W_i=w\right\}}{\left({\widehat{\mu }}^{\left(w,\hspace{0.17em}-i\right)}\left(X\right)-{\widehat{\mu }}^{\left(w\right)}\left(X\right)\right)}^2|n_w}\right]\le a\left(n_w\right)$$ [S1]

for some sequence $a\left(n\right)\to 0.$ The quantity inside the left-hand expectation is the popular jackknife estimate of variance for ${\widehat{\mu }}^{\left(w\right)};$ the condition requires that this variance estimate converge to 0 in expectation. We note that this condition is very weak: Most classical statistical estimators will satisfy this condition with $a\left(n\right)=\mathcal{O}\left(1/n\right),$ whereas subsampled random forests of the type studied in ref. 44 satisfy it with $a\left(n\right)=\mathcal{O}\left(s/n\right),$ where s is the subsample size.

Now, as in the proof of Theorem 4, we write $\widehat{\tau }$ as

$$\frac{1}{n_1}{\displaystyle \sum _{\left\{i:W_i=1\right\}}\left(Y_i-{\mu }^{\left(1\right)}\left(X_i\right)\right)}-\frac{1}{n_0}{\displaystyle \sum _{\left\{i:W_i=0\right\}}\left(Y_i-{\mu }^{\left(0\right)}\left(X_i\right)\right)}+\frac{1}{n}{\displaystyle \sum _{i=1}^n\left({\mu }^{\left(1\right)}\left(X_i\right)-{\mu }^{\left(0\right)}\left(X_i\right)\right)}+R,$$

where R is a residual term

$$R={\displaystyle \sum _{i=1}^n\frac{{\left(-1\right)}^{W_i}}{n_{W_i}}\left(\frac{n_0}{n}\left({\widehat{\mu }}^{\left(1,\hspace{0.17em}-i\right)}\left(X_i\right)-{\mu }^{\left(1\right)}\left(X_i\right)\right)+\frac{n_1}{n}\left({\widehat{\mu }}^{\left(0,\hspace{0.17em}-i\right)}\left(X_i\right)-{\mu }^{\left(0\right)}\left(X_i\right)\right)\right)}.$$

The main component of $\widehat{\tau }$ is an unbiased, efficient estimator for τ. It remains to show that R is asymptotically unbiased given jackknife compatibility and is moreover asymptotically negligible if $\widehat{\mu }$ is risk consistent.

For the remainder of the proof, we focus on the setting where the regression adjustments ${\widehat{\mu }}^{\left(0\right)}$ and ${\widehat{\mu }}^{\left(1\right)}$ are computed separately on samples with $W_i=0$ and $W_i=1,$ respectively. The reason R may not be exactly unbiased is that ${\widehat{\mu }}^{\left(1,\hspace{0.17em}-i\right)}\left(X_i\right)$ is a function of $n_1$ observations if $W_i=0,$ whereas it is a function of $n_1-1$ observations if $W_i=1,$ and this effect can create biases. Our goal is to show, however, that these biases are small for any jackknife-compatible estimator. To do so, we first define a “leave-two-out” approximation to R,

$$R_2={\displaystyle \sum _{i=1}^n\frac{{\left(-1\right)}^{W_i}}{n_0{n}_1}{\displaystyle \sum _{\left\{j:W_j\ne W_i\right\}}\left(\frac{n_0}{n}\left({\widehat{\mu }}^{\left(1,\hspace{0.17em}-\left\{i,\hspace{0.17em}j\right\}\right)}\left(X_i\right)-{\mu }^{\left(1\right)}\left(X_i\right)\right)+\frac{n_1}{n}\left({\widehat{\mu }}^{\left(0,\hspace{0.17em}-\left\{i,\hspace{0.17em}j\right\}\right)}\left(X_i\right)-{\mu }^{\left(0\right)}\left(X_i\right)\right)\right)}},$$

where the ${\widehat{\mu }}^{\left(w,\hspace{0.17em}-\left\{i,\hspace{0.17em}j\right\}\right)}\left(x\right)$ are predictions obtained without either the ith or the jth training examples. We see that ${\widehat{\mu }}^{\left(w,\hspace{0.17em}-\left\{i,\hspace{0.17em}j\right\}\right)}\left(X_i\right)$ is always computed on $n_w-1$ observations and, moreover, is independent of $W_i$ conditionally on $n_1.$ Thus, we see that $\mathbb{E}\left[R_2\right]=0$ by randomization and moreover that $\mathbb{E}\left[R_2^2\right]=o\left(1/n\right)$ under risk consistency.

To establish our desired result, it remains to show that $\mathbb{E}\left[{\left(R_2-R\right)}^2\right]=o\left(1/n\right).$ In the case where ${\widehat{\mu }}^{\left(0\right)}$ and ${\widehat{\mu }}^{\left(1\right)}$ are computed separately on samples with $W_i=0$ and $W_i=1,$ respectively, we can write

$$R_2-R={\displaystyle \sum _{\left\{i:W_i=0\right\}}{\displaystyle \sum _{\left\{j:W_j=1\right\}}\left(\frac{n_0}{n}\left({\widehat{\mu }}^{\left(1,\hspace{0.17em}-j\right)}\left(X_i\right)-{\widehat{\mu }}^{\left(1\right)}\left(X_i\right)\right)-\frac{n_1}{n}\left({\widehat{\mu }}^{\left(0,\hspace{0.17em}-i\right)}\left(X_j\right)-{\widehat{\mu }}^{\left(0\right)}\left(X_j\right)\right)\right)/\left(n_0{n}_1\right)}}.$$

The jackknife-compatibility condition implies that

$$\mathbb{E}\left[{\left({\widehat{\mu }}^{\left(1,\hspace{0.17em}-j\right)}\left(X_i\right)-{\widehat{\mu }}^{\left(1\right)}\left(X_i\right)\right)}^2|n_1\right]\le \frac{a\left(n_1\right)}{n_1},$$

for any j with $W_j=1,$ etc., and so we find that

$$\frac{1}{2}\mathbb{E}\left[{\left(R_2-R\right)}^2|n_1\right]\le \frac{n_0^2\hspace{0.17em}a\left(n_1\right)}{n^2{n}_1}+\frac{n_1^2\hspace{0.17em}a\left(n_0\right)}{n^2{n}_0},$$

which converges to 0 in probability at a rate faster than $1/n,$ as desired. In the case where ${\widehat{\mu }}^{\left(0\right)}$ and ${\widehat{\mu }}^{\left(1\right)}$ are not computed separately, we attain the same result by applying the jackknife-compatibility condition on the leave-one-out estimators.

B. Additional Simulation Results

We now present additional simulation results for coverage rates of cross-validated cross-estimation, with a focus on settings with treatment heterogeneity and potentially non-Gaussian designs. In an effort to challenge our method, we used signals with very high signal-to-noise ratio. We estimate ${\beta }^{\left(0\right)}$ and ${\beta }^{\left(1\right)}$ jointly, using the procedure described in ref. 22.

In these simulations, we always used an autoregressive covariance structure with ${\mathrm{\Sigma }}_{ij}={\rho }^{\left|i-j\right|},$ where we interpret $0^0=1.$ The design matrices X were generated as $X={\mathrm{\Sigma }}^{1/2}Z,$ where Z was either Gaussian $Z_{ij}\stackrel{\text{i.i.d.}}{\sim }\mathcal{N}\left(0,\hspace{0.17em}1\right)$ or Bernoulli $Z_{ij}\stackrel{\text{i.i.d.}}{\sim }\pm 1$ uniformly at random. The treatment assignment was random with $\mathrm{\mathbb{P}}\left[W_i=1\right]=\mathrm{0.5.}$ Conditionally on $X_i$ and $W_i,$ we generated $Y_i\stackrel{\text{i.i.d.}}{\sim }\mathcal{N}\left(X_i\cdot {\beta }^{\left(W_i\right)},{\sigma }^2\right).$ Finally, we considered three different settings for the signal,

$$\begin{array}{c}\text{Dense:}\hspace{0.17em}{\beta }_j^{\left(0\right)}=p^{-1},\hspace{1em}{\beta }_j^{\left(1\right)}=1.1\hspace{0.17em}{p}^{-1},\\ \text{Geometric:}\hspace{0.17em}{\beta }_j^{\left(0\right)}={10}^{-10j/p},\hspace{1em}{\beta }_j^{\left(1\right)}={11}^{-10j/p},\\ \text{Sparse:}\hspace{0.17em}{\beta }_j^{\left(0\right)}=10\cdot 1\left(\left\{j=1\mathrm{mod}20\right\}\right),\\ {\beta }_j^{\left(1\right)}=9\cdot 1\left(\left\{j=1\mathrm{mod}20\right\}\right)+1\left(\left\{j=1\mathrm{mod}10\right\}\right),\end{array}$$

for $j=1,\hspace{0.17em}\dots ,\hspace{0.17em}p.$ Results are presented in Table S1. Overall, the coverage rates appear quite promising, especially noting the wide variety of simulation settings.