Variance Reduction Comparison: No Covariates vs Raw Covariates vs CUPAC vs MLRATE¶

In this example target depends on two covariates in different ways:

• x1 — linearly related to target.
• x2 — non-linearly related to target (via x2^2).

We compare four ways of using this information:

1. No covariates — plain difference-in-means, x1/x2 are ignored entirely.
2. Raw covariates (no ML) — add x1 and x2 directly as regression covariates, no model involved. OLS can capture the linear x1 relationship, but not the non-linear x2 one.
3. CUPAC — fit a model on pre-experiment data to predict target from x1/x2, then use the prediction as a covariate on experiment data.
4. MLRATE — K-fold cross-fit a model on the experiment data itself (Guo et al., NeurIPS 2021) to predict target from x1/x2; no pre-experiment data required.

Expectation: since CUPAC and MLRATE can both capture the full (linear + non-linear) relationship via a gradient-boosted model, while the raw covariate can only capture the linear half, we should see CUPAC ≈ MLRATE >> raw covariates >> no covariates.

In [1]:

import numpy as np
import pandas as pd
import plotnine as p9
from sklearn.ensemble import HistGradientBoostingRegressor

from cluster_experiments import (
    ConstantPerturbator,
    NonClusteredSplitter,
    OLSAnalysis,
    PowerAnalysis,
)

import numpy as np import pandas as pd import plotnine as p9 from sklearn.ensemble import HistGradientBoostingRegressor from cluster_experiments import ( ConstantPerturbator, NonClusteredSplitter, OLSAnalysis, PowerAnalysis, )

Data generation¶

target = 3 * x1 + 2 * x2^2 + noise: x1 enters linearly, x2 enters only through its square. x2 itself has (close to) zero linear correlation with target, so a linear model can't pick up its contribution at all — only a model flexible enough to learn x2^2 can.

We keep a pre-experiment slice — only needed by CUPAC, to fit its model — and an experiment slice used by all four methods for the actual power simulation.

In [2]:

np.random.seed(2026)

N = 4_000
x1 = np.random.normal(size=N)
x2 = np.random.normal(size=N)
target = 3 * x1 + 2 * x2**2 + np.random.normal(scale=2.6, size=N)
df = pd.DataFrame({"x1": x1, "x2": x2, "target": target})

is_pre_experiment = np.random.rand(N) < 0.4
df_pre = df[is_pre_experiment].reset_index(drop=True)
df_analysis = df[~is_pre_experiment].reset_index(drop=True)

print(f"{len(df_pre) = }, {len(df_analysis) = }")
df_analysis.head()

np.random.seed(2026) N = 4_000 x1 = np.random.normal(size=N) x2 = np.random.normal(size=N) target = 3 * x1 + 2 * x2**2 + np.random.normal(scale=2.6, size=N) df = pd.DataFrame({"x1": x1, "x2": x2, "target": target}) is_pre_experiment = np.random.rand(N) < 0.4 df_pre = df[is_pre_experiment].reset_index(drop=True) df_analysis = df[~is_pre_experiment].reset_index(drop=True) print(f"{len(df_pre) = }, {len(df_analysis) = }") df_analysis.head()

len(df_pre) = 1604, len(df_analysis) = 2396

Out[2]:

	x1	x2	target
0	-0.013235	0.344164	0.655488
1	1.449708	-0.093372	5.541413
2	-0.829896	1.096221	-3.258865
3	-1.596159	1.840610	1.499819
4	0.613438	-0.532967	1.634634

1. No covariates¶

Plain difference-in-means: x1/x2 are not used at all.

In [3]:

perturbator = ConstantPerturbator(average_effect=0.38)
splitter = NonClusteredSplitter()

pw_none = PowerAnalysis(
    perturbator=perturbator,
    splitter=splitter,
    analysis=OLSAnalysis(),
    n_simulations=200,
    seed=2026,
)
power_none = pw_none.power_analysis(df_analysis)
print(f"No covariates: {power_none = }")

perturbator = ConstantPerturbator(average_effect=0.38) splitter = NonClusteredSplitter() pw_none = PowerAnalysis( perturbator=perturbator, splitter=splitter, analysis=OLSAnalysis(), n_simulations=200, seed=2026, ) power_none = pw_none.power_analysis(df_analysis) print(f"No covariates: {power_none = }")

No covariates: power_none = np.float64(0.475)

2. Raw covariates, no ML¶

x1 and x2 are included directly as regression covariates. No model is fit anywhere: OLS picks up the linear x1 effect just fine, but has no way to represent x2^2, so it gets none of the benefit x2 actually carries.

In [4]:

pw_raw = PowerAnalysis(
    perturbator=perturbator,
    splitter=splitter,
    analysis=OLSAnalysis(covariates=["x1", "x2"]),
    n_simulations=200,
    seed=2026,
)
power_raw = pw_raw.power_analysis(df_analysis)
print(f"Raw covariates (no ML): {power_raw = }")

pw_raw = PowerAnalysis( perturbator=perturbator, splitter=splitter, analysis=OLSAnalysis(covariates=["x1", "x2"]), n_simulations=200, seed=2026, ) power_raw = pw_raw.power_analysis(df_analysis) print(f"Raw covariates (no ML): {power_raw = }")

Raw covariates (no ML): power_raw = np.float64(0.675)

3. CUPAC¶

Fit a HistGradientBoostingRegressor on pre-experiment data (df_pre) to predict target from x1 and x2, then use that prediction (estimate_target) as the regression covariate on experiment data. A GBM can capture both the linear x1 term and the x2^2 shape.

In [5]:

pw_cupac = PowerAnalysis(
    perturbator=perturbator,
    splitter=splitter,
    analysis=OLSAnalysis(covariates=["estimate_target"]),
    cupac_model=HistGradientBoostingRegressor(),
    features_cupac_model=["x1", "x2"],
    n_simulations=200,
    seed=2026,
)
power_cupac = pw_cupac.power_analysis(df_analysis, df_pre)
print(f"CUPAC: {power_cupac = }")

pw_cupac = PowerAnalysis( perturbator=perturbator, splitter=splitter, analysis=OLSAnalysis(covariates=["estimate_target"]), cupac_model=HistGradientBoostingRegressor(), features_cupac_model=["x1", "x2"], n_simulations=200, seed=2026, ) power_cupac = pw_cupac.power_analysis(df_analysis, df_pre) print(f"CUPAC: {power_cupac = }")

CUPAC: power_cupac = np.float64(0.905)

4. MLRATE¶

K-fold cross-fit a HistGradientBoostingRegressor on the experiment data itself — each row's prediction comes from a fold that never trained on it, so there's no pre-experiment data requirement and no overfitting bias.

In [6]:

pw_mlrate = PowerAnalysis(
    perturbator=perturbator,
    splitter=splitter,
    analysis=OLSAnalysis(covariates=["estimate_target"]),
    cupac_model=HistGradientBoostingRegressor(),
    ml_option="mlrate",
    features_cupac_model=["x1", "x2"],
    n_simulations=200,
    seed=2026,
)
# Note: no pre_experiment_df passed in — MLRATE never needs it.
power_mlrate = pw_mlrate.power_analysis(df_analysis)
print(f"MLRATE: {power_mlrate = }")

pw_mlrate = PowerAnalysis( perturbator=perturbator, splitter=splitter, analysis=OLSAnalysis(covariates=["estimate_target"]), cupac_model=HistGradientBoostingRegressor(), ml_option="mlrate", features_cupac_model=["x1", "x2"], n_simulations=200, seed=2026, ) # Note: no pre_experiment_df passed in — MLRATE never needs it. power_mlrate = pw_mlrate.power_analysis(df_analysis) print(f"MLRATE: {power_mlrate = }")

MLRATE: power_mlrate = np.float64(0.935)

Comparison¶

In [7]:

results = pd.DataFrame(
    {
        "method": ["No covariates", "Raw covariates\n(no ML)", "CUPAC", "MLRATE"],
        "power": [power_none, power_raw, power_cupac, power_mlrate],
    }
)
results

results = pd.DataFrame( { "method": ["No covariates", "Raw covariates\n(no ML)", "CUPAC", "MLRATE"], "power": [power_none, power_raw, power_cupac, power_mlrate], } ) results

Out[7]:

	method	power
0	No covariates	0.475
1	Raw covariates\n(no ML)	0.675
2	CUPAC	0.905
3	MLRATE	0.935

In [8]:

(
    p9.ggplot(results, p9.aes(x="method", y="power"))
    + p9.geom_col(fill="#4a86e8")
    + p9.theme_minimal()
    + p9.labs(x="", y="Power", title="Power by variance-reduction strategy")
)

( p9.ggplot(results, p9.aes(x="method", y="power")) + p9.geom_col(fill="#4a86e8") + p9.theme_minimal() + p9.labs(x="", y="Power", title="Power by variance-reduction strategy") )

Takeaways¶

With target = 3*x1 + 2*x2^2 + noise, at N=4,000 and 200 simulations we see the expected ordering: CUPAC ≈ MLRATE >> raw covariates >> no covariates.

• No covariates: no variance reduction at all — the baseline.
• Raw covariates: OLS recovers the linear x1 relationship, so it clearly beats no covariates — but it captures none of x2's (non-linear) contribution, so it falls well short of CUPAC/MLRATE.
• CUPAC / MLRATE: by plugging a HistGradientBoostingRegressor into the same covariate slot instead of a linear term, both capture the entire relationship — linear x1 part and non-linear x2^2 part alike — and land at essentially the same (much higher) power.
• CUPAC needed a separate pre-experiment sample to fit its model; MLRATE reached the same result using only the experiment data, via cross-fitting, with no pre_experiment_df required.