Causal Inference I(i) — A/B Testing: ML-Based Variance Reduction (CUPAC)¶
Beyond CUPED — using a machine-learning prediction as the control variate¶
CUPED (notebook 1b) cut experiment variance using one pre-experiment covariate, with the reduction capped at $\rho^2$ — the squared correlation between the outcome and that covariate. But an outcome usually depends on many pre-experiment features, often nonlinearly, and a single covariate leaves most of that predictable structure on the table. CUPAC (Control Using Predictions As Covariates; Tang et al., DoorDash 2020) closes the gap: train a machine-learning model on pre-experiment features to predict the outcome, and use that prediction $\hat Y$ as the control variate. Because $\hat Y$ can track the outcome far better than any single feature, the variance reduction — still exactly $\rho^2(Y,\hat Y)$ — is much larger.
This notebook builds it:
- Control variates and CUPED — the general principle (subtract $\beta$ times a mean-zero, treatment-independent covariate) and the $\rho^2$ ceiling with one feature.
- CUPAC — replace the single covariate with a cross-fitted ML prediction built from all pre-experiment features; the variance reduction jumps, and the estimate stays unbiased.
- The one rule that keeps it honest — the control variate must be a function of pre-treatment data only, and cross-fitting prevents overfitting bias.
Simulation-first with a known effect. Python leads (from-scratch CUPED, CUPAC via cross-fitted gradient boosting, coverage check); the R companion mirrors it. This ties the experimentation arc to the ML arc (the predictor is any regressor) and to Double ML (the same residualize-with-ML idea).
1. Control variates and the CUPED ceiling¶
A control variate is any pre-experiment quantity $C$ that (i) is correlated with the outcome and (ii) is independent of the treatment (so subtracting it can't bias the effect). Replacing $Y$ with $Y-\beta\,(C-\bar C)$, where $\beta=\operatorname{Cov}(Y,C)/\operatorname{Var}(C)$, leaves the treatment effect unchanged in expectation but shrinks its variance by a factor $\rho^2(Y,C)$. CUPED uses a single pre-experiment covariate as $C$.
We simulate an outcome that depends nonlinearly on six pre-experiment features plus the treatment. The best single feature correlates only moderately with the outcome, so CUPED captures a modest slice of the variance — its $\rho^2$ ceiling.
import numpy as np, matplotlib.pyplot as plt, warnings
from sklearn.ensemble import GradientBoostingRegressor as GBR
from sklearn.model_selection import cross_val_predict
warnings.filterwarnings("ignore")
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; GREY="#a0aec0"; PURP="#6b46c1"
def gen(seed, n=6000, delta=0.5):
r=np.random.default_rng(seed); Z=r.normal(0,1,(n,6))
gZ=2*np.sin(Z[:,0]*Z[:,1])+Z[:,2]**2+1.5*Z[:,3]+0.8*Z[:,4]*Z[:,5] # nonlinear prognostic function
T=r.binomial(1,0.5,n); Y=gZ+delta*T+r.normal(0,1,n)
return Z,T,Y
Z,T,Y=gen(0); delta=0.5
def eff_se(Yv):
d=Yv[T==1].mean()-Yv[T==0].mean(); se=np.sqrt(Yv[T==1].var(ddof=1)/(T==1).sum()+Yv[T==0].var(ddof=1)/(T==0).sum()); return d,se
d0,se0=eff_se(Y)
cors=[abs(np.corrcoef(Y,Z[:,j])[0,1]) for j in range(6)]; best=int(np.argmax(cors))
theta=np.cov(Y,Z[:,best])[0,1]/np.var(Z[:,best],ddof=1)
Ycuped=Y-theta*(Z[:,best]-Z[:,best].mean()); d1,se1=eff_se(Ycuped)
print(f"true effect delta = {delta}")
print(f" unadjusted : effect {d0:.3f} SE {se0:.4f}")
print(f" CUPED (best 1 covariate, feature {best}): effect {d1:.3f} SE {se1:.4f} var-reduction {100*(1-(se1/se0)**2):.0f}% (rho^2 = {cors[best]**2:.2f})")
fig,ax=plt.subplots(1,2,figsize=(13,4.2))
ax[0].bar(range(6),[c**2 for c in cors],color=BLUE); ax[0].set_xlabel("pre-experiment feature"); ax[0].set_ylabel("rho^2 with outcome")
ax[0].set_title("No single feature explains much of the outcome")
ax[1].scatter(Z[:,best],Y,s=5,alpha=.2,color=GREY); ax[1].set_xlabel(f"best single feature (Z{best})"); ax[1].set_ylabel("outcome Y")
ax[1].set_title(f"CUPED uses one covariate -> rho^2={cors[best]**2:.2f} ceiling")
plt.tight_layout(); plt.show()
print("The outcome is driven by nonlinear combinations of several features, so any ONE covariate leaves most of the")
print("predictable variance unexplained -- and CUPED's variance reduction is capped at that single-covariate rho^2.")
true effect delta = 0.5 unadjusted : effect 0.427 SE 0.0697 CUPED (best 1 covariate, feature 3): effect 0.512 SE 0.0567 var-reduction 34% (rho^2 = 0.33)
The outcome is driven by nonlinear combinations of several features, so any ONE covariate leaves most of the predictable variance unexplained -- and CUPED's variance reduction is capped at that single-covariate rho^2.
2. CUPAC — an ML prediction as the control variate¶
CUPAC replaces the single covariate with $\hat Y$, an ML model's prediction of the outcome from all pre-experiment features. The control-variate math is identical — $Y-\beta(\hat Y-\bar{\hat Y})$ with $\beta=\operatorname{Cov}(Y,\hat Y)/\operatorname{Var}(\hat Y)$ — but now $\rho^2(Y,\hat Y)$ can be large because a flexible learner captures the nonlinear structure a single covariate misses. The prediction is built by cross-fitting (out-of-fold predictions), which keeps $\hat Y$ from overfitting its own outcome and so preserves unbiasedness. Here CUPAC more than doubles CUPED's variance reduction while returning the same effect estimate.
Yhat=cross_val_predict(GBR(n_estimators=300,max_depth=3,learning_rate=0.05), Z, Y, cv=5) # cross-fitted ML prediction
beta=np.cov(Y,Yhat)[0,1]/np.var(Yhat,ddof=1)
Ycupac=Y-beta*(Yhat-Yhat.mean()); d2,se2=eff_se(Ycupac)
rho2_ml=np.corrcoef(Y,Yhat)[0,1]**2
print(f"true effect delta = {delta}")
print(f" unadjusted : effect {d0:.3f} SE {se0:.4f}")
print(f" CUPED (1 covar): effect {d1:.3f} SE {se1:.4f} var-reduction {100*(1-(se1/se0)**2):.0f}% (rho^2 {cors[best]**2:.2f})")
print(f" CUPAC (ML pred): effect {d2:.3f} SE {se2:.4f} var-reduction {100*(1-(se2/se0)**2):.0f}% (rho^2 {rho2_ml:.2f})")
print(f" -> all estimates ~ {delta} (unbiased); CUPAC cuts the SE by {100*(1-se2/se0):.0f}% because the ML prediction tracks Y far better")
fig,ax=plt.subplots(1,2,figsize=(13,4.4))
ax[0].scatter(Yhat,Y,s=5,alpha=.2,color=GREEN); lim=[Y.min(),Y.max()]; ax[0].plot(lim,lim,"k--",lw=1)
ax[0].set_xlabel("ML prediction Yhat (pre-experiment features)"); ax[0].set_ylabel("outcome Y"); ax[0].set_title(f"CUPAC control variate: rho^2(Y,Yhat)={rho2_ml:.2f}")
red=[0,100*(1-(se1/se0)**2),100*(1-(se2/se0)**2)]; rho=[0,cors[best]**2*100,rho2_ml*100]
x=np.arange(3); w=0.38
ax[1].bar(x-w/2,rho,w,color=GREY,label="rho^2 (predicted)"); ax[1].bar(x+w/2,red,w,color=GREEN,label="actual var-reduction")
ax[1].set_xticks(x); ax[1].set_xticklabels(["unadjusted","CUPED\n(1 covariate)","CUPAC\n(ML)"],fontsize=8)
ax[1].set_ylabel("% variance reduction"); ax[1].set_title("Variance reduction = rho^2; ML lifts the ceiling"); ax[1].legend(fontsize=8)
plt.tight_layout(); plt.show()
print("CUPAC = CUPED with an ML-learned control variate. The variance reduction still equals rho^2(Y,Yhat), but a good")
print("predictor pushes that far above any single covariate -- the same 'residualize with ML' idea as Double ML (subsection 9).")
true effect delta = 0.5 unadjusted : effect 0.427 SE 0.0697 CUPED (1 covar): effect 0.512 SE 0.0567 var-reduction 34% (rho^2 0.33) CUPAC (ML pred): effect 0.522 SE 0.0352 var-reduction 75% (rho^2 0.74) -> all estimates ~ 0.5 (unbiased); CUPAC cuts the SE by 50% because the ML prediction tracks Y far better
CUPAC = CUPED with an ML-learned control variate. The variance reduction still equals rho^2(Y,Yhat), but a good predictor pushes that far above any single covariate -- the same 'residualize with ML' idea as Double ML (subsection 9).
3. The rule that keeps it honest — pre-treatment features and cross-fitting¶
Variance reduction is free precision only if it introduces no bias, which requires one discipline: the control variate must be a function of pre-treatment data alone. If $\hat Y$ used any post-treatment signal (or the treatment indicator), subtracting it would soak up part of the effect itself and bias the estimate toward zero — the mediator trap of the DAG notebook, in disguise. Two safeguards: (i) build $\hat Y$ only from features fixed before assignment, and (ii) use cross-fitting so the prediction for each unit is made by a model that never saw that unit's outcome (preventing overfit-induced bias). We confirm across repeated experiments that both CUPED and CUPAC are unbiased with correct 95% coverage — CUPAC simply has a much tighter sampling distribution.
def run_once(seed):
Z,T,Y=gen(seed)
se=lambda Yv: np.sqrt(Yv[T==1].var(ddof=1)/(T==1).sum()+Yv[T==0].var(ddof=1)/(T==0).sum())
d_un=Y[T==1].mean()-Y[T==0].mean()
j=int(np.argmax([abs(np.corrcoef(Y,Z[:,k])[0,1]) for k in range(6)])); th=np.cov(Y,Z[:,j])[0,1]/np.var(Z[:,j],ddof=1)
Yc=Y-th*(Z[:,j]-Z[:,j].mean()); d_cu=Yc[T==1].mean()-Yc[T==0].mean()
Yh=cross_val_predict(GBR(n_estimators=150,max_depth=3,learning_rate=0.08),Z,Y,cv=3); b=np.cov(Y,Yh)[0,1]/np.var(Yh,ddof=1)
Ya=Y-b*(Yh-Yh.mean()); d_ca=Ya[T==1].mean()-Ya[T==0].mean()
return (d_un,se(Y)),(d_cu,se(Yc)),(d_ca,se(Ya))
res=[run_once(s) for s in range(120)]
for i,nm in enumerate(["unadjusted","CUPED","CUPAC"]):
ds=np.array([r[i][0] for r in res]); ses=np.array([r[i][1] for r in res])
cov=np.mean(np.abs(ds-delta)<1.96*ses)
print(f" {nm:11s}: mean effect {ds.mean():.3f} (bias {ds.mean()-delta:+.3f}), SD {ds.std():.4f}, 95% coverage {cov:.2f}")
print(f" -> all unbiased and calibrated; CUPAC's SD is smallest (most precise) -- variance reduction without bias.")
fig,ax=plt.subplots(figsize=(8,4.2))
for i,(nm,c) in enumerate(zip(["unadjusted","CUPED","CUPAC"],[GREY,ORANGE,GREEN])):
ds=np.array([r[i][0] for r in res]); ax.hist(ds,bins=22,alpha=.6,color=c,label=f"{nm} (SD {ds.std():.3f})")
ax.axvline(delta,color=RED,lw=2,ls="--",label=f"true effect {delta}")
ax.set_xlabel("estimated effect across experiments"); ax.set_ylabel("count"); ax.set_title("All unbiased; CUPAC has the tightest sampling distribution"); ax.legend(fontsize=8)
plt.tight_layout(); plt.show()
print("Because the control variate uses only pre-treatment features (and is cross-fitted), CUPAC adds precision with no bias --")
print("the estimate still centers on the true effect, just with far less spread.")
unadjusted : mean effect 0.504 (bias +0.004), SD 0.0677, 95% coverage 0.94 CUPED : mean effect 0.504 (bias +0.004), SD 0.0606, 95% coverage 0.94 CUPAC : mean effect 0.502 (bias +0.002), SD 0.0402, 95% coverage 0.94 -> all unbiased and calibrated; CUPAC's SD is smallest (most precise) -- variance reduction without bias.
Because the control variate uses only pre-treatment features (and is cross-fitted), CUPAC adds precision with no bias -- the estimate still centers on the true effect, just with far less spread.
4. Summary¶
CUPAC generalizes CUPED from one covariate to a machine-learning prediction of the outcome, lifting the variance-reduction ceiling:
- Any pre-treatment control variate cuts the effect's variance by $\rho^2(Y,C)$; CUPED uses one covariate and is capped at that single-feature correlation (here ~34%).
- CUPAC uses a cross-fitted ML prediction from all pre-experiment features as the control variate; because the model captures the nonlinear prognostic structure, $\rho^2(Y,\hat Y)$ is much larger and the variance reduction more than doubled (~75%) — the same effect estimate, a far tighter interval.
- It stays honest under one rule: the control variate must be built from pre-treatment features only, with cross-fitting to avoid overfit bias. A repeated-experiment check confirmed both CUPED and CUPAC are unbiased with correct coverage; CUPAC just has the smallest spread.
Practical guidance: when you have rich pre-experiment data, use CUPAC — train any good regressor on pre-period features to predict the metric, cross-fit it, and use the prediction as the CUPED covariate; it is the cheapest large power gain available, stacking naturally with a fixed sample size or sequential design. Cross-links: the predictor is any model from the ML arc (trees/boosting/nets); the "residualize the outcome with an ML model" logic is exactly Double Machine Learning (subsection 9) and the R-learner (subsection 9, meta-learners), here used for precision rather than identification; and it extends CUPED / covariate adjustment (1b). The R companion reproduces CUPED and CUPAC with a gradient-boosting control variate.