Heterogeneous Effects and Meta-Learners

55.1 From one effect to one effect per customer (40 min)

The conditional average treatment effect is τ(x) = E[spend(1) − spend(0) | X = x]: the average effect for customers with features x. In the coupon simulator, τ ranges from about −$26 to +$69: deal-seekers respond strongly, app users more, big spenders less (they'd have bought anyway, and the coupon just gives away margin).

Knowing τ(x) answers the business question directly: coupon the customers with the largest τ(x). Three warnings before we start:

1. τ(x) is never observed for anyone, so no learner can be checked the way a prediction can. In real work, Day 57's validation on an experiment is the only check. Here the simulator lets us cheat and compare with the truth: PEHE, the root mean squared error against the true CATE.
2. CATE estimates are much noisier than an ATE: each one rests on far fewer comparable customers.
3. All Days 31–36 assumptions still apply. Heterogeneous effects estimated from confounded data are only as good as the confounders you measured.

55.2 Four meta-learners (60 min)

A meta-learner builds a CATE model out of ordinary ML models (Künzel et al. 2019; Nie & Wager 2021; Kennedy 2023):

Create day55/meta_learners.py. It fits on half the customers and scores on the other half:

# day55/meta_learners.py   (python -m day55.meta_learners)
"""
Who responds to the coupon? Four meta-learners (Kunzel et al., 2019; Kennedy, 2023), each building a CATE
model tau(x) out of ordinary ML models. Fit on half the customers, scored on the other half against the
simulator's true effect for each customer.
  S-learner:  one model of spend on (X, coupon); tau(x) = f(x, 1) - f(x, 0)
  T-learner:  separate models for couponed and non-couponed customers; tau(x) = f1(x) - f0(x)
  X-learner:  imputes each customer's effect from the OTHER group's model, then models those effects
  DR-learner: regresses the doubly robust (AIPW) score of Day 36 on X
"""
import numpy as np
import polars as pl
from econml.dr import DRLearner
from econml.metalearners import SLearner, TLearner, XLearner
from sklearn.model_selection import train_test_split

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

full = coupons().with_row_index("row")
train, test = train_test_split(full.to_pandas(), test_size=0.5, random_state=0)
Xtr, Xte = train[X_COLS], test[X_COLS]

models = {
    "S-learner": SLearner(overall_model=reg()),
    "T-learner": TLearner(models=reg()),
    "X-learner": XLearner(models=reg(), propensity_model=clf()),
    "DR-learner, flexible final model": DRLearner(model_propensity=clf(), model_regression=reg(),
                                                  model_final=reg(), cv=5, min_propensity=0.02, random_state=0),
    "DR-learner, simple final model": DRLearner(model_propensity=clf(), model_regression=reg(),
                                                model_final=reg(n_estimators=100, num_leaves=4, min_child_samples=300),
                                                cv=5, min_propensity=0.02, random_state=0),
}
rows, preds = [], {}
for name, m in models.items():
    m.fit(train.spend, train.coupon, X=Xtr)
    tau_hat = m.effect(Xte)
    preds[name] = tau_hat
    rows.append({"learner": name, "ATE": round(tau_hat.mean(), 2),
                 "PEHE (RMSE vs true CATE)": round(float(np.sqrt(np.mean((tau_hat - test._tau) ** 2))), 2),
                 "correlation with true CATE": round(float(np.corrcoef(tau_hat, test._tau)[0, 1]), 3)})
rows.append({"learner": "TRUTH", "ATE": round(test._tau.mean(), 2), "PEHE (RMSE vs true CATE)": 0.0,
             "correlation with true CATE": 1.0})
with pl.Config(tbl_width_chars=110, fmt_str_lengths=40):
    print(pl.DataFrame(rows))
print(f"True CATEs in the test half: SD {test._tau.std():.1f}, from {test._tau.min():.0f} to {test._tau.max():.0f}")
np.savez("data/day55_cate_preds.npz", rows=test.row.to_numpy(),
         **{k.split(",")[0].replace("-", "_"): v for k, v in preds.items() if "flexible" not in k})

python -m day55.meta_learners

│ learner                          ┆ ATE   ┆ PEHE (RMSE vs true CATE) ┆ correlation with true CATE │
│ S-learner                        ┆ 14.36 ┆ 7.4                      ┆ 0.867                      │
│ T-learner                        ┆ 15.86 ┆ 13.74                    ┆ 0.677                      │
│ X-learner                        ┆ 16.14 ┆ 8.28                     ┆ 0.826                      │
│ DR-learner, flexible final model ┆ 16.24 ┆ 16.45                    ┆ 0.62                       │
│ DR-learner, simple final model   ┆ 16.05 ┆ 6.17                     ┆ 0.896                      │
│ TRUTH                            ┆ 16.47 ┆ 0.0                      ┆ 1.0                        │
True CATEs in the test half: SD 13.8, from -26 to 69

• The S-learner ranks customers well (correlation 0.87) but its average effect is biased ($14.4): shrinkage again.
• The T-learner is the noisiest: each group's model makes its own errors, and the difference doubles them.
• The DR-learner with a flexible final model is the worst of all; with a simple one (100 shallow trees, at least 300 customers per leaf), it's the best. AIPW scores are unbiased but extremely noisy, so the last step must average a lot of them.

55.3 What each learner found (45 min)

Create day55/cate_plot.py:

# day55/cate_plot.py   (python -m day55.cate_plot)
"""What did each learner find? Predicted effect against deal_share (the main driver), with the truth.
Run day55.meta_learners first: it saves the test-half predictions."""
import matplotlib.pyplot as plt
import numpy as np

from causalkit.coupons import coupons

saved = np.load("data/day55_cate_preds.npz")
test = coupons().to_pandas().iloc[saved["rows"]]
fig, axes = plt.subplots(1, 4, figsize=(14, 3.6), sharey=True)
for ax, name in zip(axes, ["S_learner", "T_learner", "X_learner", "DR_learner"]):
    ax.scatter(test.deal_share, saved[name], s=2, alpha=0.25, color="#2563eb", label="predicted")
    ax.scatter(test.deal_share, test._tau, s=2, alpha=0.15, color="#94a3b8", label="true")
    ax.set_title(name.replace("_", "-"))
    ax.set_xlabel("deal_share")
axes[0].set_ylabel("effect of the coupon ($)")
axes[0].legend(markerscale=5)
fig.tight_layout()
fig.savefig("figures/day55_cate_by_deal_share.png", dpi=120)
print("Saved figures/day55_cate_by_deal_share.png")

python -m day55.cate_plot

Saved figures/day55_cate_by_deal_share.png

Open the figure: predicted effect against deal_share, with the truth in grey. The DR-learner and S-learner track the rising trend; the T-learner's predictions scatter far above and below it. The S-learner's cloud is visibly flatter than the truth: shrinkage again.

55.4 How much data do they need? (40 min)

Try it. Rerun the four learners with 4,000 customers instead of 20,000 (coupons(n=4_000), half for training). Which learner suffers most? Create day55/small_sample.py.

# day55/small_sample.py   (python -m day55.small_sample)
import numpy as np
from econml.dr import DRLearner
from econml.metalearners import SLearner, TLearner, XLearner
from sklearn.model_selection import train_test_split

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

for n in [4_000, 20_000]:
    full = coupons(n=n).to_pandas()
    tr, te = train_test_split(full, test_size=0.5, random_state=0)
    models = {"S": SLearner(overall_model=reg()), "T": TLearner(models=reg()),
              "X": XLearner(models=reg(), propensity_model=clf()),
              "DR": DRLearner(model_propensity=clf(), model_regression=reg(),
                              model_final=reg(n_estimators=100, num_leaves=4, min_child_samples=300), cv=5,
                              min_propensity=0.02, random_state=0)}
    out = []
    for name, m in models.items():
        m.fit(tr.spend, tr.coupon, X=tr[X_COLS])
        out.append(f"{name} {np.sqrt(np.mean((m.effect(te[X_COLS]) - te._tau) ** 2)):.1f}")
    print(f"n = {n:>6,}: PEHE  " + "   ".join(out))

n =  4,000: PEHE  S 11.6   T 23.2   X 15.2   DR 16.9
n = 20,000: PEHE  S 7.4   T 13.7   X 8.3   DR 6.2

55.5 Day 55 mini-project: who should get next month's coupons? (65 min)

Next month's budget covers 30% of customers. Rank the test-half customers by each model, coupon the top 30%, and score each policy with the true effects. Create day55/targeting_value.py:

# day55/targeting_value.py   (python -m day55.targeting_value)
"""
Mini-project: the reason to estimate CATEs is to decide WHO gets a coupon. Next month's budget covers
30% of customers. Rank the test-half customers by each learner's predicted effect, coupon the top 30%,
and score the policy with the simulator's true effects. Compare with coupons at random and with the
current targeting model (which ranks customers by its own score, the propensity).
Run day55.meta_learners first.
"""
import numpy as np
import polars as pl

from causalkit.coupons import coupons

saved = np.load("data/day55_cate_preds.npz")
test = coupons().to_pandas().iloc[saved["rows"]]
k = int(0.3 * len(test))
tau = test._tau.to_numpy()
scores = {"random": np.random.default_rng(0).random(len(test)), "current targeting model": test._ps.to_numpy()}
scores |= {name.replace("_", "-"): saved[name] for name in ["S_learner", "T_learner", "X_learner", "DR_learner"]}
scores["oracle (true effects)"] = tau
rows = []
for name, s in scores.items():
    top = np.argsort(-s)[:k]
    rows.append({"policy": name, "effect per coupon": round(tau[top].mean(), 2),
                 "extra $ per 1,000 coupons": round(1000 * tau[top].mean()),
                 "negative effects": f"{(tau[top] < 0).mean():.1%}"})
with pl.Config(tbl_width_chars=120, fmt_str_lengths=30):
    print(pl.DataFrame(rows))

python -m day55.targeting_value

│ policy                  ┆ effect per coupon ┆ extra $ per 1,000 coupons ┆ negative effects │
│ random                  ┆ 16.27             ┆ 16270                     ┆ 10.4%            │
│ current targeting model ┆ 26.67             ┆ 26665                     ┆ 3.3%             │
│ S-learner               ┆ 31.3              ┆ 31295                     ┆ 0.0%             │
│ T-learner               ┆ 28.11             ┆ 28107                     ┆ 0.5%             │
│ X-learner               ┆ 30.5              ┆ 30501                     ┆ 0.2%             │
│ DR-learner              ┆ 31.74             ┆ 31739                     ┆ 0.0%             │
│ oracle (true effects)   ┆ 32.97             ┆ 32973                     ┆ 0.0%             │

Targeting on estimated effects raises the return per coupon by about 18% over the current model (which targets the customers most likely to get a coupon, not the ones who respond most), and gets within 5% of the oracle. Notice the S-learner: biased on average, but nearly as good as the best for ranking. For targeting, the order matters more than the level.

55.6 Save your work with Git (20 min)

git add .
git commit -m "Day 55: S, T, X and DR meta-learners, PEHE, targeting policy"
git push

Day 55 checklist

• I can define the CATE and explain why it can't be checked directly
• I can fit S-, T-, X- and DR-learners with EconML and explain each one's weakness
• I know that the DR-learner's final model must be simple
• I can turn CATE estimates into a targeting policy and value it