Regularization Bias and the Partially Linear Model

Summary

The partially linear regression (PLR) model $Y=D{\theta }_0+g_0(X)+U$, $D=m_0(X)+V$ (Robinson 1988) is the lead example of Chernozhukov et al. (2018). A naive ML estimator — learn $g_0$ with a regularized learner, then regress $Y-{\hat{g}}_0(X)$ on $D$ — fails to be $\sqrt{n}$-consistent: its scaled error contains a term $b$ of order $\sqrt{n}{n}^{-{\phi }_g}\to \infty$ because $D$ is centered at $m_0(X)\ne 0$ and correlates with the bias in ${\hat{g}}_0$. Partialling $X$ out of $D$ as well (“double prediction”) replaces this with a term proportional to the product $({\hat{m}}_0-m_0)({\hat{g}}_0-g_0)$, bounded by $\sqrt{n}{n}^{-({\phi }_m+{\phi }_g)}$, which vanishes whenever both learners beat the $n^{-1/4}$ rate. This is the origin of the name double/debiased machine learning.

Overview

The PLR model captures a causal design with selection on observables: if $D$ is as good as randomly assigned given $X$ (the Conditional Independence Assumption), ${\theta }_0$ is the average treatment effect of $D$ — “the ‘lift’ parameter in business applications” (§1.1, p. 2). The first equation is the outcome model; the second “keeps track of confounding, namely the dependence of the treatment variable on controls” and is “not of interest per se but is important for characterizing and removing regularization bias.” The dimension $p$ of $X$ is modeled as growing with $N$, so the nuisance ${\eta }_0=(m_0,g_0)$ lives in a space whose entropy grows with the sample — outside classical semiparametrics.

Why not just fit $Y\sim D\theta +g(X)$ with a flexible learner that leaves $\theta$ unpenalized (e.g. alternate between a random forest for $g$ and OLS for $\theta$)? Because every useful high-dimensional learner trades variance for bias, and that bias leaks into ${\hat{\theta }}_0$ through the correlation between $D$ and $X$. This is Omitted Variables Bias in a new guise: the “omitted variable” is the part of $g_0(X)$ that regularization shrank away, and it is correlated with $D$ through $m_0(X)$.

Main Content

Partially linear regression model ^def-plr

$$Y=D{\theta }_0+g_0(X)+U,\mathbb{E}[U\mid X,D]=0,$$ $$D=m_0(X)+V,\mathbb{E}[V\mid X]=0,$$

with outcome $Y$, scalar policy/treatment variable $D$, controls $X=(X_1,\dots ,X_p)$, and nuisance ${\eta }_0=(m_0,g_0)$. (Vector $D$: add one equation like the second per component.)

Failure of the naive plug-in estimator ^thm-naive-failure

Split the sample into a main part $I$ (size $n$) and auxiliary part $I^c$; fit ${\hat{g}}_0$ on $I^c$ and set

$${\hat{\theta }}_0=\left(\frac{1}{n}\sum _{i\in I}{D}_i^2{)}^{-1}\frac{1}{n}\sum _{i\in I}{D}_i\right(Y_i-{\hat{g}}_0(X_i)).$$

Then $\sqrt{n}({\hat{\theta }}_0-{\theta }_0)=a+b$ with

$$a=(\frac{1}{n}\sum D_i^2{)}^{-1}\frac{1}{\sqrt{n}}\sum _{i\in I}{D}_i{U}_i⇝\mathcal{N}(0,\bar{\Sigma }),$$ $$b=(\mathbb{E}{D}_i^2{)}^{-1}\frac{1}{\sqrt{n}}\sum _{i\in I}{m}_0(X_i)(g_0(X_i)-{\hat{g}}_0(X_i))+o_P(1).$$

Term $b$ is a sum of $n$ terms with non-zero mean divided by $\sqrt{n}$. If ${\hat{g}}_0$ converges at rate $n^{-{\phi }_g}$ with ${\phi }_g<1/2$, then $b$ is of stochastic order $\sqrt{n}{n}^{-{\phi }_g}\to \infty$, so $|\sqrt{n}({\hat{\theta }}_0-{\theta }_0)|{\to }_P\infty$ (eq. 1.4).

Note that sample splitting alone does not rescue the naive estimator — ${\hat{g}}_0$ above was already fit on an independent sample. The bias is a first-order sensitivity problem, not an overfitting problem.

The orthogonalized (double ML) estimator ^thm-dml-plr-decomp

Also fit ${\hat{m}}_0$ on $I^c$, form ${\hat{V}}_i=D_i-{\hat{m}}_0(X_i)$ and

$${\check{\theta }}_0=\left(\frac{1}{n}\sum _{i\in I}{\hat{V}}_i{D}_i{)}^{-1}\frac{1}{n}\sum _{i\in I}{\hat{V}}_i\right(Y_i-{\hat{g}}_0(X_i)).$$

Then $\sqrt{n}({\check{\theta }}_0-{\theta }_0)=a^{\ast }+b^{\ast }+c^{\ast }$ where

$$a^{\ast }=(\mathbb{E}{V}^2{)}^{-1}\frac{1}{\sqrt{n}}\sum _{i\in I}{V}_i{U}_i⇝\mathcal{N}(0,\Sigma ),$$ $$b^{\ast }=(\mathbb{E}{V}^2{)}^{-1}\frac{1}{\sqrt{n}}\sum _{i\in I}({\hat{m}}_0(X_i)-m_0(X_i)\left)\right({\hat{g}}_0(X_i)-g_0(X_i)),$$

so $|b^{\ast }|\lesssim \sqrt{n}{n}^{-({\phi }_m+{\phi }_g)}$, which vanishes e.g. when ${\phi }_m,{\phi }_g>1/4$. The remainder $c^{\ast }$ contains terms like $\frac{1}{\sqrt{n}}{\sum }_{i\in I}{V}_i({\hat{g}}_0(X_i)-g_0(X_i))$ and is $o_P(1)$ because of sample splitting (see Cross-Fitting and Sample Splitting).

${\check{\theta }}_0$ is a linear IV estimator with instrument $\hat{V}$: the residualized treatment is, by construction, (approximately) uncorrelated with anything that is a function of $X$, including the error in ${\hat{g}}_0$. A second, first-order-equivalent version is Robinson’s pure partialling-out estimator,

$${\check{\theta }}_0=\left(\frac{1}{n}\sum _{i\in I}{\hat{V}}_i{\hat{V}}_i{)}^{-1}\frac{1}{n}\sum _{i\in I}{\hat{V}}_i\right(Y_i-{\hat{\ell }}_0(X_i)),{\ell }_0(X)=\mathbb{E}[Y\mid X],$$

i.e. regress the outcome residual on the treatment residual — an ML-powered Frisch–Waugh–Lovell. Its practical advantage: both nuisances, ${\ell }_0$ and $m_0$, are plain conditional means that any off-the-shelf learner can target, whereas $g_0(X)=\mathbb{E}[Y-D{\theta }_0\mid X]$ depends on the unknown ${\theta }_0$.

Theorem 4.1 — DML inference in the PLR model ^thm-dml-plr

Under Assumption 4.1 (moment bounds with $q>4$; $\mathbb{E}[V^2]\ge c$; bounded conditional variances; nuisance estimators with $\Vert {\hat{\eta }}_0-{\eta }_0{\Vert }_{P,2}\le {\delta }_N$ and the product-rate condition $\Vert {\hat{m}}_0-m_0{\Vert }_{P,2}\times \Vert {\hat{g}}_0-g_0{\Vert }_{P,2}\le {\delta }_N{N}^{-1/2}$ for score (4.3), or $\Vert {\hat{m}}_0-m_0{\Vert }_{P,2}\times (\Vert {\hat{m}}_0-m_0{\Vert }_{P,2}+\Vert {\hat{\ell }}_0-{\ell }_0{\Vert }_{P,2})\le {\delta }_N{N}^{-1/2}$ for score (4.4)), the cross-fit DML1 and DML2 estimators are first-order equivalent and

$${\sigma }^{-1}\sqrt{N}({\tilde{\theta }}_0-{\theta }_0)⇝\mathcal{N}(0,1),{\sigma }^2=[\mathbb{E}{V}^2{]}^{-1}\mathbb{E}[V^2{U}^2][\mathbb{E}{V}^2{]}^{-1},$$

uniformly over $P\in \mathcal{P}$; the plug-in ${\hat{\sigma }}^2$ is consistent and $[{\tilde{\theta }}_0\pm {\Phi }^{-1}(1-\alpha /2)\hat{\sigma }/\sqrt{N}]$ is uniformly valid.

Two remarks sharpen this. Efficiency (Remark 4.2): under homoscedasticity ${\sigma }^2=\mathbb{E}[V^2{]}^{-1}\mathbb{E}[U^2]$, the semiparametric efficiency bound. Tightness (Remark 4.3): if $g_0,m_0$ are sparse with indices $s_g,s_m$ and estimated by ${\ell }_1$-penalization at rates $\sqrt{s_g/N},\sqrt{s_m/N}$, the product-rate condition reads $s_g{s}_m\ll N$ — far weaker than the $(s_g{)}^2+(s_m{)}^2\ll N$ needed without sample splitting. A very sparse propensity permits a dense outcome model and vice versa; if $m_0$ is known (a randomized experiment), mere consistency of ${\hat{g}}_0$ suffices.

The variance tells you where identification comes from. ${\sigma }^2\propto 1/\mathbb{E}[V^2]$: precision is driven by the variation in treatment not predicted by $X$. If a flexible ${\hat{m}}_0$ predicts $D$ almost perfectly, $\mathbb{E}[V^2]\approx 0$ and no method can help — the continuous-treatment analogue of an overlap failure.

Examples

Figure 1 of the paper (n = 500, p = 20). $g_0$ is “a very smooth function of a small number of variables” — a setting favorable to random forests. The naive forest-based ${\hat{\theta }}_0$ has a histogram “badly biased, shifted much to the right” of the truth and poorly described by its nominal normal approximation. On the same simulated data, the orthogonal estimator ${\check{\theta }}_0$ with forest nuisances is centered at ${\theta }_0$ and matches the $\mathcal{N}(0,\Sigma )$ curve.

Code sketch (partialling-out score with cross-fitting; any scikit-learn regressors):

import numpy as np
from sklearn.model_selection import KFold
from sklearn.ensemble import RandomForestRegressor
 
def dml_plr(Y, D, X, K=5, seed=0):
    res_y, res_d = np.zeros_like(Y, float), np.zeros_like(D, float)
    for train, test in KFold(K, shuffle=True, random_state=seed).split(X):
        l_hat = RandomForestRegressor(500, min_samples_leaf=5).fit(X[train], Y[train])
        m_hat = RandomForestRegressor(500, min_samples_leaf=5).fit(X[train], D[train])
        res_y[test] = Y[test] - l_hat.predict(X[test])   # Y - l_hat(X)
        res_d[test] = D[test] - m_hat.predict(X[test])   # V_hat = D - m_hat(X)
    theta = (res_d @ res_y) / (res_d @ res_d)            # DML2: pooled moment
    psi = (res_y - theta * res_d) * res_d                # score at theta
    J = np.mean(res_d ** 2)
    se = np.sqrt(np.mean(psi ** 2) / J ** 2 / len(Y))    # sandwich, Thm 3.2
    return theta, se

In the 401(k) application, this PLR estimator gives an eligibility effect of $7,717–$9,247 across learners (s.e. $\approx$ $1,300–$1,750 with the split-adjusted median method).

Connections

• Neyman Orthogonality — the general principle of which ”$b$ vs $b^{\ast }$” is the special case: ${\partial }_g\mathbb{E}[(Y-\theta D-g)D]\ne 0$ but ${\partial }_{\eta }\mathbb{E}[(Y-D\theta -g)(D-m)]=0$.
• Cross-Fitting and Sample Splitting — controls $c^{\ast }$ and restores full-sample efficiency.
• DML Estimators for ATE and the Interactive Model — drops the additive-separability restriction for binary $D$.
• R-Learner and Orthogonal CATE Estimation — lets ${\theta }_0$ become a function $\tau (X)$ in the same residual-on-residual equation.
• Omitted Variables Bias — regularization bias is OVB where the omitted term is $g_0-{\hat{g}}_0$.
• Horseshoe and Regularized Horseshoe Priors — shrinkage priors on control coefficients create the same leak into a treatment coefficient; residualizing the treatment (or including a propensity estimate) is the analogous Bayesian remedy.

See Also

• Table 2 Fallacy — $g_0$ and $m_0$ are nuisances; their fitted “effects” have no causal reading.
• Instrumental Variables — the partially linear IV model (§4.2) replaces $D-m(X)$ with $Z-m(X)$ in the score: $\psi =(Y-D\theta -g(X))(Z-m(X))$.
• Frequentist Causal Estimation — outcome-regression estimators are the “naive plug-in” family that this note shows to be fragile under ML.
• Regression and the CEF — the classical regression-anatomy logic that partialling-out generalizes.