From one number to a function
The ATE answers "does the programme work on average?". A ministry with a limited budget needs more: for whom does it work? The conditional average treatment effect is
τ(x) = E[Y(1) − Y(0) | X = x],
the average effect for people with characteristics x. Two programmes with the same ATE can have very different CATEs: one helps everybody a little, the other helps a few people a lot.
!pip install -q numpy pandas scikit-learn lightgbm matplotlib seaborn
import numpy as np, pandas as pd
from lightgbm import LGBMRegressor, LGBMClassifier
import matplotlib.pyplot as plt, seaborn as sns
rng = np.random.default_rng(10)
sns.set_theme(style="whitegrid")
A retraining programme with heterogeneous effects
A retraining stipend for unemployed workers. The effect on earnings one year later is large for the long-term unemployed and for workers with little schooling, and close to zero for recent job losers with degrees. Participation is partly self-selected.
def make_data(n):
X = pd.DataFrame({
"months_unemployed": rng.gamma(2, 4, n),
"years_education": rng.integers(8, 21, n).astype(float),
"age": rng.uniform(20, 60, n),
"prior_earnings": rng.lognormal(10.3, 0.4, n) / 1_000, # thousands of dollars
"female": rng.binomial(1, 0.5, n).astype(float),
})
tau = 500 + 250 * np.minimum(X.months_unemployed, 24) - 150 * (X.years_education - 12)
p = 1 / (1 + np.exp(-(-0.5 + 0.05 * X.months_unemployed - 0.03 * (X.age - 40))))
d = rng.binomial(1, p)
y0 = 15_000 + 400 * X.prior_earnings + 300 * (X.years_education - 12) - 120 * X.months_unemployed \
+ rng.normal(0, 3_000, n)
y = y0 + d * tau
return X, d.astype(float), y.to_numpy(), tau.to_numpy()
X, d, y, tau = make_data(6_000)
X_test, _, _, tau_test = make_data(2_000)
print(f"True ATE: {tau.mean():,.0f} | effects range from {tau.min():,.0f} to {tau.max():,.0f}")
S-, T- and X-learners
Meta-learners turn any prediction algorithm into a CATE estimator.
| Learner | Recipe | Strength / weakness |
|---|---|---|
| S (single) | fit μ(x, d) with d as a feature; τ̂(x) = μ(x,1) − μ(x,0) | simple, but regularisation may ignore d and shrink τ̂ to 0 |
| T (two) | fit μ₁ on treated, μ₀ on controls; τ̂ = μ₁ − μ₀ | flexible, but noisy when one group is small |
| X (cross) | impute individual effects with the other group's model, model them, blend with e(x) | good with unbalanced groups |
def gbm():
return LGBMRegressor(n_estimators=300, learning_rate=0.05, num_leaves=15, verbose=-1)
# S-learner
s = gbm().fit(X.assign(d=d), y)
cate_s = s.predict(X_test.assign(d=1.0)) - s.predict(X_test.assign(d=0.0))
# T-learner
mu1 = gbm().fit(X[d == 1], y[d == 1]); mu0 = gbm().fit(X[d == 0], y[d == 0])
cate_t = mu1.predict(X_test) - mu0.predict(X_test)
# X-learner
d1 = y[d == 1] - mu0.predict(X[d == 1]) # imputed effects for the treated
d0 = mu1.predict(X[d == 0]) - y[d == 0] # imputed effects for the controls
tau1 = gbm().fit(X[d == 1], d1); tau0 = gbm().fit(X[d == 0], d0)
e = LGBMClassifier(n_estimators=200, learning_rate=0.05, verbose=-1).fit(X, d).predict_proba(X_test)[:, 1]
cate_x = e * tau0.predict(X_test) + (1 - e) * tau1.predict(X_test)
print(f"Share treated: {d.mean():.2f}")
Evaluating and summarising CATEs
In a simulation we know τ(x), so we can score each learner with the PEHE (precision in estimating heterogeneous effects): the root mean squared error of τ̂(x) against τ(x). In real data you cannot, so you rely on the honest validation tools in Lesson 8.3.
pehe = {name: np.sqrt(np.mean((est - tau_test) ** 2))
for name, est in [("S-learner", cate_s), ("T-learner", cate_t), ("X-learner", cate_x)]}
print(pd.Series(pehe).round(0).rename("PEHE ($)"))
fig, axes = plt.subplots(1, 3, figsize=(15, 4), sharey=True)
for ax, (name, est) in zip(axes, [("S", cate_s), ("T", cate_t), ("X", cate_x)]):
ax.scatter(tau_test, est, s=5, alpha=.3); ax.plot([tau_test.min(), tau_test.max()], [tau_test.min(), tau_test.max()], "r--")
ax.set(title=f"{name}-learner", xlabel="True CATE", ylabel="Estimated CATE")
plt.tight_layout(); plt.show()
Individual CATE estimates are noisy. Policy-makers usually need group average treatment effects (GATEs), for example by unemployment duration:
groups = pd.cut(X_test.months_unemployed, [0, 3, 12, 100], labels=["< 3 months", "3–12 months", "> 12 months"])
gate = pd.DataFrame({"true": tau_test, "X-learner": cate_x, "group": groups}).groupby("group", observed=True).mean().round(0)
print(gate)
Exercise
Make the treated group small: change the intercept of the participation model so only about 10% are treated. Re-run the three learners. Which one degrades most, and why?
Show a solution
def pehe_with_share(intercept):
Xs, _, _, _ = make_data(6_000)
p = 1 / (1 + np.exp(-(intercept + 0.05 * Xs.months_unemployed)))
ds = rng.binomial(1, p).astype(float)
taus = 500 + 250 * np.minimum(Xs.months_unemployed, 24) - 150 * (Xs.years_education - 12)
ys = 15_000 + 400 * Xs.prior_earnings + rng.normal(0, 3_000, len(Xs)) + ds * taus
m1 = gbm().fit(Xs[ds == 1], ys[ds == 1]); m0 = gbm().fit(Xs[ds == 0], ys[ds == 0])
t_hat = m1.predict(X_test) - m0.predict(X_test)
t1 = gbm().fit(Xs[ds == 1], ys[ds == 1] - m0.predict(Xs[ds == 1]))
t0 = gbm().fit(Xs[ds == 0], m1.predict(Xs[ds == 0]) - ys[ds == 0])
x_hat = ds.mean() * t0.predict(X_test) + (1 - ds.mean()) * t1.predict(X_test)
return ds.mean(), np.sqrt(np.mean((t_hat - tau_test) ** 2)), np.sqrt(np.mean((x_hat - tau_test) ** 2))
print("share treated %.2f | PEHE T %.0f | PEHE X %.0f" % pehe_with_share(-2.8))
# The T-learner's treated model is fitted on few observations and is noisy; the X-learner borrows
# strength from the large control group and usually does better.
Checklist and recap
- I can define the CATE and explain why the ATE is not enough for targeting.
- I can implement S-, T- and X-learners.
- I summarise heterogeneity with GATEs rather than over-reading individual estimates.