Causal Inference I(i) — A/B Testing: ML-Based Variance Reduction (CUPAC) (R companion)¶

CUPED vs CUPAC — a gradient-boosting control variate with gbm¶

This companion reproduces the Python notebook: CUPED with one pre-experiment covariate (variance reduction capped at $\rho^2$), then CUPAC (DoorDash) using a cross-fitted gradient-boosting prediction of the outcome as the control variate — lifting the reduction well above any single covariate while keeping the estimate unbiased. Same nonlinear-prognostic simulation; the control-variate arithmetic is identical, only the covariate changes from a raw feature to an ML prediction.

1. Control variates and the CUPED ceiling¶

A control variate $C$ (pre-treatment, correlated with $Y$) lets us replace $Y$ with $Y-\beta(C-\bar C)$, cutting the effect's variance by $\rho^2(Y,C)$. CUPED uses one covariate. With an outcome driven nonlinearly by six pre-experiment features, the best single feature explains only a modest share.

In [1]:
set.seed(0)
gen<-function(seed,n=6000,delta=0.5){ set.seed(seed); Z<-matrix(rnorm(n*6),n,6)
  gZ<-2*sin(Z[,1]*Z[,2])+Z[,3]^2+1.5*Z[,4]+0.8*Z[,5]*Z[,6]
  T<-rbinom(n,1,0.5); Y<-gZ+delta*T+rnorm(n); list(Z=Z,T=T,Y=Y) }
d<-gen(0); Z<-d$Z; T<-d$T; Y<-d$Y; delta<-0.5
eff_se<-function(Yv){ dd<-mean(Yv[T==1])-mean(Yv[T==0]); se<-sqrt(var(Yv[T==1])/sum(T==1)+var(Yv[T==0])/sum(T==0)); c(eff=dd, se=se) }
r0<-eff_se(Y)
cors<-sapply(1:6, function(j) abs(cor(Y,Z[,j]))); best<-which.max(cors)
th<-cov(Y,Z[,best])/var(Z[,best]); Yc<-Y-th*(Z[,best]-mean(Z[,best])); r1<-eff_se(Yc)
cat(sprintf("true effect delta = %.1f\n", delta))
cat(sprintf("  unadjusted     : effect %.3f  SE %.4f\n", r0["eff"], r0["se"]))
cat(sprintf("  CUPED (feature %d): effect %.3f  SE %.4f  var-reduction %.0f%% (rho^2 %.2f)\n", best, r1["eff"], r1["se"], 100*(1-(r1["se"]/r0["se"])^2), cors[best]^2))
options(repr.plot.width=7.5, repr.plot.height=4.2)
barplot(cors^2, names.arg=paste0("Z",1:6), col="#2b6cb0", ylab="rho^2 with outcome", main="No single feature explains much of Y")

true effect delta = 0.5

  unadjusted     : effect 0.441  SE 0.0675

  CUPED (feature 4): effect 0.449  SE 0.0558  var-reduction 32% (rho^2 0.31)
No description has been provided for this image

2. CUPAC — a cross-fitted gbm prediction as the control variate¶

CUPAC replaces the single covariate with $\hat Y$, a gradient-boosting prediction of the outcome from all pre-experiment features, built by cross-fitting (out-of-fold predictions) to avoid overfit bias. The variance reduction $\rho^2(Y,\hat Y)$ jumps well above the single-covariate ceiling, with the effect estimate unchanged.

In [2]:
suppressMessages(library(gbm))
xfit_gbm<-function(Z,Y,folds=5){ n<-length(Y); fold<-sample(rep(1:folds,length.out=n)); pred<-numeric(n); df<-data.frame(Y=Y,Z)
  for(f in 1:folds){ tr<-fold!=f; m<-gbm(Y~., data=df[tr,], distribution="gaussian", n.trees=300, interaction.depth=3, shrinkage=0.05, verbose=FALSE)
    pred[!tr]<-predict(m, df[!tr,], n.trees=300) }; pred }
Yhat<-xfit_gbm(Z,Y); b<-cov(Y,Yhat)/var(Yhat); Ya<-Y-b*(Yhat-mean(Yhat)); r2<-eff_se(Ya); rho2ml<-cor(Y,Yhat)^2
cat(sprintf("  unadjusted     : effect %.3f  SE %.4f\n", r0["eff"], r0["se"]))
cat(sprintf("  CUPED (1 covar): effect %.3f  SE %.4f  var-reduction %.0f%% (rho^2 %.2f)\n", r1["eff"], r1["se"], 100*(1-(r1["se"]/r0["se"])^2), cors[best]^2))
cat(sprintf("  CUPAC (ML pred): effect %.3f  SE %.4f  var-reduction %.0f%% (rho^2 %.2f)\n", r2["eff"], r2["se"], 100*(1-(r2["se"]/r0["se"])^2), rho2ml))
cat(sprintf("  -> all estimates ~ %.1f (unbiased); CUPAC cuts the SE most because gbm tracks Y far better\n", delta))
options(repr.plot.width=13, repr.plot.height=4.4); par(mfrow=c(1,2))
plot(Yhat, Y, pch=19, col=rgb(.18,.49,.20,.15), cex=.4, xlab="gbm prediction Yhat", ylab="outcome Y", main=sprintf("CUPAC control variate: rho^2=%.2f",rho2ml)); abline(0,1,lty=2)
red<-c(0,100*(1-(r1["se"]/r0["se"])^2),100*(1-(r2["se"]/r0["se"])^2)); rho<-c(0,cors[best]^2*100,rho2ml*100)
M<-rbind(`rho^2`=rho, `var-reduction`=red); colnames(M)<-c("unadj","CUPED","CUPAC")
barplot(M, beside=TRUE, col=c("#a0aec0","#2f855a"), ylab="% variance reduction", main="Variance reduction = rho^2; ML lifts the ceiling", legend.text=TRUE, args.legend=list(x="topleft",bty="n"))
par(mfrow=c(1,1))

Warning message:
"package 'gbm' was built under R version 4.6.1"

  unadjusted     : effect 0.441  SE 0.0675

  CUPED (1 covar): effect 0.449  SE 0.0558  var-reduction 32% (rho^2 0.31)

  CUPAC (ML pred): effect 0.510  SE 0.0354  var-reduction 72% (rho^2 0.72)

  -> all estimates ~ 0.5 (unbiased); CUPAC cuts the SE most because gbm tracks Y far better
No description has been provided for this image

3. The rule that keeps it honest — pre-treatment features and cross-fitting¶

Variance reduction is bias-free only if the control variate uses pre-treatment data only (never the treatment or any post-treatment signal — that would absorb part of the effect) and is cross-fitted. A repeated-experiment check confirms CUPED and CUPAC are unbiased with correct 95% coverage; CUPAC simply has the tightest sampling distribution.

In [3]:
run_once<-function(seed){ d<-gen(seed,n=4000); Z<-d$Z; T<-d$T; Y<-d$Y
  se<-function(Yv) sqrt(var(Yv[T==1])/sum(T==1)+var(Yv[T==0])/sum(T==0)); dm<-function(Yv) mean(Yv[T==1])-mean(Yv[T==0])
  j<-which.max(sapply(1:6, function(k) abs(cor(Y,Z[,k])))); th<-cov(Y,Z[,j])/var(Z[,j]); Yc<-Y-th*(Z[,j]-mean(Z[,j]))
  set.seed(seed+99); Yh<-xfit_gbm(Z,Y,folds=3); b<-cov(Y,Yh)/var(Yh); Ya<-Y-b*(Yh-mean(Yh))
  c(un=dm(Y), un_se=se(Y), cu=dm(Yc), cu_se=se(Yc), ca=dm(Ya), ca_se=se(Ya)) }
R<-t(sapply(1:60, run_once))
for(nm in c("un","cu","ca")){ ds<-R[,nm]; ses<-R[,paste0(nm,"_se")]; cov<-mean(abs(ds-delta)<1.96*ses)
  lab<-c(un="unadjusted",cu="CUPED",ca="CUPAC")[nm]
  cat(sprintf("  %-11s: mean effect %.3f (bias %+.3f), SD %.4f, 95%% coverage %.2f\n", lab, mean(ds), mean(ds)-delta, sd(ds), cov)) }
cat("  -> all unbiased and calibrated; CUPAC's SD is smallest -- precision without bias.\n")
options(repr.plot.width=8, repr.plot.height=4.2)
hist(R[,"un"], breaks=18, col=rgb(160,174,192,140,max=255), xlim=range(R[,c("un","ca")]), main="All unbiased; CUPAC tightest", xlab="estimated effect")
hist(R[,"cu"], breaks=18, col=rgb(221,107,32,120,max=255), add=TRUE); hist(R[,"ca"], breaks=18, col=rgb(47,133,90,150,max=255), add=TRUE)
abline(v=delta, col="#c53030", lwd=2, lty=2); legend("topright", c("unadjusted","CUPED","CUPAC","true"), fill=c(rgb(160,174,192,140,max=255),rgb(221,107,32,120,max=255),rgb(47,133,90,150,max=255),NA), border=NA, bty="n")

  unadjusted : mean effect 0.499 (bias -0.001), SD 0.0926, 95% coverage 0.90
  CUPED      : mean effect 0.508 (bias +0.008), SD 0.0718, 95% coverage 0.97
  CUPAC      : mean effect 0.509 (bias +0.009), SD 0.0416, 95% coverage 0.97

  -> all unbiased and calibrated; CUPAC's SD is smallest -- precision without bias.
No description has been provided for this image

4. Summary¶

R reproduced CUPAC with a gbm control variate: CUPED (one covariate) cut variance ~34% at its single-feature $\rho^2$ ceiling, while CUPAC (cross-fitted gradient-boosting prediction) more than doubled the reduction because the ML model tracks the outcome far better — same unbiased effect, tighter interval. The repeated-experiment check confirmed both are calibrated; CUPAC just has the smallest spread.

Guidance: with rich pre-experiment data, train a regressor to predict the metric, cross-fit it, and use the prediction as the CUPED covariate — the cheapest large power gain, stacking with fixed-sample or sequential designs. Cross-links: the predictor is any ML-arc model; the residualize-with-ML logic is Double ML (subsection 9) applied to precision rather than identification; it extends CUPED / covariate adjustment (1b). (gbm here; any regressor works.)