Debiased Machine Learning without Sample-Splitting for Stable Estimators

Debiased Machine Learning without Sample-Splitting for Stable Estimators

Introduction

A large variety of problems in causal inference and more generally semi-parametric inference can be framed as finding a solution to a moment condition:

where $Z\in {\mathcal Z}$ is a vector of random variables that, apart from the target parameter $\theta_0\in \Theta\subset\mathbb{R}^p$ of interest, also depends on unknown nuisance functions $g_0\in {\mathcal G}$, which need to be estimated in a flexible manner from the data. This framework has a long history in the literature on semi-parametric inference levit1976efficiency,hasminskii1979nonparametric,ibragimov1981statistical,pfanzagl1982lecture,klaassen1987consistent,robinson1988root,van1991differentiable,bickel1993efficient,newey1994asymptotic,robins1995semiparametric,vaart,bickel1988estimating,newey1998undersmoothing,ai2003efficient,newey2004twicing,ai2007estimation,tsiatis2007semiparametric,kosorok2007introduction,ai2012semiparametric, which analyzes the following two-stage estimation process with sample re-use, when having access to $n$ iid samples $\{Z_1, \ldots, Z_n\}$ and as $n$ grows, treating the target parameter dimension $p$ as a constant:

Given that the estimation of function $\hat{g}$ is a complex non-parametric problem, it will typically only satisfy slower than parametric error rates. Moreover, in high dimensional settings, when machine learning techniques are used to estimate $\hat{g}$, then the regularization bias of $\hat{g}$ will propagate to the final estimate $\hat{\theta}$, leading to non-regular estimates and the inability to construct confidence intervals.

The literature on efficient semi-parametric estimation provides conditions on the moment function so that the influence of the estimation error of the nuisance function $\hat{g}$ is of second order importance and does not alter the distributional properties of the second stage hasminskii1979nonparametric,bickel1988estimating,zhang2014confidence,belloni2011inference,belloni2014inference,belloni2014uniform,belloni2014pivotal,javanmard2014confidence,javanmard2014hypothesis,javanmard2018debiasing,van2014asymptotically,ning2017general,chernozhukov2015valid,neykov2018unified,ren2015asymptotic,jankova2015confidence,jankova2016confidence,jankova2018semiparametric,bradic2017uniform,zhu2017breaking,zhu2018linear. This approach dates back to the classical work on doubly robust estimation and targeted maximum likelihood robins1995analysis,robins1995semiparametric,van2006targeted,van2011targeted,luedtke2016statistical,toth2016tmle as well as the more recent work on locally robust or Neyman orthogonal conditions on the moment function neyman1959,neyman1979c,chernozhukov2016locally,belloni2017program,chernozhukov2018double, typically referred to as double or debiased machine learning. At a high-level, the earlier literature on semi-parametric inference shows that if the moment function satisfies some form of robustness to nuisance perturbations and, importantly, as long as the space ${\mathcal G}$ used in the first stage estimation is a relatively simple function class, in terms of statistical complexity, typically referred to as a Donsker function class, then the second stage estimate is root-$n$ consistent and asymptotically normal. Hence, one can easily construct confidence intervals for the target parameter of interest.

Crucially the recent literature on debiased machine learning alters the standard two-stage estimation algorithm to introduce the idea of sample-splitting, bickel1982adaptive,zheng2010asymptotic. In particular, instead of estimating $\hat{g}$ on all the samples, the recent work on debiased machine learning chernozhukov2018double,chernozhukov2016locally,chernozhukov2021simple estimates $\hat{g}$ on a separate sample, or for better sample efficiency invokes “cross-fitting,” where we train a nuisance model on half the data and evaluate it in the second stage on the other half and vice versa. Sample splitting avoids the Donsker conditions that were prevalent in the classic semi-parametric inference literature and only requires a root-mean-squared-error (RMSE) guarantee of the estimate $\hat{g}$ of $o_p(n^{-1/4})$.

However, sample splitting or cross-fitting still leads to poorer sample usage, as we can lose half of our data when training complex non-parametric or machine learning models, which can be problematic in small and moderate sample regimes. Our main result is to show that root-$n$ consistency and asymptotic normality of the standard algorithm, without sample splitting, can be achieved without the Donsker property, but solely if one assumes that the first stage estimation algorithm is $o(n^{-1/2})$ leave-one-out stable, a relatively widely studied property in the statistical machine learning and generalization theory literature bousquet2002stability,Kale11cross-validationand,elisseeff2003leave,Hardt2016. As a leading example we show that our stability conditions are satisfied by bagging estimators formed with sub-sampling without replacement.

Recent prior work of chernozhukov2020adversarial also considered asymptotic normality based on stability conditions, but as we expand in the main text, their requirement on the stability property is much harsher than the one we derive here. Moreover, chernozhukov2020adversarial analyzes only a special case of the class of moment problems that we consider here. For instance, in our leading example of bagging estimators, the prior result of chernozhukov2020adversarial would require that the bias of the base estimator decays faster than $1/n$, where $n$ is the sample size, which is typically not the case. In contrast, our stability condition does not impose any explicit assumption on the bias and solely requires that the sub-sample size $m$ is $o(\sqrt{n})$.

To simplify the regularity assumptions required for asymptotic normality, we focus on the case where $m(Z;\theta,g)$ is linear in $\theta$, i.e.

where $a(Z;g)\in \mathbb{R}^{p\times p}$ is a $p\times p$ matrix and $\nu(Z;g)\in \mathbb{R}^p$ is a $p$-vector, and we denote with:

Many of the leading examples in semi-parametric problems that arise in causal inference correspond to linear moment problems. We present below a representative set of problems that are widely used in the practice of causal inference.

For completeness, we also include in Appendix (ref) an extension of our results to nonlinear moment problems.

Asymptotic Normality without Sample Splitting

We start by providing an asymptotic normality theorem for semi-parametric moment estimators without sample splitting and where the moment satisfies the well-studied property of Neyman orthogonality. Our theorem requires four main conditions: i) root-mean-squared-error (RMSE) rates for the nuisance function estimates of $o_p(n^{-1/4})$, ii) Neyman orthogonality of the moment with respect to the nuisances, iii) second-order smoothness of the moment with respect to the nuisance functions and, iv) stochastic equicontinuity of the Jacobian and the offset part of the linear moment function as the nuisance estimate $\hat{g}$ converges to $g_0$. For a vector $x\in \mathbb{R}^p$ we denote with $\|x\|_2$ the $\ell_2$ norm and for a matrix $X\in \mathbb{R}^{p\times p}$ we denote with $\|X\|_{op}$ the operator norm with respect to the $\ell_2$ norm.

The first three conditions are standard assumptions in the literature on debiased machine learning. The final condition (stochastic equicontinuity) is exactly where sample splitting comes very handy in the literature. To illustrate the reason why sample splitting helps with the stochastic equicontinuity condition, let us consider the first part of the condition (the reasoning is analogous for the second part). It asks that the difference of two centered empirical processes, namely $A_n(\hat{g})-A(\hat{g})$ and $A_n(g_0) - A(g_0)$, goes to zero faster than $n^{-1/2}$. We expect each empirical process to go down to zero at exactly $n^{-1/2}$ and so this condition asks, since $\hat{g}$ converges to $g_0$ is the empirical process continuous in its argument and for that reason does the difference converge to zero faster than each individual component. If the estimate $\hat{g}$ was fitted on a separate sample, then conditional on $\hat{g}$, we have that each element $t=(i,j)\in [p]\times[p]$ of $A_n(\hat{g}) - A_n(g_0)$ is an empirical average of iid random variables with mean $A(\hat{g})-A(g_0)$. Thus a simple Bernstein inequality would show that the difference of the two empirical processes would converge to zero at the order of:

where we also invoked a mean-squared-continuity property of $a_t(Z;g)$ and the fact that $\|\hat{g}-g_0\|_2=o_p(1)$. Thus, with sample splitting, no further constraint is required from $\hat{g}$, other than a convergence rate on $\|\hat{g}-g_0\|_2$. In fact, as was noted in recent work of chernozhukov2021simple, in the above step it suffices to assume that $\mathbb{E}[(a_t(Z;\hat{g}) - a_t(Z;g_0))^2] = O\left(\|\hat{g}-g_0\|_2^q\right)$ for any $q<\infty$, which is a much weaker mean-squared-continuity assumption, and the property would still hold, since $\|\hat{g}-g_0\|_2^{q/2} n^{-1/2}=o(n^{-1/2})$, whenever $\|\hat{g}-g_0\|_2=o_p(1)$.

Without sample splitting, note that $\hat{g}$ is now correlated with the samples in the empirical averages and hence $A_n(\hat{g}) - A_n(g_0)$ is no longer an average of i.i.d. random variables. Typical approaches would try to prove a uniform stochastic equicontinuity property over the function space ${\mathcal G}$, typically referred to as a Donsker property of the function space ${\mathcal G}$. In particular, if we could show that with high probability:

then the above property would also hold for $\hat{g}$. Subsequently, since we know that $\|\hat{g}-g_0\|_2^2=o_p(n^{-1/2})$, by our convergence rate assumptions on $\hat{g}$, then it would suffice that $\delta_n^2 = o(n^{-1/2})$. Such localized concentration inequalities have been known to hold for Donsker classes, which are typically defined via entropy integrals, and more recently it was also noted that such inequalities are satisfied with $\delta_n$ being the critical radius of the space ${\mathcal G}$, defined via localized Rademacher complexities.

However, the latter approach is conservative as it requires a uniform control over the function space ${\mathcal G}$ and does not utilize at all the properties of the estimation algorithm itself. In particular, as we will show in the next section, the main result of our work is that this stochastic equicontinuity condition follows from $o(n^{-1/2})$ leave-one-out stability conditions on our estimation algorithm, which are typical in the machine learning literature and in the excess risk and generalization bounds literature.

Stochastic Equicontinuity via Stability

We will show that the Condition (ref) is satisfied, whenever the estimate $\hat{g}$ satisfies leave-one-out stability properties and the moment satisfies the weak mean-squared-continuity property of chernozhukov2021simple. We start by some preliminary definitions required to state our stability conditions. Define $Z^{(-l)}$ as the data $Z_1,\ldots,Z_n$ with the $l$-th data point $Z_l$ replaced with an independent copy $\tilde{Z}_l$. Define $Z^{(-l_1,-l_2)}$ as the data $Z_1,\ldots,Z_n$ with both the $l_1$-th and the $l_2$-th data points $Z_{l_1},Z_{l_2}$ replaced with independent copies $\tilde{Z}_{l_1},\tilde{Z}_{l_2}$. Define similarly for $Z^{(-l_1,-l_2,-l_3)}$ and so on. Let $\hat g^{(-l)}$ be the estimator trained on $Z^{(-l)}$ instead of $Z_1,\ldots,Z_n.$ Similarly, let $\hat g^{(-l_1,-l_2)}$ be trained on $Z^{(-l_1,-l_2)}$ and so on. Moreover, we will always denote with $Z$ a fresh random variable drawn from the distribution of the samples, but which is not part of any training sample. For any random variable $X$, we denote with $\|X\|_1:=\mathbb{E}[|X|]$, with $\|X\|_2:=\sqrt{\mathbb{E}[X^2]}$, and with $\|X\|_p:=\left(\mathbb{E}[|X|^p]\right)^{1/p}$ for any $p\ge 1$ in general.

Tightness of Stability Condition

We present here an example that shows that, without further structural constraints (on the moment or the function space ${\mathcal G}$), the stability condition we impose is required for stochastic equicontinuity condition to hold. Let $(X_i)\overset{i.i.d}{\sim}\rm{unif}[0,1]$ be i.i.d uniform random variables. We set $Y_i:=\mathbb{I}(X_i\le 0.5)$ and set $Z_i:=(X_i,Y_i)$. For any $x\in [0,1]$, we define $c(Z_{1:n},x):=\operatorname*{arg\,min}_{i\le n} |X_i-x|$ the function that returns the index of the nearest example to $x$ in $\{X_1,\dots,X_n\}$ and note that the quantity $Y_{c(Z_{1:n},x)}$ is its nearest neighbour estimator. Let $$\hat{g}(Z_{1:n})(x,y):=n^{1/6}\mathbb{I}(y\ne Y_{c(Z_{1:n},x)}),\qquad \nu(Z,g):= g(Z)^3.$$ We remark that $\nu(\cdot,\hat g)$ does not satisfy our stability conditions as $ \|\nu(Z,\hat g)-\nu(Z,\hat g^{(-1)})\|_2$ is exactly of order $1/\sqrt{n}$, neither does it respect the stochastic equicontinuity property.

Application: Bagging Estimators

We remark that if the functions $a$ and $\nu$ satisfy certain $L^p$-Lipchitz conditions then the stability conditions of (ref) are implied by the algorithmic stability of the estimator $\hat g$. Those conditions are only marginally stronger than the condition of mean-squared-continuity found in (ref)

The uniform $L^{2r}$-continuity conditions are going to be satisfied by most moment functions $m(\cdot;\cdot,\cdot)$. We notably show in the appendix that (ref), (ref) and (ref) with general moment conditions satisfy our conditions.

The condition of (ref) is a commonly assumed condition in recent literature austern2020asymptotics, bayle2020cross and is satisfied by various regularized empirical risk minimization estimators and stochastic gradient descent estimators bousquet2002stability, hardt2016train. Notably, it is satisfied by bagged estimators. We show in (ref) that under very general conditions a bagged ensemble of any machine learning estimator is stable. Let $Z_1,\ldots,Z_n\in\mathcal{Z}$ be an independent and identically distributed (i.i.d.) sample of size $n$. We sample uniformly randomly without replacement from these observations repeatedly and independently, each time taking a sample of size $m$, for a number of $B$ times. We denote the resulting samples as $$Z_{1:m}^b:=\left\{Z_1^b,\ldots,Z_m^b\right\},b\in[B].$$

Let $\hat{h}_m:\mathcal{Z}^m\rightarrow \mathbb{R}^P$ be a base machine learning estimator trained on $m$ observations. This base estimator can be tree, a CNN, a nearest-neighbour classifier, or any other type of machine learning estimator. The corresponding bagged estimator is $$\hat{g}(\cdot)=\frac{1}{B}\sum_{b=1}^B\hat{h}_m(Z_{1:m}^b)(\cdot).$$

We will show that subject to $B$ and $m$ being sufficiently big and mild moment conditions on the base estimator $\hat{h}_m$, the bagged estimator satisfies condition (ref).

In particular, we can take the base machine learning estimator to be the 1-nearest neighbor estimator. This specific choice leads to a bagged estimator that can be proved to satisfy both our stability conditions and our consistency rate condition in Theorem (ref), in regression settings where covariates are of small intrinsic dimension.

Let $\left\{Z_i=\left(X_i, Y_i\right)\right\}_{i=1}^n$ be a sample of size $n$, drawn independently and identically distributed from $Z=(X,Y)$. Here $X\in\mathcal{X}\subset \mathbb{R}^D$ are known covariates and $Y\in\mathcal{Y}\subset \mathbb{R}^P$ is the response variable. As in our bagging setting, let $Z_{1:m}^1,\ldots,Z_{1:m}^B$ be B independent samples of size $m$ drawn without replacement from observations $\left\{Z_i\right\}_{i=1}^n$. A bagged 1-nearest neighbor (1-NN) estimator takes the following form: $$\hat g(x):=\frac{1}{B}\sum_{b=1}^B\sum_{i=1}^n \mathbbm{1}_{\left\{X_i= S_b(x)\right\}}Y_i,$$ where $S_b(x)$ is the 1-NN of $x$ in the set $Z_{1:m}^b.$ The estimator is used to estimate the conditional expectation $g_0(x):=\mathbb{E}[Y\mid X=x].$

In particular, we note that when $0<d<2$ and $m=O(n^{\frac{1}{2}-\epsilon})$ for some $0<\epsilon\leq\frac{1}{2}-\frac{1}{4}d,$ we immediately have that $$\sqrt{\mathbb{E}\left[\left\|\hat{g}(X)-g_0(X)\right\|_2^2\right]}=o(n^{-1/4}),$$ achieving the convergence rate assumed in Theorem (ref). Moreover, provided that we choose $B$ to be sufficiently large such that $$B>>n^{\epsilon+\frac{1}{2}} \hspace{1cm} B\geq n^{1-\frac{2\epsilon+1}{k}},$$ the required stability conditions can also be satisfied.

Experimental Evaluation

We consider a synthetic experimental evaluation of our main theoretical findings. We focus on the partially linear model with a scalar outcome $Y \in \mathbb{R}$, a scalar continuous treatment $T\in \mathbb{R}$ and many controls $X\in \mathbb{R}^{n_x}$, where:

Our goal is the estimation of the treatment effect $\theta_0$, while estimating the nuisance functions $p_0(X)$ and $q_0(X)=\theta_0 p_0(X) + f_0(X)$ in a flexible manner. We will consider estimation based on the orthogonal moment presented in Example (ref).

We considered sub-sampled $1$-nearest neighbor regression (NN) and sub-sampled fully grown (with only $1$ sample minimum leaf size) random forest (RF) regression for the estimation of $p_0$ and $q_0$, regressing $T$ on $X$ and $Y$ on $X$ correspondingly.

The true functions $p_0$ and $f_0$ are actually linear in our data generating process, with $p_0(X)=\beta_0'X$ and $f_0(X)=\gamma_0'X$ and where $\beta_0$ and $\gamma_0$ have only one non-zero coefficient ($1$-sparse), and which is the same coefficient for both $\beta_0$ and $\gamma_0$. In other words, only one of the $n_x$ potential confounding variables $X$ is actually a confounder.

We evaluate the performance of the estimate for $\theta_0$ for a range of values of the sample size $n$ and the dimension of the controls $n_x$ and with or without cross-fitting. For the cross-fitted estimates we used $2$ splits. For each specification we draw $1000$ experimental samples to evaluate the distributional properties of the estimate.

We considered sub-sample sizes for the nuisance regressions based on our theoretical $n^{0.49}$ specification and for larger specifications too. We find that the estimate without cross-fitting typically has almost equal bias and smaller (and always comparable) variance (due most probably to the smaller mean squared error of the nuisances, since they are trained on larger sizes), especially in smaller sample sizes, and has better coverage properties, when the sub-sample size is $m=o(n^{1/2})$. Moreover, the estimate is approximately normally distributed, even without cross-fitting as is verified qualitatively via quantile-quantile (Q-Q) plots.

We also report the mean of the estimate of the standard error across the $1000$ experiments, to evaluate the bias in the estimation of the standard error. We find that the estimate of the standard error is more accurate without cross-fitting, potentially because the estimate of the standard error does not incorporate the extra variance that stems from the sample-splitting process. This inaccuracy of the standard error estimate is most probably the reason for the worst coverage properties of the confidence intervals with cross-fitting.

In summary, we verify experimentally that for stable estimators, with the theoretically required level of stability, sample splitting or cross-fitting is not needed to maintain asymptotic normality, small bias and nominal coverage.

As expected however, we do verify that as the sub-sample size becomes $m=O(n)$, where the estimate is both not stable and also does not optimize over hypothesis spaces with small critical radius or that satisfy the Donsker property, then the estimate without cross-fitting breaks down, while the estimate with cross-fitting maintains a decent performance, despite the high-variance of the nuisance estimate.

The experiments were run on a normal PC laptop with Processor Intel(R) Core(TM) i7-8650U CPU @ 1.90GHz, 2112 Mhz, 4 Core(s), 8 Logical Processor(s), 16GB RAM. It took around 1.5 hours to run all the experiments.