The partially linear model
DML's workhorse is the partially linear regression (PLR) model:
Y = θ·D + g(X) + ε, D = m(X) + v
θ is the causal effect; g and m are unknown, possibly very complicated functions of many covariates. In Lesson 3.2 you saw that θ can be recovered by regressing residualised Y on residualised D (FWL). DML does exactly that, using flexible learners to estimate E[Y | X] and E[D | X].
!pip install -q numpy pandas scikit-learn statsmodels matplotlib
import numpy as np, pandas as pd
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import KFold
import statsmodels.api as sm
rng = np.random.default_rng(2)
High-dimensional, non-linear confounding
A digital-skills training programme, 20 background covariates from a census-style record, with participation and wages depending on a few of them non-linearly. True effect: $5,000.
N, P = 3_000, 20
X = rng.normal(0, 1, size=(N, P))
m_x = 1 / (1 + np.exp(-(0.8 * X[:, 0] - 0.5 * X[:, 1] ** 2 + 0.4 * X[:, 2] * X[:, 3]))) # P(D=1 | X)
D = rng.binomial(1, m_x)
g_x = 2_000 * np.sin(2 * X[:, 0]) + 3_000 * X[:, 3] + 1_500 * X[:, 1] ** 2 + 1_000 * X[:, 2] * X[:, 3]
Y = 5_000 * D + g_x + rng.normal(0, 4_000, N)
TRUE = 5_000
lin = sm.OLS(Y, sm.add_constant(np.column_stack([D, X]))).fit()
print(f"OLS with all 20 controls (linear): {lin.params[1]:,.0f} (truth {TRUE:,})")
Linear OLS cannot absorb sin(x), x² or interactions, so confounding leaks through.
Naive ML plug-in vs residual-on-residual
A tempting idea: fit a random forest of Y on (D, X) and read off how predictions change with D. Forests regularise, shrinking the effect of D, so this "plug-in" is biased. Residualising is better, but if we fit the learners and compute residuals on the same data, the residuals are too small (overfitting) and bias creeps back.
def rf():
return RandomForestRegressor(n_estimators=300, min_samples_leaf=5, max_features=0.5, n_jobs=-1, random_state=0)
# Plug-in: one model for Y on (D, X)
full = rf().fit(np.column_stack([D, X]), Y)
plug_in = (full.predict(np.column_stack([np.ones(N), X])) - full.predict(np.column_stack([np.zeros(N), X]))).mean()
# Residual-on-residual, in-sample (no cross-fitting)
res_y = Y - rf().fit(X, Y).predict(X)
res_d = D - rf().fit(X, D).predict(X)
in_sample = (res_d @ res_y) / (res_d @ res_d)
print(f"Plug-in forest: {plug_in:,.0f}")
print(f"Residual-on-residual, in-sample: {in_sample:,.0f}")
Cross-fitting by hand
Cross-fitting fixes overfitting: split the data into K folds; for each fold, train the learners on the other folds and compute residuals on the held-out fold. Every residual is then an honest out-of-sample error. Combine the residuals and run the final regression, with a standard error from the usual heteroskedasticity-robust formula.
def dml_plr(Y, D, X, learner, k=5, seed=0):
res_y, res_d = np.zeros_like(Y, dtype=float), np.zeros_like(D, dtype=float)
for train, test in KFold(k, shuffle=True, random_state=seed).split(X):
res_y[test] = Y[test] - learner().fit(X[train], Y[train]).predict(X[test])
res_d[test] = D[test] - learner().fit(X[train], D[train]).predict(X[test])
theta = (res_d @ res_y) / (res_d @ res_d)
eps = res_y - theta * res_d
se = np.sqrt(np.mean(res_d ** 2 * eps ** 2) / np.mean(res_d ** 2) ** 2 / len(Y))
return theta, se
theta, se = dml_plr(Y, D, X, rf)
print(f"Cross-fitted DML: {theta:,.0f} (SE {se:,.0f}; 95% CI [{theta - 1.96*se:,.0f}, {theta + 1.96*se:,.0f}])")
Why does this work when random forests are biased? Because the final step uses the product of two residuals. Errors in E[Y | X] and E[D | X] enter the estimate only through their product (Neyman orthogonality, Lesson 7.1), so if each learner is moderately accurate (error shrinking faster than n^(-1/4)), their product becomes negligible and θ is √n-consistent and approximately normal. Cross-fitting removes the own-observation overfitting bias that orthogonality alone cannot.
Run dml_plr with k=2 and k=10, and with a weaker learner (a forest with max_depth=2). Which choice
moves the estimate most? What does that say about tuning the nuisance models?
Show a solution
for k in (2, 10):
print(f"k={k:>2}: {dml_plr(Y, D, X, rf, k=k)[0]:,.0f}")
shallow = lambda: RandomForestRegressor(n_estimators=200, max_depth=2, n_jobs=-1, random_state=0)
print(f"Shallow forest: {dml_plr(Y, D, X, shallow)[0]:,.0f}")
# The number of folds matters little; a poor learner leaves confounding in the residuals.
# Nuisance-model quality is what matters: tune and compare learners by out-of-sample loss.
Checklist and recap
- I can write the partially linear model and connect it to FWL.
- I can explain the two sources of bias DML removes: regularisation (via orthogonality) and overfitting (via cross-fitting).
- I can implement cross-fitted DML by hand.