EconBase
← Back to paper

A New Central Limit Theorem for the Augmented IPW Estimator: Variance Inflation, Cross-Fit Covariance and Beyond

Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.

89,372 characters · 9 sections · 78 citation commands

Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.

A new central limit theorem for the augmented IPW estimator: variance inflation, cross-fit covariance and beyond

frontmatter\runtitle{A new CLT for the augmented IPW estimator} \begin{aug} , , \footnotemark[1] \and \footnotemark[1] \address[A]{Department of Statistics, Harvard University\printead[presep={,\ }]{e1,e3,e4}} \address[B]{Department of Biostatistics, Harvard T.H. Chan School of Public Health\printead[presep={,\ }]{e2}} \end{aug} \begin{abstract} Estimation of the average treatment effect (ATE) is a central problem in causal inference. In recent times, inference for the ATE in presence of high-dimensional covariates has been extensively studied. Among diverse approaches that have been proposed, augmented inverse propensity weighting (AIPW) with cross-fitting has emerged a popular choice in practice. In this work, we study this cross-fit AIPW estimator under well-specified outcome regression and propensity score models in a high-dimensional regime where the number of features and samples are both large and comparable. Under assumptions on the covariate distribution, we establish a new central limit theorem for the suitably scaled cross-fit AIPW that applies without any sparsity assumptions on the underlying high-dimensional parameters. Our CLT uncovers two crucial phenomena among others: (i) the AIPW exhibits a substantial variance inflation that can be precisely quantified in terms of the signal-to-noise ratio and other problem parameters, (ii) the asymptotic covariance between the pre-cross-fit estimators is non-negligible even on the $\sqrt{n}$ scale. These findings are strikingly different from their classical counterparts. On the technical front, our work utilizes a novel interplay between three distinct tools—approximate message passing theory, the theory of deterministic equivalents, and the leave-one-out approach. We believe our proof techniques should be useful for analyzing other two-stage estimators in this high-dimensional regime. Finally, we complement our theoretical results with simulations that demonstrate both the finite sample efficacy of our CLT and its robustness to our assumptions. \end{abstract}

Introduction

Causal inference based on observational studies poses a problem of intrinsic interest in the natural and social sciences. Unmeasured confounders pose a major challenge in this regard. We recall that an unmeasured confounder is a variable that affects both the exposure and the outcome of interest, and thus invalidates causal effect estimates based on observational data. Fortunately, with rapid advances in modern data collection technologies, the statistician often hopes to overcome this barrier by collecting data on a large number of potential confounders. While this provides an attractive strategy to mitigate unmeasured confounding, it necessitates causal effect estimation in presence of high-dimensional confounders. This has inspired rapid methodological advances over the past decade at the intersection of statistics, machine learning, computer science, epidemiology etc. on high-dimensional causal inference. This manuscript contributes to this crucial area of research.

Before proceeding further, we describe our problem of interest formally. We observe $n$ i.i.d. observations of a tuple $(y,A,x)$ from some joint distribution $\mathbb{P}$, where $y \in \mathbb{R}$ denotes an outcome of interest, $A \in \{0,1\}$ denotes the binary exposure or treatment and $x \in \mathbb{R}^p$ denotes the measured confounders. We seek to estimate the Average Treatment Effect (ATE), a canonical estimand in this context. The ATE is defined as $\mathbb{E}(y(1)-y(0))$, where $y(a)$ denotes the potential outcome corresponding to $A=a\in \{0,1\}$ pearl2009causality,hernan2010causal,imbens2015causal. Throughout the manuscript, we assume that conditions necessary for identification of the ATE are satisfied, that is, we have (i) no unmeasured confounding, ($y(1),y(0)\perp \!\!\! \!\perp A|x$), (ii) consistency, ($y=Ay(1)+(1-A)y(0)$) and (iii) positivity ($\mathbb{P}(A=1|x)>0$ for all $x\in \mathbb{R}^p$). Under these assumptions, the ATE can be identified from the observed data distribution using $\mathbb{E}(y(1)-y(0))=\mathbb{E}(\mathbb{E}(Y|A=1,x)-\mathbb{E}(Y|A=0,x))$ pearl2009causality,hernan2010causal,imbens2015causal.

Varied approaches exist for ATE estimation athey2019machine and two nuisance functions arise naturally in this context---(i) the Outcome Regression (OR) given by $m(A,x)=\mathbb{E}(y|A,x)$ and (ii) the Propensity Score (PS) given by $\pi(x)=\mathbb{E}(A|x)=\mathbb{P}(A=1|x)$. The regression coefficient vectors underlying the PS and OR models form nuisance parameters for the problem of ATE estimation. Classical approaches include those based on outcome regression robins1986new,hernan2010causal,snowden2011implementation, vansteelandt2011invited, propensity score horvitz1952generalization,rosenbaum1983central,hahn1998role,hirano2003efficient, augmented inverse probability weighting (AIPW) bang2005doubly,scharfstein1999adjusting, to name a few---these utilize suitable modeling assumptions for at least one of the nuisance functions. Recent state-of-the-art methods, including Double Machine Learning chernozhukov2017double, Covariate Balancing imai2014covariate, li2018balancing, athey2018approximate,zubizarreta2015stable,fong2018covariate,ning2020robust, Matching methods rubin1973use,rosenbaum1984reducing,rubin1996matching,stuart2010matching, abadie2011bias, abadie2016matching, and calibration-based procedures sun2021high,tan2020regularized,tan2020model, also estimate at least one of the nuisance functions on way to estimating the ATE. In particular, these approaches assume some structure, e.g., sparsity, in one (or both) of the nuisance parameters in high dimensions.

Among the aforementioned approaches, Double Machine Learning style estimators allow particular flexibility in choosing nuisance functions. To accommodate this flexibility and facilitate theoretical analyses in high dimensions, one additionally employs the idea of cross-fitting chernozhukov2017double,newey2018cross,smucler2019unifying. In this scheme, the statistician initially splits the observed data into (a few) distinct folds. The nuisances are computed based on one fold and the ATE estimate is obtained from an independent fold. The nuisance estimates from the initial step are plugged into the final ATE estimate as appropriate. Subsequently, additional estimators are obtained by permuting roles of the folds. The final cross-fitted estimator is obtained by averaging these distinct estimators. Under high-dimensional sparse models, this strategy allows one to establish consistency and asymptotic normality of the proposed estimator. Remarkably, this approach yields efficient estimators in high dimensions under appropriate sparsity assumptions chernozhukov2017double. (see Section (ref) for a detailed review).

We note that verifying structural assumptions, e.g. sparsity, in the relevant nuisance parameters can be difficult in practice in high dimensions. In addition, the results obtained under such assumptions may suffer from gross inaccuracies when the assumptions are violated. To illustrate, we present Figure (ref). Here we focus on the augmented inverse probability weighting (AIPW) estimator scharfstein1999adjusting,bang2005doubly that exhibits a number of fascinating features, and is arguably the most widely used Double Machine Learning style estimators in practice (see (ref) for a formal definition). In low dimensions, the estimator has the desirable double robustness property, that is, one can estimate the ATE consistently even if one of the OR or PS models is misspecified scharfstein1999adjusting,bang2005doubly. Furthermore, cross-fitted versions of this estimator retain similar robustness properties in ultra-high-dimensions when suitable conditions on signal sparsities are met (see Section (ref) for details).

In Figure (ref), we consider a setting with $n=10000$ i.i.d. samples and $p=700$ covariate dimension. We plot the standard errors of a centered and scaled cross-fit AIPW where the following cross-fitting mechanism is employed: split the data into three equal folds, estimate PS, OR from separate folds, plug into the third fold to estimate the AIPW, switch the role of folds, and average the resulting $3!$ estimators. For this figure, the coordinates of the covariates are drawn i.i.d. from suitably normalized mean zero standard gaussian and the nuisance parameter coordinates (for both the OR and PS models) are drawn i.i.d. from mean-zero gaussian and subsequently considered fixed. We use a linear-logistic model specification for the OR and PS respectively. Theoretical calculations using existing results bang2005doubly,smucler2019unifying,chernozhukov2017double show that the classical estimate of SE in this setting equals 2.06 (the red line). Note these works showed that the theoretical value remains the same in low dimensions, and ultra-high-dimensions under suitable sparsity assumptions. The blue histogram represents the true empirical variability of the estimator. We observe the empirical SE to be much larger, concentrating around 5.5. Thus, as soon as the dimension is moderate compared to the sample size, the cross-fitted AIPW estimator exhibits a massive variance inflation compared to its classical variance, when the underlying signals are not sparse. This means if we use the red line to provide uncertainty quantification, it will lead to gross errors in settings where assumptions from chernozhukov2017double,smucler2019unifying might be violated. Therefore, there is an urgent need for theory and methods that explain Figure (ref), and allow for causal effect estimation with suitable uncertainty quantification in analogous such settings. In this paper, we fill this critical gap in the literature.

figure[figure omitted — 768 chars of source]

We study ATE estimation in the absence of sparsity-type assumptions on nuisance parameters, in an arguably high-dimensional regime. In our subsequent analysis, we assume a linear model for the outcome regression, and a logistic model for the propensity score. We analyze the cross-fitted AIPW estimator in the “proportional asymptotic regime", where the number of observations $n$ and features $p$ both diverge, with the ratio $p/n$ converging to some constant $\kappa>0$. This regime has attracted considerable recent attention in high-dimensional statistics johnstone2001distribution,donoho2009message,bayati2011lasso,el2013robust,bean2013optimal,wang2017bridge,el2018impact,donoho2016high,thrampoulidis2018precise,lei2018asymptotics,sur2019likelihood,bellec2019biasing,candes2020phase,sur2019modern,celentano2020estimation,celentano2020lasso,feng2021unifying,salehi2019impact,bu2019algorithmic,hu2019asymptotics,bellec2021asymptotic,bellec2022observable,xu2019consistent,rad2018scalable,javanmard2013state,bellec2020out,thrampoulidis2015regularized,dobriban2018high,patil2021uniform,yadlowsky2022causal, statistical machine learning and analysis of algorithms mei2018mean,montanari2019generalization,mei2022generalization,deng2019model,liang2020precise,javanmard2020precise,li2021minimum,chandrasekher2021sharp,mignacco2020role, econometrics bekker1994alternative,hahn2002optimal,andrews2005identification,cattaneo2018inference,cattaneo2018alternative,cattaneo2019two,anatolyev2019many etc, and shares roots with probability theory and statistical physics zdeborova2016statistical,montanari2022short. Asymptotic approximations derived under this regime demonstrate commendable performance even under moderate sample sizes (c.f. sur2019modern,liang2020precise as well as the aforementioned references)---this renders the proportional asymptotics regime particularly attractive from a practical perspective.

In our analysis, we trade structural assumptions on model parameters for specific distributional assumptions on observed covariates; intuitively, our setup complements the sparse models studied in the recent literature. Under sparsity, one assumes that either the OR or the PS is governed by relatively few strong features. In contrast, our setting allows both to be potentially influenced by all the covariates, but their individual influences must be of a comparable scale (Section (ref) formalizes this notion). We emphasize that we do not debate the relative merits of these two classes of assumptions. Instead, we seek to provide novel alternate approximations that can be valuable to practitioners in settings where sparsity assumptions from the recent high-dimensional causal inference literature may be violated.

In our framework, consistent estimation of the high-dimensional nuisance parameter vectors (e.g. the regression coefficient vectors for the PS and OR) are impossible in $L_2$ norm. However, low-dimensional functionals such as the ATE can still be estimated at the classical $O_p(1/\sqrt{n})$ rate. The recent work yadlowsky2022causal noted this possibility and compared certain high-level properties of common ATE estimators in the absence of sparsity under proportional asymptotics. However, yadlowsky2022causal focused only on the possibility of $\sqrt{n}$-consistent estimation, without any uncertainty quantification. In this paper, we derive an explicit CLT for the AIPW estimator, and in sharp contrast to the analysis of yadlowsky2022causal, we study the AIPW estimator with cross-fitting. The analysis of the cross-fitted AIPW estimator is relatively straight-forward under sparsity---several pairs of estimators obtained from permuting the splits turn out to be (asymptotically) independent. The averaging operation therefore reduces the variance by a constant factor to gain back the efficiency lost due to sample splitting chernozhukov2017double. In our setting, the behavior shows far more nuances--- these estimators exhibit non-trivial dependencies across the splits that we characterize precisely. To the best of our knowledge, this is the first instance where such non-trivial cross-covariances have been identified. Indeed, we believe this to be one of our main contributions. We hope our analysis will inspire follow-up analyses of similar two-stage estimators under this proportional asymptotics regime.

Throughout this paper, we analyze the 3-split version of the cross-fitted estimator that we used for Figure (ref). To keep things tractable, we consider that the OR is fit using maximum likelihood whereas the PS is fit using either maximum likelihood or its ridge regularized version. We next describe our main contributions in this paper.

Our contributions

Our main contributions are as follows:

enumerate• First, we establish that the cross-fit AIPW estimator converges to a Gaussian limit after centering and $\sqrt{n}$-scaling under the high-dimensional asymptotics $p/n \rightarrow \kappa >0$. Though our assumption on the covariate distribution is stylized, to the best of our knowledge, this is the first CLT for the celebrated AIPW that applies in an arguably high-dimensional regime without any sparsity condition. We hope our analysis will motivate further investigations into properties of other ATE estimators in this regime. • We provide a precise characterization of the asymptotic variance of the appropriately centered and scaled cross-fit AIPW in terms of the problem parameters. Empirically, we observe that this limiting variance is higher than the classical variance. This is indeed expected per prior observations noted in el2018impact,bean2013optimal,el2013robust,donoho2016high,sur2019modern,sur2019likelihood,cattaneo2018alternative,yadlowsky2022causal. However, the exact form of the variance allows one to carefully study effects of (i) the signal-to-noise ratios of the underlying parameters, (ii) the degree of high-dimensionality as quantified by $\kappa$, and (iii) the relations among the underlying parameters, on the asymptotic variance. • Next, cross-fitting leads to intriguing phenomena in our setting. In the existing ultra-high-dimensional literature, certain pairs of estimators obtained by permuting the folds are asymptotically independent on the $\sqrt{n}$ scale, and cross-fitting leads to constant gains in the asymptotic variance---thus yielding an efficient estimator chernozhukov2017double. In sharp contrast, the corresponding pairs of estimators are asymptotically correlated in our setting. We provide an (asymptotically) exact characterization of these cross-covariances as a function of our problem parameters. This once again allows one to study the effects of the parameters on the magnitude of these cross-covariances. In fact, we uncover that in many settings these cross-covariances are, in fact, negative. Complementing earlier works in the literature chernozhukov2017double,newey2018cross,smucler2019unifying, our work thus suggests further benefits of cross-fitting in high dimensions, at least in some scenarios. • On the technical front, we develop our proofs based on the following three distinct techniques: approximate message passing theory, the theory of deterministic equivalents, and the leave-one-out approach. As the reader will see, dealing with the cross-fit estimator and in particular, characterizing the cross-covariances requires a novel conjunction of all of the aforementioned tools. To the best of our knowledge, we have not encountered high-dimensional problems in the literature, broadly speaking, that demand the full strengths of all of these approaches. We expect that the our proof ideas should be useful for studying several other high-dimensional estimators---particularly those involving two-stage procedures that start with nuisance estimation followed by a plug-in step. • To study the practical merits of this work, we complement our results with substantial simulations that demonstrate the finite sample efficacy of our theory. This is perhaps another fascinating feature of the proportional asymptotics regime--- the asymptotic theory based on this regime usually demonstrates remarkable performance even in moderate sample sizes. The recent literature in high-dimensional statistics shows ample evidence in this regard across a variety of problems, and we observe this once again for the AIPW CLT characterized in our work. We also provide extensive comparisons of our work with classical results and demonstrate that we recover classical results when $p/n$ becomes vanishingly small. Finally, our experiments demonstrate that optimizing for predictive accuracy during propensity score estimation via ridge-regularized logistic regression fails to yield optimal downstream variance for the AIPW estimator. This calls for other approaches that would be necessary for choosing the optimal regularization parameter in terms of the AIPW variance.

{ Organization:} The rest of the paper is organized as follows. We describe our precise setting and the recent literature in Section (ref). We present our main result together with empirical studies on its finite sample performance in Section (ref). We complement this via further simulations in Section (ref), where we investigate the effects of cross-fitting in high dimensions and test the robustness of our assumptions. Finally, we discuss key ideas involved in the proof in Section (ref), and finish with a discussion of directions for future research in Section (ref).

{\bf{Notation}}

The results in this paper are mostly asymptotic (in $n$) in nature and thus requires some standard asymptotic notations. If $a_n$ and $b_n$ are two sequences of real numbers then $a_n \gg b_n$ (and $a_n \ll b_n$) implies that ${a_n}/{b_n} \rightarrow \infty$ (and ${a_n}/{b_n} \rightarrow 0$) as $n \rightarrow \infty$, respectively. Similarly $a_n \gtrsim b_n$ (and $a_n \lesssim b_n$) implies that $\liminf_{n \rightarrow \infty} {{a_n}/{b_n}} = C$ for some $C \in (0,\infty]$ (and $\limsup_{n \rightarrow \infty} {{a_n}/{b_n}} =C$ for some $C \in [0,\infty)$). Alternatively, $a_n = o(b_n)$ will also imply $a_n \ll b_n$ and $a_n=O(b_n)$ will imply that $\limsup_{n \rightarrow \infty} \ a_n / b_n = C$ for some $C \in [0,\infty)$). If $C>0$ then we write $a_n=\Theta(b_n)$. If $a_n/b_n\rightarrow 1$, then we say $a_n \sim b_n$.

We use $\stackrel{p}{\to}$ and $\stackrel{d}{\to}$ to denote convergence in probability and distribution respectively. We use $o_p(1)$ to denote sequences of random variables which converge to zero in probability. For any sequences of probability measures $\mu_n$ and another probability measure $\mu$, we say that $\mu_n \stackrel{W_2}{\to} \mu$ if the following holds: there exists a sequence of couplings $\Pi_n$ with marginals $\mu_n$ and $\mu$ respectively, so that if $(X_n,X) \sim \Pi_n$, then $\mathbb{E}[(X_n - X)^2] \to 0$ as $n\to \infty$.

Setup

We study the AIPW estimator using the following working model. Throughout we assume that we observe $n$ i.i.d. samples $\{(y_i, A_i, x_i): 1\leq i \leq n\}$, where the conditional distribution of the treatment given the covariates follows a logistic regression, and the conditional distribution of the outcome given the treatment and the covariates satisfy a linear model. We wish to work in a high-dimensional regime where the covariate dimension is allowed to grow with the sample size. To model this formally, we consider a sequence of problem instances, $\{y_i,A_i,x_i,\epsilon_i^{(0)}, \epsilon_i^{(1)}, 1 \leq i \leq n, \beta(n),\alpha^{(0)}, \alpha^{(1)},\beta^{(0)}(n),\beta^{(1)}(n)\}_{n \geq 1}$ such that

align[align omitted — 181 chars of source]

where $\epsilon_i^{(A_i)} \sim \mathcal{N}\left(0,\left(\sigma^{(A_i)}\right)^2\right)$, independent of everything else. Above, $x_i(n),1\leq i \leq n, \beta(n),\beta^{(0)}(n),\beta^{(1)}(n)$ all lie in $ \mathbb{R}^{p(n)}$ and we allow $p(n), n \rightarrow \infty$ with $p(n)/n \rightarrow \kappa > 0 $. We assume that the covariates satisfy $x_i(n) \sim \mathcal{N}(0,I_p/n)$. Naturally, this is a stylized setting, but we will see that the setting uncovers novel high-dimensional phenomena that should motivate further studies into this regime. We also check robustness to our assumption on the covariate distribution in Section (ref). In the sequel, we drop the dependence on $n$ whenever it is clear from context.

Under the outcome regression model (ref), the population average treatment effect is given by

align[align omitted — 149 chars of source]

We seek to study estimation and inference for $\Delta$, without invoking sparsity type conditions on the propensity score/outcome regression model parameters. This is of course challenging in high dimensions---thus, to keep the problem meaningful we assume that the signal strengths remain finite in the limit, after appropriate scaling. This reduces to requiring that

align[align omitted — 333 chars of source]

for some $\gamma, \sigma_{0 \beta}, \sigma_{1\beta} \in \mathbb{R}^{+}, \rho_{01} \in [-1,1]$. Finally, we require a regularity condition on the structure of the signals given as follows:

align[align omitted — 284 chars of source]

where $W_2$ denotes Wasserstein-2 convergence.

This assumption says that the empirical distributions constructed out of the deterministic sequence of vectors $\{\beta(n) \}_{n \geq 1}, \{\beta^{(0)}(n) \}_{n \geq 1}, \{\beta^{(1)}(n) \}_{n \geq 1}$ converges to a weak limit and the corresponding second moments converge. This is a rather common assumption in the proportional asymptotics regime donoho2009message,bayati2011lasso,javanmard2013state, and intuitively, it ensures that the entries of each of these vectors do not differ wildly from each other. To keep a specific example in mind, the reader may consider a random effects setting, where each entry of the vector $\beta$ is i.i.d., that is, $\beta_i \stackrel{\text{i.i.d}}{\sim} \mu_{\beta}$ and analogously for $\beta^{(1)}_i, \beta^{(0)}_i $. Note that we can allow $\mu_{\beta}$ to contain a spike at $0$, meaning that $\beta$ would then be a sparse vector with sparsity linear in $n$ or $p$. Once again, this is true for $\beta^{(1)}_i, \beta^{(0)}_i $ as well.

We seek to study the cross-fitted AIPW estimator in the aforementioned regime, focusing on the 3-split version:

itemize• Split the data into 3 groups $S_1, S_2, S_3$ with sizes $n_1, n_2, n_3$ respectively such that $$n_{1}+n_{2}+n_{3}=n, \quad \lim _{n \rightarrow \infty} \frac{n_{i}}{n}=r_{i} \in(0,1), \quad \lim _{n \rightarrow \infty} \frac{p}{n_{i}} =\kappa_i > 0 \quad \forall i=1,2,3.$$ • Let $ (a,b,c) $ be a permutation of $(1,2,3)$. \begin{enumerate} • Use $S_{a}$ to obtain an estimate for $\beta$. Here we consider either the logistic MLE or its ridge regularized counterpart. We denote these using $\hat{\beta}_{S_a}$ or $\hat{\beta}_{S_a}^{(\lambda)}$ respectively. Note that $\hat{\beta}_{S_a}^{(\lambda)}$ is obtained by solving the following strongly convex minimization problem $$ \hat{\beta}_{S_a}^{(\lambda)} =\text{argmin}_{b \in \mathbb{R}^{p}} \sum_{i \in S_a}\left\{\log \left(1 + e^{{x}_{i}^{\top} {b}}\right)-A_{i}\left({x}_{i}^{\top} {b}\right)\right\} + \frac{\lambda}{2} \|b\|^2 .$$ • Use $S_b$ to estimate ${\alpha}^{(0)}, {\alpha}^{(1)}, {\beta}^{(0)}, {\beta}^{(1)}$. In particular, we consider the least squares estimators \begin{align} (\hat{\alpha}^{(0)}, \hat{\beta}^{(0)}) = \mathrm{argmin}_{(\alpha, \beta)} \sum_{i \in S_b} (1-A_i) (y_i - \alpha - x_i^{\top} \beta)^2 , \nonumber \\ (\hat{\alpha}^{(1)}, \hat{\beta}^{(1)}) = \mathrm{argmin}_{(\alpha, \beta)} \sum_{i \in S_b} A_i (y_i - \alpha - x_i^{\top} \beta)^2. \end{align} • Use $S_c$ to obtain the final estimator \begin{equation} \hat{\Delta}_{AIPW}=\hat{\Delta}_{AIPW,{1}}-\hat{\Delta}_{AIPW,0} \end{equation} for the ATE, where $$ \hat{\Delta}_{AIPW,{1}} = \frac{1}{n_c} \sum_{i \in S_c} \left \{ \frac{A_iy_i}{\sigma \left( x_i^{\top} \hat{\beta}_{S_a} \right )} - \frac{A_i - \sigma\left( x_i^{\top} \hat{\beta}_{S_a} \right)}{\sigma\left( x_i^{\top} \hat{\beta}_{S_a}\right)} \left(\hat{\alpha}^{(1)}_{S_b} +x_i^{\top}\hat{\beta}^{(1)}_{S_b}\right)\right \} ,$$ $$ \hat{\Delta}_{AIPW,{0}} = \frac{1}{n_c} \sum_{i \in S_c} \left \{ \frac{\left(1-A_i\right)y_i}{1-\sigma\left( x_i^{\top} \hat{\beta}_{S_a}\right)} + \frac{A_i - \sigma\left( x_i^{\top} \hat{\beta}_{S_a}\right)}{1-\sigma\left( x_i^{\top} \hat{\beta}_{S_a}\right)} \left (\hat{\alpha}_{S_b}^{(0)} +x_i^{\top}\hat{\beta}^{(0)}_{S_b}\right )\right \} .$$ \end{enumerate} • For each permutation of $(1,2,3)$, we obtain an estimator $\hat{\Delta}_{AIPW}$. The final estimator of the population treatment effect is obtained by averaging all such estimators. We denote the cross-fitted estimator as $\hat{\Delta}_{cf}$.

Note that we use OLS estimators for $(\alpha^{(0)}, \beta^{(0)}, \alpha^{(1)}, \beta^{(1)})$, so we need to restrict to a regime where these are unique. Of course this is not guaranteed, especially when the feature dimension $p$ is reasonably large compared to sample size $n$. In Theorem (ref) below, we derive an explicit characterization of the regime where unique OLS estimators exist with high probability for our aforementioned problem. The theorem shows that it suffices to have $\kappa_i < 1/2$ for $i=1,2,3$. We will implicitly restrict ourselves to this region in the rest of the paper. Similarly, when we use the logistic MLE we will restrict to a regime where it exists w.h.p. We will clarify this further in Section (ref).

\bf Background

In this section, we review strategies for ATE estimation, focusing primarily on the recent literature on ATE estimation with high-dimensional covariates. Along the way, we describe some key ideas facilitating these recent methodological breakthroughs, and contrast them with our approach.

{\bf ATE estimation in low dimensions:} In the classical setting ($p$-fixed, $n \to \infty$), the ATE can be estimated at the $\sqrt{n}$ rate, and asymptotically normal semi-parametric efficient estimators are well-known. In this context, AIPW estimators are particularly attractive scharfstein1999adjusting,bang2005doubly,van2006targeted. These estimators were originally introduced for mean estimation in missing data problems robins1994estimation,robins1995analysis,robins1995semiparametric,scharfstein1999adjusting, before being used for causal effect estimation. The interest in these estimators stems from the well-known "Double Robustness" (DR) property. Formally, AIPW estimators facilitate consistent estimation of the ATE even if one of the PS or OR is misspecified. Additionally, such estimators are also asymptotically gaussian under potential model misspecifications described above scharfstein1999adjusting,bang2005doubly,van2006targeted, and thus facilitates robust inference of the ATE. Indeed, this attractive combination of properties has established AIPW estimators as a trusted tool for causal effect estimation in the modern statistician's toolkit.

{\bf ATE estimation in high dimensions:} We now turn to the extensive recent advances in causal effect estimation in high dimensions (i.e. both $n,p \to \infty$). Ideally, one still wishes to design estimators that enable consistent and asymptotically normal (CAN) inference for the ATE under misspecification of either the PS or OR model. Unfortunately, this presents challenges in high dimensions, and such estimators are usually available under strong structural assumptions on the PS and/or OR models. Over the past decade, the scope of allowed model misspecifications expanded significantly and at the same time, structural constraints imposed on the “well-specified" part of the model reduced steadily. Such remarkable progress occurred due to a number of creative methodological ideas such as penalized regression followed by de-biasing, sample splitting and cross-fitting etc. In the subsequent discussion, we will touch upon some of these key ideas, and discuss why they fail to apply in our setting.

First, we review the state-of-the-art in terms of allowed model misspecification, and survey the modern causal effect estimators that enjoy these robustness guarantees (along the way, we will indicate the structural assumptions imposed on the well-specified part of the model by these respective strategies). In terms of tolerated model misspecifications, two recent notions have gained prominence: (i) rate double robustness---here one assumes that both the PS and OR models have approximately sparse expansions, and establishes that CAN estimation is possible as long as the product of the underlying sparsity parameters is sufficiently small, (ii) model double robustness---here one allows one of the PS or OR model to be misspecified, as long as the other well-specified nuisance component is sufficiently sparse.

Rate Double Robustness: In the context of rate double robustness, belloni2014inference,farrell2015robust,chernozhukov2017double,chernozhukov2018biased,smucler2019unifying,chernozhukov2021automatic employ somewhat parallel strategies where one first estimates the nuisance functions and thereby requires the product of their errors (in root mean squared error) in estimating the true functions to be $o_{p}(1)$. Translating to exact sparsity classes, since one can typically estimate OR and PS at a rate $\sqrt{s_{m}\log{p}/n}$ and $\sqrt{s_{\pi}\log{p}/n}$ (see e.g. buhlmann2011statistics) respectively (where $s_{\pi}$ is the sparsity of $\pi(\mathbf{x})$ and $s_{m}$ is the maximum sparsity of $m(1,\mathbf{x})$ and $m(0,\mathbf{x})$ respectively), one obtains a requirement of $s_{m}\vee s_{\pi}\ll \sqrt{n}/\log{p}$ for CAN estimation of ATE. More carefully constructed estimators have obtained sharper results through various approaches that lower the requirement on the sparsities of $s_m$ and $s_{\pi}$. For instance, bradic2019sparsity constructs an estimator that requires either $(s_{\pi}\ll n/\log{p}, s_m\ll \sqrt{n}/\log{p})$ or $(s_{\pi}\ll \sqrt{n}/\log{p}, s_m\ll n^{3/4}/\log{p})$.

Model Double Robustness: We now turn to the model double robustness literature. In this regard, (a) athey2018approximate bypasses correct specification on PS by exploiting the structure of the bias in estimation of the sparse OR (which is required to satisfy $s_{m}\ll \sqrt{n}/\log{p}$); (b) wang2020debiased bypasses correct specification of OR by correcting the bias in estimation of the sparse PS (which is required to satisfy $s_{\pi}\ll \sqrt{n}/\log{p}$); (d) tan2020model constructs estimators of ATE based on calibrated OR and PS estimation. This allows valid CAN inference on ATE when the PS model is correctly specified and the OR model is misspecified (under a linear representation in a feature space), but the product of sparsities of the PS and the limit of the OR estimator is smaller than $n/\log^2{p}$; (e) ning2020robust employs a covariate balancing technique to allow for similar results to tan2020model but also provides asymptotic normality of their estimator at a rate slower than $\sqrt{n}$ when the PS model is misspecified; and (f) smucler2019unifying provides a unified view of construction of rate and model doubly robust estimators of quantities similar in essence to ATE, using ideas from semiparametric theory.

{\bf Key Methodological Ingredients and Principles:} The impressive advances surveyed above rest on a few key insights. First, the aforementioned estimators allow $\sqrt{n}$-consistent, asymptotically normal estimation of the ATE, as long as at least one of the PS or OR models is consistently estimable belloni2014inference,athey2018approximate,tan2020model,tan2020regularized,bradic2019sparsity,wang2020debiased,smucler2019unifying,chernozhukov2017double,farrell2015robust in $L_2$ norm. Furthermore, while constructing CAN estimators using Neyman orthogonalization, an approach that encompasses AIPW-type estimators, one first establishes an asymptotic expansion farrell2015robust,chernozhukov2017double,smucler2019unifying under suitable regularity conditions (e.g. sparsity). This expansion implies a limiting gaussian distribution for the estimator prior to cross fitting. Finally, one establishes that the individual estimators obtained from the permutation of the splits are asymptotically independent on the $\sqrt{n}$ scale, and thus a CLT for the cross-fit estimator follows immediately (see e.g. chernozhukov2017double,kennedy2022semiparametric).

{\bf Key distinctions in our setting:} It is particularly instructive to evaluate the utility of the aforementioned ideas in our context. First and foremost, consistent estimation of the PS and OR models in $L_2$ norm is impossible in our framework \textcolor{blue}{mourtada2019exact,sur2019modern,donoho2016high}. This immediately invalidates the technical ingredients underlying the prior methods. Moreover, the aforementioned expansion of the AIPW estimator fails to hold in our case. Finally, as mentioned previously, the estimators obtained from permuting different splits are asymptotically dependent in our setting. This crucially affects our analysis, and necessitates a radically different approach. We emphasize that although we assume well-specified PS and OR models, CAN estimation of the ATE is known to be challenging even under these additional simplifications bang2005doubly,wager2016high,ning2020robust,tan2020model.

Main Results

Recall from Section (ref) that we use OLS for fitting the outcome regression model. As a first step, we characterize the sample size regimes that ensure the existence of these least squares estimators with high probability.

theoremFor any $(a,b,c)$ permutation of $(1,2,3)$, the estimates $(\hat{\alpha}^{(0)}, \hat{\beta}^{(0)})$, $(\hat{\alpha}^{(1)}, \hat{\beta}^{(1)})$ are unique with high probability if and only if $\kappa_b<1/2$.

When we use maximum likelihood for the propensity score estimation, we need to ensure that this exists in our setting. The precise asymptotic threshold for the existence of the logistic MLE has been recently characterized in candes2020phase. Specifically, candes2020phase provides an explicit formula for a function $h(\cdot)$ such that when $\kappa = \lim p/n < h(\gamma^2)$ (resp. $\kappa > h(\gamma^2)$), the logistic MLE exists (resp. does not exist) with high probability. Combining these two requirements, we introduce the notion of a feasible tuple that refers to any combination of problem parameters for which both the OLS for the outcome regression model and the MLE for the propensity score model exist w.h.p.

defn[Feasible] We call a tuple $(n,p,r_1, r_2, r_3,\beta,\alpha^{(0)},\beta^{(0)}, \alpha^{(1)}, \beta^{(1)})$ to be feasible if \begin{itemize} • The logistic regression MLE estimates $\hat{\beta}_{S_1}$, $\hat{\beta}_{S_2}$, $\hat{\beta}_{S_3}$ exist with probability converging to $1$, and • The OLS estimates $\{(\hat{\alpha}^{(0)}_{S_i}, \hat{\beta}^{(0)}_{S_i}): i =1,2,3\}$ and $\{(\hat{\alpha}^{(1)}_{S_i}, \hat{\beta}^{(1)}_{S_i}): i =1,2,3\}$ exist with probability converging to $1$. \end{itemize}

We now introduce the first of our two main results that establishes the asymptotic distribution of the cross-fit AIPW estimator for every feasible tuple, when the propensity score model is fit using maximum likelihood.

theoremAssume that the tuple $(n,p,r_1, r_2, r_3,\beta,\alpha^{(0)},\beta^{(0)}, \alpha^{(1)}, \beta^{(1)})$ is feasible and that the logistic MLE is used for propensity score estimation. Under the conditions specified in Section (ref), as $p,n \to \infty$ with $p/n \rightarrow \kappa > 0$, \begin{align} & \sqrt{n}(\hat{\Delta}_{cf} - \Delta) \stackrel{d}{\to} \mathcal{N}(0, \sigma_{cf}^2), \quad with \\ & \sigma_{cf}^2 =\left[ \left(\sigma^{(0)}\right)^2 + \left(\sigma^{(1)}\right)^2 \right] f(\kappa,\gamma^2)+ \frac{\kappa}{9} \left( \sigma_{0\beta}^2 + \sigma_{1\beta}^2 - 2 \rho_{01} \sigma_{0 \beta} \sigma_{1 \beta} \right) \Big( \frac{1}{r_{1}} + \frac{1}{r_2} + \frac{1}{r_3} \Big).\nonumber \end{align}

The effect of fitting the propensity score and the noise level in the observed outcomes appear in the first summand in the variance, while the second summand concerns the signal strengths underlying the two outcome regression models. The function $f(\cdot)$ takes a complicated form so we defer its details to Appendix (ref) (Eqn. (ref)).

Our new formula (ref) warrants an immediate comparison with its classical counterpart. To this end, we consider a simplified setting where $r_i = 1/3$ and $\sigma^{(0)}=\sigma^{(1)}=\sigma_{\varepsilon}$. If the dimension were fixed, the classical asymptotic (in large sample limit) variance for the AIPW bang2005doubly in this case reduces to

align[align omitted — 189 chars of source]

The ultra-high-dimensional settings in chernozhukov2017double,smucler2019unifying also admit the same variance form, apart from an additional limit (in $p$) on the RHS to account for the divergence of $p$. Here, we restrict our discussion to the fixed $p$ case for simplicity. Note that the second term in $\sigma^2_{\text{cf}}$ (Eq. (ref)) is the limit, under our regime, of $ \text{Var}\{x_{i}^{\top}(\beta^{(1)}-\beta^{(0)}) \}$, the second term in the classical formula (ref). Thus, the differences induced by our high-dimensional regime manifests through differences between $f(\kappa,\gamma^2)$ from (ref) and $\mathbb{E}[1/\sigma(x_i^{\top}\beta)]$ from (ref). To visualize this difference, we plot the ratio $\log(f(\kappa,\gamma^2)/\mathbb{E}[1/\sigma(x_i^{\top}\beta)])$ as a function of $p/n$, for a few choices of $\gamma$ in Figure (ref). Note that the ratio tends to zero as $p/n$ approaches zero, indicating that our variance formula recovers the classical formula when the dimensionality decreases. Whereas the ratio deviates further from 1 as $p/n$ grows larger. We investigate our formula for $f(\kappa,\gamma^2)$ further and formally show in Appendix (ref) that $f(\kappa,\gamma^2)$ reduces to $\mathbb{E}[1/\sigma(x_i^{\top}\beta)]$ in the classical regime (fixed $p$, large $n$). We further plot the ratio between the total variance in our regime versus the classical regime in Figure (ref), and observe similar trends.

figure[figure omitted — 1,231 chars of source]

Note that Theorem (ref) uses maximum likelihood for both the OR and PS models, thereby restricting the parameter range where the Theorem applies. To overcome this restriction, we next establish an analogous CLT where the propensity scores are estimated via ridge regularized logistic regression.

theoremFix any $\lambda \in \mathbb{R}^{+}$. Assume that the OLS estimates $\{(\hat{\alpha}^{(0)}_{S_i}, \hat{\beta}^{(0)}_{S_i}): i =1,2,3\}$ and $\{(\hat{\alpha}^{(1)}_{S_i}, \hat{\beta}^{(1)}_{S_i}): i =1,2,3\}$ exist with probability converging to $1$, that is, $\kappa_i < 1/2$ for all $i$. Under the conditions specified in Section (ref), as $p,n \to \infty$ with $p/n \rightarrow \kappa > 0$, \begin{align} \sqrt{n}(\hat{\Delta}_{cf}^{(\lambda)} - \Delta) \stackrel{d}{\to} \mathcal{N}\left(0, \left(\sigma_{cf}^{(\lambda)}\right)^2\right),\nonumber \end{align} where $$\left(\sigma_{cf}^{(\lambda)}\right)^2 =\left[ \left(\sigma^{(0)}\right)^2 + \left(\sigma^{(1)}\right)^2 \right] f^{(\lambda)}(\kappa,\gamma^2)+ \frac{\kappa}{9} \left( \sigma_{0\beta}^2 + \sigma_{1\beta}^2 - 2 \rho_{01} \sigma_{0 \beta} \sigma_{1 \beta} \right) \Big( \frac{1}{r_{1}} + \frac{1}{r_2} + \frac{1}{r_3} \Big).$$

Once again, $f^{(\lambda)}(\cdot)$ takes a complicated form so we defer its details to Appendix (ref). Note the limiting variance has a similar structure as in Theorem (ref). On examining the Appendix one would observe that $f^{(\lambda)}(\kappa,\gamma^2)$ equals $f(\kappa,\gamma^2)$ when $\lambda=0$, as we would expect.

We next study the finite sample efficacy of our result. Through the rest of this section and the subsequent section, we set $n = 10,000, n_1=3,333, n_2 = 3,333, n_3 = 3,334, p = 700$ so that the “dimensionalities” $p/n_1, p / n_2, p/n_3$ are approximately 0.21. The matrix of covariates has i.i.d. $ \mathcal{N}\left(0, 1/n\right)$ entries unless otherwise specified, and the regression coefficients $\beta, \beta^{(1)}, \beta^{(0)}$ are drawn from normal distributions with zero mean and scaled in such that $ \gamma=0.1$ and $ \sigma_{0\beta},\sigma_{1\beta},\rho_{01}$ remain the same as in Figure (ref).

In the aforementioned setting, Figure (ref) shows two overlaid normal Q-Q plots of $\sqrt{n}\left(\hat{\Delta}_{c f}-\Delta\right)$. In both cases, we compute the sample quantiles from 30,000 simulation runs. The darker blue points represent the theoretical quantiles based on our theory, when the logistic MLE is used for propensity score estimation, while the lighter cyan points represent those computed based on the classical theory. Observe that our theory captures the true sample quantiles accurately. The plot exhibits some deviation from the reference line near the tails. This occurs due to the presence of $\sigma(\cdot)$ and $1-\sigma(\cdot)$ in the denominator of the AIPW estimator. It is expected that if either of these terms is extremely small, this would manifest as outliers in the QQ-plot. To alleviate this issue, we winsorize the sigmoid function to satisfy $0.005 \leq \sigma(\cdot) \leq 0.995$. This winsorizing step is commonly used in the implementation of the AIPW estimator. Figure (ref) demonstrates that after winsorizing, our theoretical variance matches the empirical value exceptionally well. We discuss the possibilities of rigorously quantifying an analogous CLT for the winsorized estimator in Section (ref). In Section (ref), we further study the effects of regularized estimation of the propensity scores (Theorem (ref)).

figure[figure omitted — 655 chars of source]

Numerical Experiments

This section pursues important additional empirical investigations. First, we present a crucial phenomenon that can be studied as an upshot of our theory. Next, we study the effects of cross-validation, and finally, we test the robustness of our results to the covariate distribution assumptions.

Effects of Cross-fitting in high dimensions? The existing literature on cross-fit AIPW tells us the following important fact: at the $\sqrt{n}$-scale, the covariances between certain pairs of pre-cross-fit estimators are asymptotically negligible. Thus, averaging the pre-cross-fit estimators leads to a variance reduction. In our setting, we observe that these cross-covariances admit non-trivial limits, and our proof for Theorem (ref) precisely characterizes the asymptotic values of these cross-covariances.

To describe further, denote $\hat{\Delta}_{\text{pre$_$fit}}(S_a, S_b, S_c)$ to be the pre-cross-fit AIPW estimator where the PS is estimated using $S_a$, the OR is estimated using $S_b$, and the AIPW is calculated on $S_c$, plugging in the preceding nuisance estimates. Suppose we group the $3 !$ pre-cross-fit AIPW estimators into 3 pairs, where each pair consists of two estimators of the form $\Big(\hat{\Delta}_{\text{pre$_$fit}}(S_a, S_b, S_c), \hat{\Delta}_{\text{pre$_$fit}}(S_b, S_a, S_c)\Big)$. With this grouping, we may split our asymptotic variance $\sigma^2_{\text{cf}}$ into the following parts:

align[align omitted — 683 chars of source]

where the second covariance term captures sum of the total covariance within each pair, and the sum of the last two terms capture the overall between-pair covariances. On examining each term in the decomposition (ref)--(ref), we observe that both (ref) and (ref) contribute in our setting and in the classical low-dimensional setting. But, their magnitude is higher in our regime due to high-dimensional effects. In fact, if we were to plot ratios of these terms under the two regimes, we would once again observe trends similar to those reported in Figure (ref). Thus, we refrain from investigating these further and instead turn to the between-pair covariance, that is, sum of (ref) and (ref).

In the classical regime, the total between-pair covariance is negligible at the $\sqrt{n}$-scale. However, these contribute non-trivially in our regime even in the large sample and large dimensional limit. The reader should view this phenomenon as an additional effect of cross-fitting in high dimensions. When we fit propensity scores using maximum likelihood, we observe that the total between-pair covariance is negative, as demonstrated via Figure (ref). This illustrates that cross-fitting helps in high dimensions in such settings, in addition to its usual advantages discussed in chernozhukov2017double,newey2018cross. However, on using ridge regression for estimating the propensity score, we observe that this between-pair covariance could be positive in some cases. Thus, one needs to investigate this phenomenon further to characterize the interplay between the problem parameters, e.g. signal strength, tuning parameter, etc. that determines regimes where the between-pair covariance is negative in our high-dimensional setting. We defer these additional investigations to future work. To our knowledge, our work uncovers such non-trivial between-pair covariances for the first time in the literature on high-dimensional causal inference and cross-fitting.

figure[figure omitted — 810 chars of source]

Does Cross-validation Find the Optimal Regularization Parameter? In this paper, we allow regularized estimation of the propensity score model via ridge penalized logistic regression. This naturally requires suitable choice of the tuning parameter. In traditional supervised learning, one seeks to tune the regularization parameters to optimize the out-of-sample prediction accuracy. In this context, it is well-known that tuning parameter selection approaches such as k-fold cross-validation (CV) suffer from large biases in high dimensions (c.f rad2020scalable), whereas leave-one-out cross validation (LOOCV) exhibits desirable properties patil2021uniform.

Note that in our setting, the tuning parameter should not be selected to optimize prediction accuracy on a test point, but rather to minimize the variance of the downstream AIPW estimator. However, traditional CV based approaches are still widely utilized in this setting. Here, we explore the impact of this choice on the ATE estimation task. Formally, we study the effects of using LOOCV for choosing the tuning parameter for the propensity score model on the downstream performance of the AIPW estimator. Note that LOOCV is computationally expensive, so we work with an approximation obtained as follows. For any given $\lambda$, LOOCV involves computing all $n$ possible leave-one-out estimates $\{\hat{\beta}^{(\lambda)}_{S_a}\}^{(-i)}$. Now, Sur14516 relates such leave-one-out estimates to the original estimator, when one uses the logistic MLE. Using the exact same computation, an analogous expression can be derived for the ridge regularized problem. This connects $\{\hat{\beta}^{(\lambda)}_{S_a}\}^{(-i)}$'s to the original ridge estimate $\{\hat{\beta}^{(\lambda)}_{S_a}\}$. Utilizing this formula, one can bypass the computational overload induced by the leave-one-out operation and obtain an approximation that is asymptotically equivalent to LOOCV (rad2020scalable,wang2018approximate studies such approximations for a variety of problems). We implement this approximate LOOCV in Figure (ref)---the dotted red line shows the standard deviation of the AIPW estimator corresponding to the tuning parameter chosen via this approximated LOOCV. The solid blue line shows the variation in the standard deviation as a function of the tuning parameter. The optimal tuning parameter (in terms of the standard deviation) reduces the variance significantly compared to the MLE, as one would expect. However, the LOOCV tuned estimator is highly sub-optimal. This clearly illustrates that optimizing the propensity score fit for predictive accuracy at the first stage does not guarantee optimal sampling variance downstream.

On the other hand, if one can develop consistent estimators for the signal strength parameters $\gamma, \sigma_{0\beta}, \sigma_{1\beta}$ and $\rho_{01}$, our theory provides an alternate route to select tuning parameters (thereby minimizing the downstream variance). We defer further discussions on the possibility of developing such estimators to Section (ref).

figure[figure omitted — 592 chars of source]

Robustness to Normality Assumptions? To conclude our empirical investigations, we test the validity of our theory under non-Gaussian covariate distributions. We consider two settings: (i) a simple Uniform distribution and (ii) a discrete distribution inspired by applications in statistical genetics. Figure (ref) shows two overlaid normal Q-Q plots of $\sqrt{n}\left(\hat{\Delta}_{c f}-\Delta\right)$, where the matrix of covariates has i.i.d. Uniform$\left(-\sqrt{\frac{3}{n}}, \sqrt{\frac{3}{n}}\right)$ entries. Observe that although our theory fails to cover this setting for the time being, the theoretical predictions match the empirical behavior of the cross-fit AIPW remarkably well. To test the validity of our theory further, Figure (ref) considers a design matrix where the $j$th feature takes values in \{0,1,2\} with probabilities $p_{j}^{2}, 2 p_{j}\left(1-p_{j}\right),\left(1-p_{j}\right)^{2}$; here, $p_j \in [0.25, 0.75]$ and $ p_{j} \neq p_{k} \text { for } j \neq k$. Features are then centered and rescaled to have unit variance. The setting is otherwise the same as for Figure (ref). The left plot depicts the quantiles for the cross-fit AIPW and compares with our theory. We see that the suitably scaled and centered cross-fit estimator still follows an approximate normal distribution whose variance can be characterized by our results far better than the classical variance. This time, we do observe deviations from our theory---this is indeed expected since several of the estimated propensity scores are either too small or too large for this particular setting. This prompts us to consider the winsorized version of the estimator, where as before, $0.005 < \sigma(\cdot) < 0.995$. The right plot shows the winsorized cross-fit AIPW, and once again we observe the empirical quantiles match those based on our CLT extremely well. This set of simulations raises an interesting question: can one characterize the class of covariate distributions under which our same CLT applies? In light of our current fairly involved proofs, we defer theoretical investigations in this direction to future work.

figure[figure omitted — 381 chars of source]
figure[figure omitted — 700 chars of source]

Proof Outline

In this section, we collect some ideas involved in the proof of Theorem (ref), and discuss the main technical ingredients. To this end, we first introduce some notation.

For $j \in \{1,2,3\}$, let $\mathcal{E}_{S_j}$ denote the vectors containing $ A_i\epsilon_{i}^{\left(1\right)} +(1-A_i)\epsilon_{i}^{\left(0\right)} $ for all $i \in S_j$.

For $j,k \in \{1,2,3\}$, define

align[align omitted — 1,128 chars of source]

Furthermore, let $V_{S_j,S_k},\tilde{ V}_{S_{j}, S_{k}} \in \mathbb{R}^{|S_j|}$ denote the vectors containing $ \frac{1}{\sqrt{n_j}} \frac{A_i}{ \sigma(x_i^{\top} \hat{\beta}_{S_k}) }$ and $ \frac{1}{\sqrt{n_{j}}} \frac{1-A_{i}}{1-\sigma\left(x_{i}^{\top} \hat{\beta}_{S_k}\right)}$ for all $i \in S_j$ respectively. We establish the following representation for the cross-fitted AIPW estimator in Lemma (ref).

align[align omitted — 531 chars of source]

In the representation above, $\mathscr{S}_3$ denotes the set of all permutations of $\{1,2,3\}$, and we use $(a,b,c)$ to denote the permutations in this set.

The representation (ref) is critical for analyzing the limiting distribution of the AIPW estimator. Note that conditioned on everything but the $\mathcal{E}$ variables, $T_1$ has a mean-zero gaussian distribution. On the other hand, $T_2$ has a mean-zero gaussian distribution as it is a linear function of the $\{x_i : 1\leq i \leq n\}$ variables. As the $\{x_i : 1\leq i \leq n\}$ and $\mathcal{E}$ variables are independent, it is not hard to see that $T_1$ and $T_2$ are asymptotically independent. As both $T_1$ and $T_2$ are mean zero gaussian and asymptotically independent, the limiting gaussian distribution of the AIPW follows immediately, once we establish that the limiting variance of $T_1$ and $T_2$ converge to well-defined constants.

The limiting variance of $T_2$ is explicit, and its convergence follows directly from our assumptions (ref) and (ref). The variance of $T_1$ is significantly more involved---we establish that $\mathrm{Var}(T_1|A,X)$ converges to a deterministic constant in the limit $n \to \infty$. This is our main theoretical contribution, and requires the bulk of the technical work in this paper.

To characterize the limit of $\mathrm{Var}(T_1|A,X)$, we carefully combine several distinct ingredients. We take this opportunity to briefly describe each tool, and motivate its usefulness in our setting. We believe these ideas can be useful for analyzing other estimators in high-dimensions, and should be of independent interest.

{\bf Approximate Message Passing and state evolution:} Approximate Message Passing (AMP) algorithms were introduced in the study of mean-field spin glasses and in compressed sensing donoho2009message,bolthausen2014iterative. In high-dimensional statistics, these algorithms provide a valuable theoretical device---they can be used to “track" the performance of specific statistical estimators e.g. the LASSO, M-estimators, the MLE etc. At a high-level, an AMP algorithm introduces an iterative system $\{\hat{\beta}^{t}: t \geq 1\}$ which “converges" to the estimator of interest $\hat{\beta}$---formally,

align[align omitted — 147 chars of source]

AMP algorithms are attractive theoretical devices in high-dimensional statistics, as their empirical distributions can be tracked using low-dimensional scalar recursions, referred to as “state-evolution". In particular, for well-behaved functions $f:\mathbb{R}^2\to \mathbb{R}$ and any $t \geq 1$, one obtains explicit expressions for the limits of empirical averages $\frac{1}{p} \sum_{i=1}^{p} f(\beta_i, \hat{\beta}_i^t)$ as $p \to \infty$. Here, $\beta \in \mathbb{R}^p$ refers to the underlying latent parameter of interest. Subsequently, using the AMP convergence property (ref) and setting $t \to \infty$, one obtains a precise characterization of the empirical distribution of the estimator $\hat{\beta}$. Specifically, this characterizes $ \frac{1}{p} \sum_{i=1}^{p} f(\beta_i, \hat{\beta}_i)$ in the limit $p \to \infty$. We do not provide a more formal discussion of AMP style algorithms and their consequences in this paper, but refer the interested reader to montanari2012graphical,feng2021unifying for an in-depth exposition of these ideas. We note in passing that similar characterizations of empirical averages can also be obtained using the parallel approach based on Gaussian comparison inequalities stojnic2013framework,thrampoulidis2015gaussian,thrampoulidis2018precise.

Instead, we turn to the importance of these ideas in our analysis. In our analysis of the conditional variance $\mathrm{Var}(T_1| A,X)$, we naturally have to deal with averages of the form

align[align omitted — 121 chars of source]

where $\hat{\beta}_{S_1}$ denotes the MLE estimate for the propensity score model based on the sample split $S_1$. Note that as $S_1$ and $S_3$ are disjoint, conditioned on the samples in $S_1$, the empirical average above is an i.i.d. average, with $(x_i^{\top} \beta, x_i^{\top} {\hat{\beta}_{S_1}})$ are bivariate gaussian with mean zero, $\mathrm{Var}(x_i^{\top} \beta|S_1) = \frac{1}{p} \|\beta\|^2$, $\mathrm{Var}(x_i^{\top} \hat{\beta}_{S_1}|S_1) = \frac{1}{p} \|\hat{\beta}_{S_1}\|^2$, and $\mathrm{cov}(x_i^{\top}\beta, x_i^{\top} \hat{\beta}_{S_1}) = \frac{1}{p} \beta^{\top} \hat{\beta}_{S_1}$. Thus conditioned on the samples in $S_1$, as $n_3 \to \infty$,

align[align omitted — 139 chars of source]

where $(Z_1,Z_2)$ is a mean zero bivariate normal with the covariance matrix described above. Observe that if one could establish that the (random) covariance matrix of $(Z_1,Z_2)$ stabilizes to a deterministic limit as $n_1 \to \infty$, it immediately follows that

align[align omitted — 135 chars of source]

where $(Z_1,Z_2)$ is a mean zero bivariate gaussian with the limiting covariance matrix. This is precisely the step where the state-evolution characterization of the MLE is invaluable. Indeed, note that both $\frac{1}{p} \| \hat{\beta}_{S_1} \|^2$ and $\frac{1}{p} \beta^{\top} \hat{\beta}_{S_1}$ are empirical averages of the form described above, and thus have well-defined, explicit, deterministic limits specified by the state-evolution description. This idea is used repeatedly in our proof to characterize the (deterministic) limits of several averages of the form (ref).

{\bf Deterministic Equivalents:} In classical random matrix theory, the limiting spectral distribution of a random matrix $M \in \mathbb{R}^{n\times n}$ is an object of central interest. The limiting spectral measure of classical random matrix ensembles such as the Wigner and the Wishart ensembles have been characterized using a number of different approaches e.g., the moment method and the method of Stieljes transforms. However, these approaches have some shortcomings---first, they are typically tractable only for very symmetric random matrix models, and second, these approaches do not shed any light on the eigenvectors. Consequently, understanding the eigenvectors often requires significant additional work.

The theory of deterministic equivalents was inspired by applications in signal processing and wireless communications hachem2007deterministic,couillet2011deterministic, but its origins can be traced to the early works of girko2012theory. Intuitively, given a random matrix, this non-asymptotic theory identifies a deterministic surrogate which has the same eigenvalue and eigenvector properties. Crucially, this yields rich spectral information about the random matrix of interest at finite problem sizes, without the restriction that these properties converge in the limit. We use the following formal definition of deterministic equivalents in this paper phdthesis.

defn[Deterministic Equivalent] We say that $\overline{{Q}} \in \mathbb{R}^{n \times n}$ is a deterministic equivalent for the symmetric random matrix $Q \in \mathbb{R}^{n \times n}$ if, for sequences of deterministic matrix ${A} \in \mathbb{R}^{n \times n}$ and vectors $a, b \in \mathbb{R}^n$ of unit norms (operator and Euclidean, respectively), we have, as $n \rightarrow \infty$, $$\frac{1}{n} \operatorname{tr} {A}({Q}-\overline{{Q}}) \rightarrow 0, \quad {a}^{\top}({Q}-\overline{{Q}}) {b} \rightarrow 0,$$ where the convergence is either in probability or almost sure.

We refer the interested reader to the recent book rmtnew for a survey of the history of deterministic equivalents in random matrix theory, and several applications. We now discuss the relevance of this notion in our analysis.

Recall that we use OLS to fit the outcome regression parameters. For concreteness, suppose we use the second split $S_2$ to fit the outcome regression. Using properties of OLS, we note that the covariance matrix of $(\hat{\alpha}^{(1)}_{S_2} ,\hat{\beta}^{(1)}_{S_2})$ is $(\sum_{i\in S_2} A_i \tilde{x}_i \tilde{x}_i^{\top})^{-1}$, where $\tilde{x}_i^{\top} = (1, x_i^{\top})$ denotes the vector $x_i$ padded with an additional entry $1$ for the intercept. Similarly, the covariance matrix of $(\hat{\alpha}^{(0)}_{S_2} ,\hat{\beta}^{(0)}_{S_2})$ is $(\sum_{i\in S_2} (1-A_i) \tilde{x}_i \tilde{x}_i^{\top})^{-1}$.

The variable $T_1$ in (ref) involves terms of the form

align[align omitted — 356 chars of source]

where $\ell$ is a random vector independent of the samples in $S_2$, and a function of the covariates $x_i$'s and the exposure $A_i$'s. Thus the conditional variance $\mathrm{Var}(T_1|A,X)$ involves quadratic forms $\ell^{\top} \Big( \sum_{i\in S_2} A_i \tilde{x}_i \tilde{x}_i^{\top} \Big)^{-1} \ell$ and $\ell^{\top} \Big( \sum_{i\in S_2} (1-A_i) \tilde{x}_i \tilde{x}_i^{\top} \Big)^{-1} \ell$. To determine the limit of the conditional variance, it suffices to establish that these quadratic forms converge to deterministic limits as $n,p \to \infty$. To this end, we derive a deterministic equivalent of the covariance matrices---this allows us to replace the quadratic forms $\ell^{\top}\Big( \sum_{i\in S_2} A_i \tilde{x}_i \tilde{x}_i^{\top} \Big)^{-1} \ell$, $\ell^{\top}\Big( \sum_{i\in S_2} (1-A_i) \tilde{x}_i \tilde{x}_i^{\top} \Big)^{-1} \ell$ by quadratic forms with deterministic interaction matrices. This is crucial for our subsequent analysis, and aids us in deriving the limits of these quadratic forms.

{\bf Leave one out:} “Leave one out" style arguments have been critical in random matrix theory bai2010spectral, as well as in high-dimensional statistics bean2013optimal,el2013robust,el2018impact. This technique is also related to the cavity method from statistical physics montanari2022short,mezard2009information. In random matrix theory, this technique is ubiquitous, and is used for example in the proof of the limiting spectral distribution of a sample covariance matrix by the Stieljes transform method bai2010spectral. This idea has also been critical in establishing asymptotic distribution of classical estimators/test statistics in high-dimensional inference problems in the proportional asymptotics regime. To the best of our knowledge, this idea was first employed in high-dimensional statistics in the works of El Karoui and collaborators to analyze M-estimators in linear models bean2013optimal,el2013robust,el2018impact. Subsequently, it has been crucial for analyzing the MLE, LRT in logistic regression sur2019modern,sur2019likelihood, as well as diverse optimization problems ma2018implicit,chen2021spectral. Finally, this technique has been recently used to prove universality of high-dimensional estimation problems to the distribution of the feature vectors hu2020universality.

The nature of the technique as employed in the random matrix literature versus the high-dimensional statistics literature has subtle differences. Our analysis crucially employs both styles of leave-one-out arguments described above. To highlight the utility of this technique for our proofs, we sketch two intermediate arguments that utilize this idea.

First, we present the leave-one-out idea applied in the context of random matrices. Lemma (ref) establishes that

align[align omitted — 192 chars of source]

where $x_i$ are iid random vectors in $\mathbb{R}^{p}$ with iid $\mathcal{N}(0,1/n)$ entries. Setting $T_1 = \Big( \sum_{i=1}^{n} x_i \Big) \Big( \sum_{i=1}^{n} x_i \Big)^{\top}$ and $T_2 = T_1 - \sum_{i=1}^{n} x_i x_i^{\top}$, we have

align*[align* omitted — 143 chars of source]

where the last step follows from the identity $(A+B)^{-1} = A^{-1} - A^{-1} B (A+B)^{-1}$ for square matrices $A,B$. Thus it suffices to establish that

align*[align* omitted — 132 chars of source]

This analysis is involved as both $T_2$ and $\sum_{i=1}^{n} x_i x_i^{\top}$ depend on all the $x_i$ vectors. A natural strategy at this point is to isolate out the dependence of this expression on the individual $x_i$'s. Applying the Sherman-Morrison identities, one obtains that

align*[align* omitted — 290 chars of source]

where $x^{\{-i\}} = \sum_{j \neq i} x_j$. This representation isolates out $x_i$ from the other vectors---the resulting sum is easy to track by direct computation. Indeed, one completes the proof by directly establishing that the sum above has mean zero and variance converging to zero. This illustrates one instance of the leave-one-out idea in the context of random matrices, as utilized in our proof. We refer the interested reader to the proof of Lemma (ref) for additional details. While the above application of the leave-one-out is straightforward to the experts (and the result can be established without this technique for Gaussian covariates), we chose this example to provide a simple illustration of the technique in action. Our proofs invoke this technique in a large number of steps and often for expressions that are far more complicated than (ref). However, the underlying basic principle mostly remains similar to the above.

In addition to the abovementioned application of the leave-one-out, we utilize the technique crucially to track the asymptotic dependence between estimators used in cross-fitting. We emphasize that this is a major challenge in our proof; in comparison, this dependence is absent in the analysis of the AIPW estimator without cross-fitting, and the associated CLT proof would be significantly simpler. To explain the issue at a high-level, note that the cross-fitted AIPW includes a term where the first split is used to estimate the propensity score model, while the final plug-in is performed on the third split. Simultaneously, it includes a term where the roles of the first and third splits are flipped (Note that there is nothing special about these two terms---the same issue arises for many pairs of terms obtained from the sample splits.). Naturally, when we compute the variance of the cross-fitted estimator, we have to control all of the cross-covariances among these terms. This covariance is implicit, as the MLE $\hat{\beta}$ is a complicated function of the individual sample points. To compute this limiting covariance, our strategy is to replace the logistic MLE $\hat{\beta}$ by a surrogate $\hat{\beta}^{\{-i\}}$---the MLE on the sample with the $i^{th}$ datapoint left-out. The surrogate is independent of the $i^{th}$ datapoint by construction, and is critical for calculating the covariance. Crucially, one cannot replace the MLE with its surrogate $\hat{\beta}^{\{-i\}}$ without paying a price---the fitted values $x_i^{\top}\hat{\beta}$ and $x_i^{\top}\hat{\beta}^{\{-i\}}$ are different, and this difference shows up in our limiting covariance calculation. This difference has been precisely characterized in sur2019modern, and is a crucial We begin our proof by replacing the MLE with the leave-one-out surrogate. However, tracking the downstream effects of this replacement is highly non-trivial and can be viewed as one of our major technical contributions. Putting these ingredients together yields a fairly explicit expression for the limiting covariances and uncovers the negative cross-covariance phenomenon described in the Introduction.

Discussions and Open Questions

We discuss follow up questions arising from our results, and collect initial thoughts regarding their resolution.

itemize• The effect of winsorizing---It is well known that the finite sample performance of the AIPW might suffer due to the inverse probability weighing involved in its evaluation. To mitigate this issue, practitioners routinely use a winsorized version of the estimator. Formally, this corresponds to replacing $\sigma(x_i^{\top} \hat{\beta}_{S_a})$ (respectively $1-\sigma(x_i^{\top} \hat{\beta}_{S_a})$ by $\max\{\sigma(x_i^{\top} \hat{\beta}_{S_a}), \varepsilon\}$ (respectively $\max\{1-\sigma(x_i^{\top} \hat{\beta}_{S_a}), \varepsilon \}$) for some small $\varepsilon>0$. We see this finite sample effect also in our simulations (see Figure (ref)). The winsorizing regularizes the estimator, and removes the outliers in the q-q plot. We believe it should be possible to track the sampling distribution of the winsorized estimator using the tools introduced in this paper. For small $\varepsilon>0$, the limiting distributions of the original estimator and the winsorized one are approximately the same. We thus do not pursue a formal theoretical treatment of the winsorized estimator in this paper. • Constructing confidence intervals for the ATE using our CLT---Given our main result, one immediately wonders if it can yield confidence intervals for the ATE. Of course, this will require a consistent estimate of the sampling variance $\sigma_{cf}^2$. The expression for $\sigma_{cf}^2$ is quite involved, so it is a priori unclear if this is possible. However, on closer inspection we notice that the limiting variance is a function of the signal-to-noise ratio type parameters (ref), and consistent estimation of such quantities are known to be feasible in the proportional asymptotics regime. This has been demonstrated in a variety of prior works bayati2013estimating,dicker2016maximum,sur2019modern,yadlowsky2021sloe,bellec2022observable,janson2017eigenprism. A combination of these techniques should yield a consistent variance estimator in our setting. We will explore this direction in-depth in future work. As an aside, we note that traditional re-sampling approaches such as the bootstrap are known to be inconsistent in simpler statistical problems under the proportional asymptotics regime el2018can---we expect similar phenomena to hold in our setup. • Beyond the assumptions on the covariates---Our result assumes that the covariates $x_i \in \mathbb{R}^p$ are iid gaussian. We believe that the gaussianity is not critical for the validity of this result---indeed, we expect our results to be valid in settings where the entries of $x_i$ are iid, as long as these have well-behaved tail properties (e.g. sub-gaussian tails). Several results of this type have by now been established in the proportional asymptotics regime bayati2015universality,abbasi2019universality,liang2020precise,hu2020universality,montanari2022universality. Furthermore, our experiments in Figures (ref) and (ref) indicate the presence of such universality phenomenon in our setting. On the contrary, extending our results to allow for correlations among the features is less straightforward. Such situations are more natural in practical applications, thus establishing analogues of our results in these settings is of intrinsic interest. We expect this direction to be feasible, at least for special covariance structures or in the case of gaussian correlated covariates, following arguments similar to liang2020precise,zhao2020asymptotic. That said, we view this paper as a stepping stone for analyzing other causal effect estimators in the proportional asymptotics regime. Our proofs are significantly involved even under the stylized covariate distribution assumed herein. In this light, we defer generalizations of this condition to future works. • More general nuisance estimators--- In this paper, we focus on simple nuisance estimators such as the MLE or ridge regression. In high dimensions, one typically wishes to employ more sophisticated estimators for the nuisance parameters e.g., those arising from modern machine learning. The performance of the AIPW with such advanced nuisance estimators has been analyzed in the recent literature chernozhukov2017double,smucler2019unifying. To the best of our knowledge, all existing analyses of this flavor assume sparsity of either the propensity score or the outcome regression model, and thus are not directly applicable to our setting. It would be interesting to explore the effect of using powerful Machine Learning based nuisance estimators in our setting. We leave this for future work. • Alternative sample splitting schemes---We employ a three sample split strategy in this paper---the two nuisance functions are estimated from distinct sample folds, while the final estimator is constructed based on the third fold. This is certainly not the only possible choice for this problem; in our case, this three sample split strategy is a conscious choice, since this aids our theoretical analysis. However, one might naturally wish to use other splitting strategies e.g., the samples could be split into two parts, the two nuisances being estimated from the first split, and the final estimator being evaluated on the second split. The most extreme example would be to use the whole data to estimate both the nuisances, and the subsequent computation of the AIPW. Analyzing these estimators are significantly more challenging, due to the subtle dependencies among the intermediate estimators. It is apriori unclear which of these sample splitting schemes yields the estimator with the best empirical performance. We believe that extending our results to settings with fewer splits will require new technical ideas, and is an interesting direction for follow-up research. • The problem of optimal estimation---Our work raises the following natural question: is some version of the AIPW (with cross-fitting) optimal in terms of the asymptotic variance in this setting? Note that in the classical low-dimensional setting ($n \to \infty$ and $p$ fixed), the AIPW is semi-parametric efficient tsiatis2006semiparametric,bickel1993efficient,van2000asymptotic,le2000asymptotics in a nonparametric model that does not restrict the distribution of the tuple $(y,A,x)$. Our analysis assumes a specific covariate distribution---this assumption allows for more efficient ATE estimation in the classical regime kallus2020role. One might naturally wonder if this improved ATE estimator might beat the cross-fitted AIPW estimator, and continue to be optimal in our setting. We emphasize that these existing comparisons do not translate directly to our proportional asymptotics regime---in fact, pinning down “efficient" estimators in our context remains an outstanding question. We defer this direction to future work, and adopt the following perspective here. The AIPW is one of the most widely used ATE estimators in practice---can its fluctuations be characterized via the classical asymptotic variance when we are neither in a classical setting, nor in the ultra-high-dimensional regime with sparsity? Our central limit theorems provide an answer in the negative and develop alternate approximations that can be used for inferring the ATE in a large class of problems.
acks[Acknowledgments] PS acknowledges support from NSF DMS-2113426 and SS acknowledges support from a Harvard Dean's Competitive Fund Award. PS would like to thank Andrea Rotnitzky for helpful discussions on an earlier version of this manuscript.
appendix