Neyman Orthogonality and Cross-Fitting

52.1 Frisch–Waugh–Lovell, again (40 min)

Every economist has seen the Frisch–Waugh–Lovell theorem: the coefficient on D in a regression of Y on D and X equals the coefficient from regressing (Y minus its prediction from X) on (D minus its prediction from X). "Partial X out of both, then compare what's left."

Robinson (1988) noticed that the predictions don't have to be linear. In the partially linear model

spend = θ · coupon + g(X) + error,         coupon = m(X) + error

g and m can be any functions. Predict spend from X (call it l(X) = E[spend | X]), predict the coupon from X (m(X) = E[coupon | X]), and regress residual on residual. Machine learning is exactly a tool for estimating unknown functions like l and m. They're called nuisance functions: we need them, but we don't care about them.

52.2 Partialling out with machine learning (50 min)

Create day52/partialling_out.py:

# day52/partialling_out.py   (python -m day52.partialling_out)
"""
Frisch-Waugh-Lovell, then Robinson (1988), then double machine learning:
  1. predict spend from the features (no coupon):  l(X) = E[spend | X]
  2. predict the coupon from the features:          m(X) = E[coupon | X]
  3. regress (spend - l(X)) on (coupon - m(X)).
With linear predictions this reproduces OLS exactly. With LightGBM and cross-fitting, it's DML.
"""
import numpy as np
import statsmodels.api as sm
import statsmodels.formula.api as smf
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import cross_val_predict

from causalkit.coupons import X_COLS, coupons, public
from causalkit.learners import clf, reg, variance_weighted

full = coupons()
df = public(full).to_pandas()
X, y, d = df[X_COLS], df.spend, df.coupon


def residual_on_residual(y_res, d_res):
    fit = sm.OLS(y_res, d_res).fit(cov_type="HC1")
    return fit.params.iloc[0], fit.bse.iloc[0]


ols = smf.ols("spend ~ coupon + " + " + ".join(X_COLS), df).fit()
lin = residual_on_residual(y - LinearRegression().fit(X, y).predict(X), d - LinearRegression().fit(X, d).predict(X))
print(f"OLS coefficient:                          {ols.params['coupon']:.2f}")
print(f"partialling out with linear models:       {lin[0]:.2f}   <- identical (Frisch-Waugh-Lovell)")

l_hat = cross_val_predict(reg(), X, y, cv=5)
m_hat = cross_val_predict(clf(), X, d, cv=5, method="predict_proba")[:, 1]
dml = residual_on_residual(y - l_hat, d - m_hat)
print(f"partialling out with LightGBM, 5 folds:   {dml[0]:.2f}  (SE {dml[1]:.2f})")
print(f"Target: {variance_weighted(full['_tau'].to_numpy(), full['_ps'].to_numpy()):.2f} "
      "(the variance-weighted effect, explained on Day 53)")
print(f"\nNuisance models: spend predicted with R2 {1 - np.var(y - l_hat) / np.var(y):.2f}; "
      f"predicted coupon probabilities range from {np.percentile(m_hat, 5):.2f} to {np.percentile(m_hat, 95):.2f} "
      "(5th to 95th percentile)")

python -m day52.partialling_out

OLS coefficient:                          13.41
partialling out with linear models:       13.41   <- identical (Frisch-Waugh-Lovell)
partialling out with LightGBM, 5 folds:   16.08  (SE 0.56)
Target: 16.39 (the variance-weighted effect, explained on Day 53)

Nuisance models: spend predicted with R2 0.66; predicted coupon probabilities range from 0.18 to 0.86 (5th to 95th percentile)

Linear partialling-out reproduces OLS to the cent. Replacing the linear predictions with LightGBM moves the estimate from $13.4 to $16.1, within half a standard error of its target. That's double machine learning in five lines.

cross_val_predict(..., cv=5) is doing something important: every customer's prediction comes from a model trained on the other four-fifths of customers.

52.3 Two ideas that make it work (45 min)

1. Neyman orthogonality: residualise BOTH spend and coupon. If you only residualise spend, any error in l(X) that correlates with the coupon goes straight into θ. Residualising the coupon too makes θ insensitive to small errors in either nuisance: errors enter only as a product of the two models' errors. Two decent models, each wrong by a little, give a much smaller error in θ. That's the "double" in double ML, and the same property made AIPW doubly robust on Day 36.

2. Cross-fitting: never predict a customer with a model that saw them. Flexible models overfit: in-sample, they partly memorise each customer's outcome, so residuals are too small, and in a way that correlates with the coupon. Splitting the data into K folds, fitting on K−1 and predicting the held-out fold (then rotating), removes this. It's the Day 36 cross-fitting, now essential.

With both, θ̂ is approximately normal around the truth with the usual standard error, as long as both nuisance models are reasonably accurate (formally: their errors shrink faster than n^(−1/4)).

52.4 How many folds? (40 min)

Try it. Rerun the LightGBM partialling-out with 2, 5 and 10 folds. Create day52/folds.py.

# day52/folds.py   (python -m day52.folds)
import statsmodels.api as sm
from sklearn.model_selection import cross_val_predict

from causalkit.coupons import X_COLS, coupons, public
from causalkit.learners import clf, reg

df = public(coupons()).to_pandas()
X, y, d = df[X_COLS], df.spend, df.coupon
for k in [2, 5, 10]:
    l_hat = cross_val_predict(reg(), X, y, cv=k)
    m_hat = cross_val_predict(clf(), X, d, cv=k, method="predict_proba")[:, 1]
    fit = sm.OLS(y - l_hat, d - m_hat).fit(cov_type="HC1")
    print(f"{k:>2} folds: {fit.params.iloc[0]:.2f} (SE {fit.bse.iloc[0]:.2f})")

 2 folds: 15.89 (SE 0.56)
 5 folds: 16.08 (SE 0.56)
10 folds: 16.20 (SE 0.56)

52.5 Day 52 mini-project: break orthogonality, break cross-fitting (75 min)

Three estimators on 40 datasets of 5,000 customers. Each estimate's error is divided by its own standard error: a valid estimator gives a standard normal, with 95% of 95% CIs covering the target. Create day52/orthogonality_mc.py (about 5 minutes):

# day52/orthogonality_mc.py   (python -m day52.orthogonality_mc)
"""
Mini-project: why BOTH residuals, and why cross-fitting? 40 simulated datasets of 5,000 customers.
  A. partial out spend only:  regress (spend - l(X)) on the raw coupon      (not orthogonal)
  B. partial out both, but predict each customer with a model trained on that same customer
     (an overfitting forest, in-sample)                                       (no cross-fitting)
  C. partial out both, out-of-fold predictions                                (DML)
Errors are divided by each estimate's standard error: a well-behaved estimator gives a standard normal.
"""
import matplotlib.pyplot as plt
import numpy as np
import polars as pl
import statsmodels.api as sm
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.model_selection import cross_val_predict

from causalkit.coupons import X_COLS, coupons, public
from causalkit.learners import clf, reg, variance_weighted


def rr(y_res, d_res):
    fit = sm.OLS(y_res, d_res).fit(cov_type="HC1")
    return float(np.asarray(fit.params)[0]), float(np.asarray(fit.bse)[0])


z = {"A. outcome residual only": [], "B. both residuals, no cross-fitting": [], "C. DML (cross-fitted)": []}
bias = {k: [] for k in z}
for s in range(40):
    full = coupons(n=5000, seed=1000 + s)
    df = public(full).to_pandas()
    X, y, d = df[X_COLS], df.spend.to_numpy(), df.coupon.to_numpy().astype(float)
    target = variance_weighted(full["_tau"].to_numpy(), full["_ps"].to_numpy())
    l_cf = cross_val_predict(reg(), X, y, cv=5)
    m_cf = cross_val_predict(clf(), X, d, cv=5, method="predict_proba")[:, 1]
    l_in = RandomForestRegressor(200, min_samples_leaf=2, random_state=0, n_jobs=-1).fit(X, y).predict(X)
    m_in = RandomForestClassifier(200, min_samples_leaf=2, random_state=0, n_jobs=-1).fit(X, d).predict_proba(X)[:, 1]
    for k, (yr, dr) in zip(z, [(y - l_cf, d - d.mean()), (y - l_in, d - m_in), (y - l_cf, d - m_cf)]):
        est, se = rr(yr, dr)
        z[k].append((est - target) / se)
        bias[k].append(est - target)

rows = [{"estimator": k, "average error ($)": round(np.mean(bias[k]), 2),
         "average error / SE": round(np.mean(v), 2), "95% CIs covering the target": f"{np.mean(np.abs(v) < 1.96):.0%}"}
        for k, v in z.items()]
with pl.Config(tbl_width_chars=120, fmt_str_lengths=40):
    print(pl.DataFrame(rows))

fig, axes = plt.subplots(1, 3, figsize=(12, 3.5), sharey=True)
grid = np.linspace(-6, 6, 200)
for ax, (k, v) in zip(axes, z.items()):
    ax.hist(np.clip(v, -6, 6), bins=np.linspace(-6, 6, 25), density=True, color="#2563eb", alpha=0.7)
    ax.plot(grid, np.exp(-grid**2 / 2) / np.sqrt(2 * np.pi), color="k", lw=1)
    ax.set_title(k, fontsize=9)
    ax.set_xlabel("(estimate - target) / SE")
fig.tight_layout()
fig.savefig("figures/day52_orthogonality.png", dpi=120)
print("Saved figures/day52_orthogonality.png")

python -m day52.orthogonality_mc

│ estimator                           ┆ average error ($) ┆ average error / SE ┆ 95% CIs covering the target │
│ A. outcome residual only            ┆ -2.81             ┆ -2.66              ┆ 20%                         │
│ B. both residuals, no cross-fitting ┆ -2.53             ┆ -2.52              ┆ 32%                         │
│ C. DML (cross-fitted)               ┆ -0.47             ┆ -0.41              ┆ 92%                         │
Saved figures/day52_orthogonality.png

Open the figure: A and B are shifted far to the left of the standard normal curve; C sits close to it. Without orthogonality, 95% intervals cover the target 20% of the time; without cross-fitting, 32%. With both, 92%. This reproduces the famous figure in Chernozhukov et al. (2018).

52.6 Save your work with Git (20 min)

git add .
git commit -m "Day 52: partialling out, orthogonality, cross-fitting, Monte Carlo"
git push

Day 52 checklist

• I can state the Frisch–Waugh–Lovell theorem and Robinson's partially linear model
• I can estimate an effect by cross-fitted partialling-out with LightGBM
• I can explain why both residuals are needed (orthogonality) and why cross-fitting is needed
• I can show by simulation that skipping either one breaks the confidence intervals