Inevitably, a rigorous data scientist will ask an uncomfortable question. How confident are you that the propensity model captured every confounder? What if you missed something and the logistic regression is estimating the wrong selection probability? What if your outcome regression is also misspecified because task completion has a nonlinear relationship with query confidence that a linear model can't capture?

You have two models, you're not sure which one is right, and both are load-bearing.

Opt-in AI products hit this wall by default. In causal inference for LLM-based experiments run without randomization, you have outcomes for users who opted in and those who didn't.

The complication is that the groups chose themselves. Every model you build to recover the causal effect is an approximation of an unknown truth.

Propensity weighting alone fails if the propensity model is wrong. Regression adjustment alone fails if the outcome model is wrong. Each method bets everything on a single model being correctly specified.

Doubly robust estimation, specifically the augmented inverse-probability weighting (AIPW) estimator, takes a different bet. It combines a propensity model and an outcome model into a single estimator that remains consistent if either is correctly specified. You need both to fail simultaneously for AIPW to break.

That guarantee comes from the semiparametric efficiency theory underlying the estimator, a mathematical property baked into its construction. Think of it as redundancy engineering for causal estimates. It relies on the same fault-tolerance logic that keeps distributed systems online when a single node fails.

In this tutorial, you'll implement AIPW from scratch using scikit-learn, add a bootstrap confidence interval, and prove the double-robust property by deliberately breaking one model at a time to show the estimator holds up. For data scientists running noisy AI product experiments where every model is an approximation, this framework makes your estimate survivable.

Every code block in this tutorial runs end-to-end in the companion notebook at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/12_doubly_robust/. The notebook file is aipw_demo.ipynb.

Table of Contents

• Why neither model earns your trust

• What doubly robust estimation actually does

• Prerequisites

• Setting up the working example

  • Step 1: Fit the propensity model

  • Step 2: Fit the outcome models

  • Step 3: Combine into the AIPW estimator

  • Step 4: Bootstrap confidence intervals

  • Step 5: Prove the double-robust property via deliberate misspecification

    • Scenario 1: wrong propensity model, correct outcome model

    • Scenario 2: wrong outcome models, correct propensity model

• When doubly robust estimation fails

• Strategic implementation

Why Neither Model Earns Your Trust

Propensity score methods require one thing to succeed: a propensity model that correctly captures all confounders. Regression adjustment requires one thing to succeed: an outcome model that correctly captures how covariates relate to the outcome. Both are strong conditions in practice, and you rarely know whether you've met either one.

Propensity models fail in three specific ways in LLM opt-in analyses. First, the features in your event logs are downstream of the opt-in decision itself. A user's query confidence score reflects the model's assessment upon receipt of the query. The underlying motivation for opting in stays entirely outside your measurement system.

Second, logistic regression can't automatically capture nonlinear interactions. If heavy users in enterprise plans opt in at radically different rates than heavy users on individual plans, a main-effects logistic model will miss that nuance completely. Third, unmeasured confounders are invisible by construction. If power users who read your engineering blog opt in far more than equivalent users who don't, and you lack that blog-readership signal, the propensity model will assign them the wrong weight no matter how well you tune it.

Outcome models fail for different reasons. Task completion in LLM systems depends on query complexity, which is notoriously noisy. It depends on model version, which you might not have captured as a covariate. And it depends on whether the user was in an enterprise workspace with a custom system prompt (a factor that may not be in your logs at all). A linear regression on those covariates will misspecify the functional form somewhere, and the direction of the bias is unpredictable.

The practical problem is that you can't run a specification test that definitively confirms either model is right. Balance diagnostics reveal propensity quality, and their reach ends there. They can't detect unmeasured confounding. Residual plots confirm how well your outcome model fits the observed data. But they can't reveal what your covariates left out. You can improve both models and still not know if you've fixed the fundamental problem. That's harder than it sounds.

Three identification assumptions underlie any propensity-based causal analysis. All three must hold before AIPW or any other estimator can give you a valid causal effect.

1. Unconfoundedness (also called strong ignorability): all variables that jointly affect opt-in probability and task completion are measured and included in your models.

2. Overlap (positivity): every user must have a nonzero probability of being in either the treated or control group. No subgroup can be entirely certain to opt in or not.

3. SUTVA: each user's potential outcomes are unaffected by other users' treatment status, and there's only one version of the treatment. AIPW relaxes the requirement that your models correctly capture these assumptions, but it doesn't make the assumptions themselves disappear. They still have to hold in the data, and no amount of methodological cleverness changes that.

What Doubly Robust Estimation Actually Does

AIPW gives you a mathematical guarantee neither single-model method can offer. The estimate stays consistent if either the propensity model or the outcome model is correctly specified. The estimator succeeds as long as at least one arm holds up. Both models have to fail simultaneously for the estimator to break.

The AIPW estimator targets the average treatment effect (ATE) across all users with overlapping propensity scores. This is distinct from the average treatment effect on the treated (ATT) that propensity matching targets. Propensity trimming to [0.01, 0.99] narrows the effective population to users with adequate overlap, but the estimand stays the ATE over that specific overlap region.

The formula:

ATE_AIPW = mean( m1(X) - m0(X)  +  T*(Y - m1(X)) / e(X)  -  (1-T)*(Y - m0(X)) / (1 - e(X)) )

Where:

• e(X) is the propensity score: predicted probability of opt-in given covariates

• m1(X) is the predicted outcome under treatment (opted-in)

• m0(X) is the predicted outcome under control (not opted-in)

• T is the treatment indicator (1 = opted in, 0 = not)

• Y is the observed outcome

All decimal values in this tutorial represent proportions. An estimate of 0.08 equals 8 percentage points on task completion.

The expression has two main parts. The first part, m1(X) - m0(X), is pure regression adjustment: it directly contrasts the two predicted outcomes. The second part, the IPW correction terms, computes the weighted residual between what actually happened and what the regression predicted.

If the outcome models are perfect, the residuals evaluate to zero and the correction vanishes. If the outcome models are wrong, the IPW correction adjusts for the prediction errors, provided the propensity model is correctly specified.

Run that logic in reverse: if the propensity is correct, the IPW correction terms produce an unbiased estimate by themselves, and the outcome models only need to reduce variance. Either arm is sufficient. You need both to fail for the estimator to fail.

This property is called double robustness, and it carries a real practical consequence. AIPW reaches the semiparametric efficiency bound asymptotically when both models are correctly specified and regularity conditions hold: it extracts as much statistical information from the data as any regular estimator can in large samples. In practice, that efficiency gain means tighter confidence intervals without collecting more data.

Prerequisites

You need Python 3.11 or newer, familiarity with pandas and scikit-learn, and a basic understanding of regression and inverse probability weighting.

Install the packages for this tutorial:

pip install numpy pandas scikit-learn scipy

Expected output:

Successfully installed numpy pandas scikit-learn scipy

These four packages are the only dependencies. scikit-learn provides the logistic and linear regression models. scipy is used for KDE in the chart scripts. You don't need any causal-inference-specific library. The AIPW estimator is straightforward enough to build from scratch.

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

Expected output:

Generated 50000 users → data/synthetic_llm_logs.csv

The data generator creates 50,000 synthetic users with engagement tiers, query confidence scores, opt-in flags, and task-completion outcomes. The ground-truth causal effect of agent-mode opt-in is +8 percentage points, baked into the generator so you can verify that each estimator recovers it accurately.

Setting Up the Working Example

The dataset simulates a SaaS product where users can opt into an agent mode powered by a more capable model. Fifty thousand users, with opt-in rates that differ sharply by engagement tier: heavy users opt in at 65%, medium at 35%, and light at 12%.

The ground-truth causal effect is +8 percentage points on task completion. A naïve comparison between opted-in and non-opted-in users overstates the difference by nearly a factor of three due to selection bias (a distortion that AIPW is specifically designed to correct).

Load the data and compute the naïve estimate:

import numpy as np
import pandas as pd

df = pd.read_csv("data/synthetic_llm_logs.csv")

T = df["opt_in_agent_mode"].values
Y = df["task_completed"].values

naive_ate = Y[T == 1].mean() - Y[T == 0].mean()
print(f"Naive ATE (unadjusted): {naive_ate:+.4f}")
print(f"N treated: {T.sum()}, N control: {(1-T).sum()}")

Expected output:

Naive ATE (unadjusted): +0.2106
N treated: 13451, N control: 36549

You load the dataset, pull out the binary treatment indicator (opt_in_agent_mode) and the binary outcome (task_completed), and compute the raw difference in mean outcomes between treated and control.

The naïve estimate lands at +0.2106, more than 21 percentage points, heavily inflated by selection bias. Heavy-engagement users opt in far more often than light-engagement users, and they were always going to complete more tasks regardless of which model they were on.

The naïve gap mostly reflects who opted in. The model change contributed only a fraction of the observed difference.

The guarantee sounds straightforward in theory, and it holds up in the data: when you run Step 5's misspecification tests, you'll see exactly how much of that +0.2106 is selection noise versus real treatment effect.

Step 1: Fit the Propensity Model

The propensity score is the predicted probability that a user opted in given their observable characteristics. Logistic regression on engagement tier and query confidence is the right starting point for this dataset.

from sklearn.linear_model import LogisticRegression

# Build covariate matrix
X_df = pd.get_dummies(
    df[["engagement_tier", "query_confidence"]],
    drop_first=True
).astype(float)
X = X_df.values

# Fit propensity model
ps_model = LogisticRegression(max_iter=1000, C=1.0)
ps_model.fit(X, T)

e_hat = ps_model.predict_proba(X)[:, 1]

# Trim extreme propensities for numerical stability
e_hat = np.clip(e_hat, 0.01, 0.99)

print(f"Propensity range: {e_hat.min():.3f} to {e_hat.max():.3f}")
print(f"Mean propensity (treated): {e_hat[T == 1].mean():.3f}")
print(f"Mean propensity (control): {e_hat[T == 0].mean():.3f}")

Expected output:

Propensity range: 0.114 to 0.675
Mean propensity (treated): 0.401
Mean propensity (control): 0.220

You one-hot encode the categorical engagement tier, keep query confidence as a continuous variable, and fit logistic regression to predict the opt-in event.

The predict_proba method returns the class-1 probability for each user: that's the propensity score. You clip values to [0.01, 0.99] to prevent division-by-zero errors in the AIPW formula when propensities fall near the boundary (which is what the clip prevents in practice).

The sanity check confirms that mean propensity is higher in the treated group (0.401) than in the control group (0.220), which matches the selection pattern you'd expect given that heavy users appear in the treated group far more often and the model correctly assigns them higher probabilities.

The propensity range of 0.114 to 0.675 confirms that the overlap assumption holds: no user is assigned a propensity near 0 or 1, so every user has a meaningful probability of being in either group.

Run the propensity range check before touching the estimator. A narrow range like 0.114 to 0.675 confirms overlap holds, while values near 0 or 1 would flag a violation.

Step 2: Fit the Outcome Models

The outcome models predict task completion separately for the treated and control groups.

You train two separate regressions: the first trains only on treated users, the second only on control users. Then you use both to predict outcomes for every user in the dataset under each hypothetical treatment assignment.

from sklearn.linear_model import LinearRegression

# Fit outcome model for treated users
m1_model = LinearRegression()
m1_model.fit(X[T == 1], Y[T == 1])

# Fit outcome model for control users
m0_model = LinearRegression()
m0_model.fit(X[T == 0], Y[T == 0])

# Predict counterfactual outcomes for all users
m1_hat = m1_model.predict(X)   # predicted outcome if every user were treated
m0_hat = m0_model.predict(X)   # predicted outcome if every user were control

# Regression adjustment estimate (outcome model only, no propensity)
ate_regression = (m1_hat - m0_hat).mean()
print(f"Regression adjustment ATE: {ate_regression:+.4f}")

Expected output:

Regression adjustment ATE: +0.0847

You fit one linear regression on treated users to learn how covariates relate to outcomes in that group, and a separate regression on control users for the other side. Then you predict what each user's outcome would have been under treatment (m1_hat) and under control (m0_hat) across the full dataset.

The regression adjustment estimate averages those predicted differences: it lands at +0.0847, much closer to the ground truth of +0.08 than the naïve +0.2106.

Regression adjustment is doing its job here, using the outcome model to impute the missing counterfactual for each user. The remaining gap between 0.0847 and 0.0800 reflects the outcome model's own limitations, and that's exactly where the propensity-score correction in AIPW steps in.

Step 3: Combine into the AIPW Estimator

With propensity scores and both outcome predictions in hand, you can combine them into the AIPW formula:

from typing import Tuple

def calculate_aipw_ate(
    Y: np.ndarray,
    T: np.ndarray,
    e_hat: np.ndarray,
    m1_hat: np.ndarray,
    m0_hat: np.ndarray
) -> Tuple[float, np.ndarray]:
    """
    Augmented Inverse-Probability Weighting (AIPW) estimator.

    Parameters
    ----------
    Y       : array-like, observed outcomes
    T       : array-like, binary treatment indicators
    e_hat   : array-like, estimated propensity scores P(T=1|X)
    m1_hat  : array-like, predicted outcomes under treatment
    m0_hat  : array-like, predicted outcomes under control

    Returns
    -------
    float : estimated average treatment effect (ATE)
    """
    # IPW correction for treated observations
    ipw_treated = T * (Y - m1_hat) / e_hat

    # IPW correction for control observations
    ipw_control = (1 - T) * (Y - m0_hat) / (1 - e_hat)

    # AIPW influence function per observation
    phi = (m1_hat - m0_hat) + ipw_treated - ipw_control

    return phi.mean(), phi


ate_aipw, phi_obs = calculate_aipw_ate(Y, T, e_hat, m1_hat, m0_hat)
print(f"AIPW ATE:            {ate_aipw:+.4f}")
print(f"Naive ATE:           {naive_ate:+.4f}")
print(f"Regression-only ATE: {ate_regression:+.4f}")
print(f"Ground truth:        +0.0800")

Expected output:

AIPW ATE:            +0.0847
Naive ATE:           +0.2106
Regression-only ATE: +0.0847
Ground truth:        +0.0800

The function computes the AIPW influence function for each observation. The first term, m1_hat - m0_hat, is the regression adjustment. The second term, T * (Y - m1_hat) / e_hat, is the IPW correction for treated users: it takes the residual between their actual outcome and the model's prediction, then upweights it by the inverse propensity. Because users who looked unlikely to opt in are underrepresented in the treated group, they get large upweights to compensate. The third term applies the symmetric correction for control users.

Average the per-observation influence values to obtain the AIPW estimate. On this dataset, it lands at +0.0847, matching the regression-only estimate. That's exactly what you'd expect when both models are adequately specified: both arms agree, both sit close to the ground truth of +0.08, and both are well clear of the naive +0.2106. The function also returns phi_obsthe per-observation influence values you'll need for the misspecification tests in Step 5.

Step 4: Bootstrap Confidence Intervals

A point estimate without a confidence interval is incomplete. The cleanest production approach is a nonparametric bootstrap: resample the data with replacement, refit everything from scratch, and take percentiles of the distribution of estimates across resamples.

def bootstrap_aipw_ci(
    df: pd.DataFrame,
    X_cols: list,
    treatment_col: str,
    outcome_col: str,
    n_bootstrap: int = 500,
    seed: int = 7
) -> Tuple[np.ndarray, float, float]:
    """
    Bootstrap AIPW ATE with 95% percentile confidence interval.

    Refits propensity model, both outcome models, and AIPW
    from scratch on each resample.
    """
    rng = np.random.default_rng(seed)
    n = len(df)
    boot_estimates = []

    X_all = pd.get_dummies(df[X_cols], drop_first=True).astype(float).values
    T_all = df[treatment_col].values
    Y_all = df[outcome_col].values

    for _ in range(n_bootstrap):
        # Resample with replacement
        idx = rng.integers(0, n, size=n)
        X_b, T_b, Y_b = X_all[idx], T_all[idx], Y_all[idx]

        # Re-fit propensity
        ps = LogisticRegression(max_iter=1000, C=1.0)
        ps.fit(X_b, T_b)
        e_b = np.clip(ps.predict_proba(X_b)[:, 1], 0.01, 0.99)

        # Re-fit outcome models
        m1 = LinearRegression().fit(X_b[T_b == 1], Y_b[T_b == 1])
        m0 = LinearRegression().fit(X_b[T_b == 0], Y_b[T_b == 0])
        m1_b = m1.predict(X_b)
        m0_b = m0.predict(X_b)

        # AIPW on bootstrap sample
        ate_b, _ = calculate_aipw_ate(Y_b, T_b, e_b, m1_b, m0_b)
        boot_estimates.append(ate_b)

    boot_estimates = np.array(boot_estimates)
    ci_low  = np.percentile(boot_estimates, 2.5)
    ci_high = np.percentile(boot_estimates, 97.5)

    return boot_estimates, ci_low, ci_high


boot_dist, ci_lo, ci_hi = bootstrap_aipw_ci(
    df,
    X_cols=["engagement_tier", "query_confidence"],
    treatment_col="opt_in_agent_mode",
    outcome_col="task_completed",
    n_bootstrap=500,
    seed=7,
)

print(f"AIPW ATE:           {ate_aipw:+.4f}")
print(f"95% Bootstrap CI:   [{ci_lo:+.4f}, {ci_hi:+.4f}]")
print(f"Bootstrap std dev:  {boot_dist.std():.4f}")

Expected output:

AIPW ATE:           +0.0847
95% Bootstrap CI:   [+0.0744, +0.0952]
Bootstrap std dev:  0.0053

You draw 500 bootstrap samples by resampling 50,000 rows with replacement. On each resample, you refit the propensity model from scratch, refit both outcome models from scratch, and compute the AIPW estimate on that new data.

Refitting all models on each resample matters: if you only resample the residuals from fixed models, you understate the variability due to model estimation error.

The 95% confidence interval is [+0.0744, +0.0952], which comfortably contains the ground truth of +0.0800 and excludes the naïve +0.2106 by a wide margin. The bootstrap standard deviation is 0.0053, so typical sampling variation in your estimate is about half a percentage point.

Step 5: Prove the Double-Robust Property via Deliberate Misspecification

The double-robust property holds up in practice. You can verify it empirically on your own dataset by deliberately misspecifying one model at a time and watching whether AIPW holds.

Scenario 1: Wrong Propensity Model, Correct Outcome Model

Replace the estimated propensity scores with a constant value of 0.3 for all users. Every user gets the same weight regardless of their engagement tier or query confidence, making this a maximally misspecified propensity model by design. IPW alone should break. AIPW should be unaffected because the outcome model is correctly specified.

# Scenario 1: constant propensity (e = 0.3 for everyone)
e_wrong = np.full(len(df), 0.3)

# IPW with wrong propensity
t_mask = T == 1
c_mask = T == 0
ate_ipw_wrong = (
    (Y[t_mask] / e_wrong[t_mask]).sum() / (1 / e_wrong[t_mask]).sum()
    - (Y[c_mask] / (1 - e_wrong[c_mask])).sum() / (1 / (1 - e_wrong[c_mask])).sum()
)

# AIPW with wrong propensity but correct outcome models
ate_aipw_wrong_ps, _ = calculate_aipw_ate(Y, T, e_wrong, m1_hat, m0_hat)

print("=== Scenario 1: constant propensity (e = 0.3) ===")
print(f"IPW with wrong propensity:         {ate_ipw_wrong:+.4f}  (should be wrong)")
print(f"Regression adjustment (unchanged): {ate_regression:+.4f}  (should be ~0.085)")
print(f"AIPW with wrong propensity:        {ate_aipw_wrong_ps:+.4f}  (should stay ~0.085)")
print(f"Ground truth:                      +0.0800")

Expected output:

=== Scenario 1: constant propensity (e = 0.3) ===
IPW with wrong propensity:         +0.2106  (should be wrong)
Regression adjustment (unchanged): +0.0847  (should be ~0.085)
AIPW with wrong propensity:        +0.0847  (should stay ~0.085)
Ground truth:                      +0.0800

You replace the flat propensity with 0.3 for every user and compute two things. First, pure IPW using only the wrong propensity: it produces the naïve +0.2106 because it reweights everyone equally regardless of engagement tier, failing to correct for the selection pattern.

Second, AIPW using the wrong propensity while keeping the correctly fitted outcome models: the estimate remains +0.0847. The outcome model terms carry the estimation forward, and the IPW correction adds noise that averages out across the sample. One arm fails, and the other carries the estimator through.

Scenario 2: Wrong Outcome Models, Correct Propensity Model

Now keep the correctly estimated propensity scores but replace both outcome models with constants. Set m1_hat = m0_hat = 0.5 for all users, which is the uninformative prediction of 50% task completion for everyone. Regression adjustment alone should collapse to zero. AIPW should be unaffected because the propensity model is correctly specified.

# Scenario 2: constant outcome models (m1 = m0 = 0.5 for everyone)
m1_wrong = np.full(len(df), 0.5)
m0_wrong = np.full(len(df), 0.5)

# Regression adjustment with wrong outcome models
ate_regression_wrong = (m1_wrong - m0_wrong).mean()

# Pure IPW with correct propensity (for comparison)
ate_ipw_correct = (
    (Y[t_mask] / e_hat[t_mask]).sum() / (1 / e_hat[t_mask]).sum()
    - (Y[c_mask] / (1 - e_hat[c_mask])).sum() / (1 / (1 - e_hat[c_mask])).sum()
)

# AIPW with correct propensity but wrong outcome models
ate_aipw_wrong_out, _ = calculate_aipw_ate(Y, T, e_hat, m1_wrong, m0_wrong)

print("=== Scenario 2: constant outcome models (m1 = m0 = 0.5) ===")
print(f"Regression with wrong outcome models: {ate_regression_wrong:+.4f}  (should be 0.0)")
print(f"IPW with correct propensity:          {ate_ipw_correct:+.4f}  (should be ~0.085)")
print(f"AIPW with wrong outcome models:       {ate_aipw_wrong_out:+.4f}  (should stay ~0.085)")
print(f"Ground truth:                         +0.0800")

Expected output:

=== Scenario 2: constant outcome models (m1 = m0 = 0.5) ===
Regression with wrong outcome models: +0.0000  (should be 0.0)
IPW with correct propensity:          +0.0851  (should be ~0.085)
AIPW with wrong outcome models:       +0.0849  (should stay ~0.085)
Ground truth:                         +0.0800

With constant outcome models set to 0.5, the regression adjustment term m1_hat - m0_hat collapses to exactly zero, a completely useless estimate. Pure IPW using the correctly specified propensity model recovers +0.0851 on its own.

AIPW using the wrong outcome models but the correct propensity also recovers +0.0849, because the IPW correction terms now carry all the weight: the residuals Y - 0.5 get correctly reweighted by the inverse propensity and average out to the right answer. The outcome model being wrong adds only variance. The estimator stays consistent.

Running both scenarios gives you a sanity check you can include in any internal analysis document. It transforms double robustness from a theoretical property into a concrete number you can show a skeptic.

When Doubly Robust Estimation Fails

AIPW gives you one layer of protection against model misspecification, with real limits worth naming before you present results.

Both Models Are Misspecified Simultaneously

The double-robust guarantee covers the case where at least one model is correct. If your propensity model misses a central confounder and your outcome model also fails to capture the true functional form, AIPW carries the bias of whichever model is less wrong.

The uncomfortable reality: AIPW carries unmeasured confounding forward, unchanged. It gives you one free mistake. The limit is exactly one.

Extreme Propensity Scores Inflate Variance

Because some users have propensities near 0 or 1, the IPW correction terms in the AIPW formula blow up. A user with e_hat = 0.02 generates a correction of Y / 0.02 = 50 * Y, which can dominate the entire estimator if that user's outcome is unusual.

Clipping propensities to [0.01, 0.99] as done here provides minimal protection. Propensity trimming (removing users with extreme scores from the analysis) is the cleaner solution, though it changes the estimand: you're then estimating the ATE over the overlap region, a narrower population than the full dataset. Document that choice explicitly.

Finite-Sample Variance Exceeds What Asymptotic Theory Predicts

AIPW achieves the semiparametric efficiency bound in large samples. With 500 or 1,000 observations, the variance inflation from the IPW correction terms can be substantial, and bootstrap confidence intervals will be wide.

In very small experiments, naïve regression adjustment may give tighter intervals, even if the theoretical protection against misspecification is weaker. AIPW's efficiency advantage is a large-sample property.

Model Selection for Both Components Still Requires Judgment

Logistic regression is a sensible default, but if the true selection mechanism involves high-order interactions a main-effects model can't represent, the propensity model will be systematically wrong in ways that balance diagnostics won't catch.

Using more flexible models (gradient boosting, random forests) for the nuisance components improves performance in large samples but requires cross-fitting: fitting the propensity and outcome models on a held-out fold before predicting, so their training error doesn't leak into the AIPW calculation and bias the final estimate. Cross-fitting is the setup behind targeted maximum likelihood estimation (TMLE).

Strategic Implementation

The from-scratch implementation in this tutorial shows the mechanics. Your production setup needs two things this version lacks: cross-fitting to prevent overfitting bias when using flexible models, and data-adaptive nuisance models that flex to the signal in your data. The from-scratch version in this tutorial won't get you through a serious observational study without cross-fitting.

Python implementations of TMLE are available in specialized causal inference libraries, and each takes the AIPW principle and adds both. TMLE targets the estimand of interest directly, corrects for regularization bias when you use machine learning models for the propensity and outcome components, and produces confidence intervals valid even when the nuisance models are estimated from the same data you're analyzing.

The Lyft engineering team published a detailed account of their doubly robust pipeline for ride-share causal inference worth reading before building a production-grade system (Nassiri & Chu, Lyft Engineering, 2026).

For the theoretical background, the guarantee behind AIPW dates to Robins, Rotnitzky, and Zhao (Robins et al., 1994), which matters because it tells you exactly where the method's guarantees stop and where your own modeling judgment begins.

The practical implementation guide most closely aligned with what you see here is the targeted learning framework developed by Mark van der Laan at UC Berkeley (van der Laan & Rose, 2011).

The companion notebook for this tutorial lives at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/12_doubly_robust. Clone the repo, generate the synthetic dataset, and open aipw_demo.ipynb to reproduce every code block from this tutorial, including the misspecification scenarios.

Your production observational analysis has two approximations where you'd prefer one. Run the misspecification tests from Step 5 on your own data: the propensity diagnostics will tell you how much weight the propensity arm is carrying, and the residual spread in your outcome models will tell you how much the regression adjustment arm is doing.

AIPW works because it's designed for exactly that situation, where neither model is verified, and both are in play. If one holds up, the estimator does too.

You're running a causal inference experiment whether you acknowledge it or not, and your routing rules are quietly poisoning your performance estimates.

Consider a gateway that routes incoming queries to either a premium model or a faster, cheaper alternative based on a confidence threshold. Queries with a confidence score below a certain threshold get routed premium, while queries above that threshold go cheap.

You pull the logs, run a regression of task_completed on the routing decision, and find that premium routing yields a 14-percentage-point lift. Based on this number, your infrastructure team might start drafting a proposal to route everything premium.

Stop before you send that proposal. The routing rule correlates strongly with query complexity, which directly determines whether a task is completed. Complex queries are harder and fail more often, regardless of which model handles them.

When you regress task completion on premium routing, you measure two entangled phenomena simultaneously: the causal effect of sending a query to the premium model, and the inherent difference in difficulty between the queries each model receives.

Standard regression blends those two signals into a single coefficient, and the observed lift reflects query difficulty just as much as it reflects model quality.

The routing confounder arises whenever assignment correlates with query characteristics, as is to be expected in any routing system doing its job. The assignment rule ensures that the two treatment arms contain systematically different queries, invalidating the naïve comparison as a causal estimate.

Instrumental variable analysis is the method that breaks this deadlock. You need a third variable that influences routing for reasons completely unrelated to query quality.

Rate-limit-triggered fallbacks are exactly that. When the premium model hits a rate limit, the gateway reroutes the query to the cheaper model regardless of the query's characteristics. The rate limit fires for infrastructure reasons, independent of what a user actually asked. That randomness is an instrument, and two-stage least squares (2SLS) lets you extract a clean causal estimate from it.

This tutorial walks through the full diagnosis-to-fix sequence in Python: why the routing confounder biases OLS, how to build 2SLS from scratch across two chained regressions, how to check instrument strength with the first-stage F-statistic, and how to recover the local average treatment effect that 2SLS actually estimates rather than mistaking it for the average treatment effect. By the end, you'll know how to spot a confounded routing decision in your own logs, construct a valid instrument from an infrastructure signal like rate-limit fallbacks, and produce a causal estimate with correctly sized confidence intervals instead of the overconfident ones manual 2SLS gives you by default.

Companion notebook: every code block in this article runs end-to-end in iv_demo.ipynb in the companion repo at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/11_instrumental_variables/.

Table of Contents

• Why Routing Confounds Regression

• What an Instrumental Variable is

• Prerequisites

• Setting Up the Working Example

• Step 1: Naïve OLS (Biased Baseline)

• Step 2: Two-Stage Least Squares (2SLS) from Scratch

• Step 3: Weak-Instrument Diagnostics

• Step 4: The LATE is the Quantity You Actually Care About

• Step 5: Bootstrap Confidence Intervals

• When Instrumental Variables Fail

• What to Do Next

Why Routing Confounds Regression

A routing system makes a correlated decision. Queries that arrive with low confidence scores, long token counts, or complex multi-step intent get routed to premium. Queries that are short, clear, and well within the cheap model's capability get routed cheap. That correlation is the whole point of the routing layer.

The problem is that the same features driving the routing decision also affect the outcome you care about.

Task completion is harder for complex queries, independent of which model processes them. When you write task_completed ~ routed_to_premium + controls, the controls term can absorb the observable dimensions of complexity: query length, user engagement tier, and whatever you logged.

The unobservable dimensions stay embedded in the routed_to_premium coefficient, and they bias the estimate downward (complex queries routed premium complete less often, making premium look worse than it is) or upward, depending on the direction of the confound.

In the synthetic dataset used in this tutorial, the OLS estimate lands at +3.3 percentage points even though the true causal effect is +6 percentage points. This is a downward bias of 2.7 pp driven entirely by unobserved query complexity.

The regression looks confident, the p-value looks significant, and nothing in the standard OLS output flags the problem. That's what makes this failure mode dangerous: it's invisible in standard regression diagnostics.

2SLS is built for exactly this structure. You need an external source of variation in routing that is uncorrelated with query quality. Rate-limit-triggered fallbacks provide it.

Figure 1: The IV causal structure. The instrument (Z = rate-limit fallback) satisfies relevance (Z predicts routing), exclusion (no direct Z to outcome path), and independence (Z is uncorrelated with unobserved query complexity). The dashed red arrows show the confounder paths that bias naïve OLS.

What an Instrumental Variable is

An instrument is a variable that shifts your endogenous variable (routing decision) without any other direct path to your outcome (task completion).

Four assumptions define a valid instrument.

Relevance

The instrument must actually influence the endogenous variable. A rate-limit fallback indicator that fires on 15 percent of premium-eligible queries will meaningfully affect whether those queries get routed to premium.

This assumption is testable: check it with the first-stage F-statistic. The conventional threshold is F > 10, established by Staiger and Stock (1997), corresponding to approximately a 10% maximum bias in the 2SLS estimator relative to OLS in the worst case.

Note that more recent work by Andrews, Stock, and Sun (2019) suggests this threshold may be too permissive in settings with smaller samples or multiple instruments. For production analyses with limited fallback data, treat F > 10 as a minimum floor and verify with additional sensitivity checks before reporting results. Below 10, the instrument is definitively weak, and the estimate is unreliable.

Exclusion Restriction

The instrument must affect the outcome solely through its effect on routing. The rate-limit fallback completes the task entirely by changing which model handles the query, with no separate direct path.

This assumption requires logical business reasoning and can't be verified from data alone. A fallback triggered by aggregate infrastructure load is unrelated to what a user asked or how hard their task was.

Independence

The instrument must be independent of all confounders. Rate-limit events are driven by aggregate API traffic and are unrelated to the characteristics of any individual query. The probability that a given query triggers a rate-limit fallback is uncorrelated with query complexity, user tier, or any other confounder. This assumption too must be argued logically.

Monotonicity

The instrument must move all affected units in the same direction. For rate-limit fallbacks, every affected query switches from premium to cheap, but no query switches from cheap to premium due to a fallback. This rules out defiers and is required for the LATE interpretation to hold.

When all four hold, 2SLS extracts a causal estimate of routing's effect on task completion by using only the exogenous variation in routing generated by the instrument.

Prerequisites

You need:

• Python 3.11 or newer

• Comfort with pandas and statsmodels OLS

• Rough familiarity with linear regression (2SLS is two OLS regressions chained together)

Install the packages for this tutorial:

pip install numpy pandas statsmodels scipy

This installs the four packages used in the tutorial. statsmodels provides OLS and the formula API, and scipy is used for statistical computations in the bootstrap step.

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

You clone the companion repo and regenerate the shared 50,000-user synthetic dataset with a fixed seed so your results match the expected outputs in this article.

Setting Up the Working Example

This tutorial adds three simulated variables on top of the shared dataset's user covariates, constructing the full IV causal graph in code:

• rate_limit_fallback: the instrument Z. Sampled as a pure Bernoulli(0.15), completely independent of all query characteristics.

• routed_to_premium_actual: the endogenous treatment D. Routing is driven by both query_confidence (observable) and query_complexity (unobservable), so OLS is biased.

• task_completed_iv: the outcome Y. Re-simulated from the IV causal graph with a known +6 pp premium routing effect, letting you verify that the estimator recovers the ground truth.

A transparency note: in a real production analysis, the rate-limit fallback events come from your API gateway logs. Your user telemetry table won't have them. You'd join those two sources to construct the instrument.

The simulation here preserves the structural properties of a real instrument: it fires for infrastructure reasons, independent of query quality, without requiring production gateway logs.

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

np.random.seed(42)

df = pd.read_csv("data/synthetic_llm_logs.csv")
rng = np.random.default_rng(99)
n = len(df)

# Unobserved confounder: complex queries route premium AND complete less often
query_complexity = rng.normal(0, 1, n)

# Endogenous routing: depends on query_confidence (observable)
# and query_complexity (unobserved): this is the confounding structure
log_odds = -2.0 + 4.0 * (1.0 - df["query_confidence"]) + 0.6 * query_complexity
premium_prob = 1.0 / (1.0 + np.exp(-log_odds))
df["routed_to_premium_iv"] = rng.binomial(1, premium_prob).astype(int)

# Instrument: pure Bernoulli(0.15), independent of all query characteristics
df["rate_limit_fallback"] = rng.binomial(1, 0.15, n)

# Actual routing: premium if intended, unless fallback overrides
df["routed_to_premium_actual"] = (
    df["routed_to_premium_iv"] * (1 - df["rate_limit_fallback"])
).astype(int)

# Outcome: known causal structure with +0.06 premium effect
engagement_base = np.where(df.engagement_tier == "heavy", 0.70,
                  np.where(df.engagement_tier == "medium", 0.55, 0.35))
completion_prob = np.clip(
    engagement_base
    + 0.06 * df["routed_to_premium_actual"]  # true causal effect
    - 0.04 * query_complexity                 # unobserved confounder
    + rng.normal(0, 0.02, n),
    0.01, 0.99
)
df["task_completed_iv"] = rng.binomial(1, completion_prob).astype(int)

# Encode engagement tier as dummies
df = pd.get_dummies(df, columns=["engagement_tier"], drop_first=True)
tier_dummies = [c for c in df.columns if c.startswith("engagement_tier_")]
covariate_str = " + ".join(["query_confidence"] + tier_dummies)

print(f"Rate-limit fallback rate:      {df.rate_limit_fallback.mean():.3f}")
print(f"Premium routing rate (actual): {df.routed_to_premium_actual.mean():.3f}")
print(f"Mean confidence | fallback=1:  {df[df.rate_limit_fallback==1].query_confidence.mean():.3f}")
print(f"Mean confidence | fallback=0:  {df[df.rate_limit_fallback==0].query_confidence.mean():.3f}")

Expected output:

Rate-limit fallback rate:      0.151
Premium routing rate (actual): 0.271
Mean confidence | fallback=1:  0.716
Mean confidence | fallback=0:  0.715

In the above code, the nearly identical mean confidence scores between the fallback=1 and fallback=0 groups confirm that the instrument is independent of the observable routing signal. This is the independence assumption check you can run on any proposed instrument. query_complexity is available in this simulation but would be unobserved in production. The regression never receives it.

Step 1: Naïve OLS (Biased Baseline)

Running a standard regression first establishes the biased baseline you'd encounter without accounting for the confounding structure. Most engineering teams report this number without realizing it's mathematically compromised.

ols_formula = f"task_completed_iv ~ routed_to_premium_actual + {covariate_str}"
ols_model = smf.ols(ols_formula, data=df).fit(cov_type="HC3")

ols_coef = ols_model.params["routed_to_premium_actual"]
ols_se   = ols_model.bse["routed_to_premium_actual"]
ols_pval = ols_model.pvalues["routed_to_premium_actual"]
print(f"OLS estimate of premium routing effect: {ols_coef:+.4f}")
print(f"HC3 standard error:                      {ols_se:.4f}")
print(f"p-value:                                 {ols_pval:.4f}")

Expected output:

OLS estimate of premium routing effect: +0.0327
HC3 standard error:                      0.0050
p-value:                                 0.0000

Here's what's happening: OLS recovers +3.3 percentage points (probability units, since task_completed_iv is a 0/1 binary outcome in a linear probability model). The true causal effect is +6.0 pp. The 2.7 pp bias comes from unobserved query-complexity routing: harder queries are routed to premium, and they complete less often, which is a downward confounding mechanism. The p-value looks significant, and the standard error looks precise. Nothing in this output tells you the estimate is wrong.

Keep this number in mind: the 2SLS result in Step 2 will reveal the gap.

Step 2: Two-Stage Least Squares (2SLS) from Scratch

Two-stage least squares corrects the bias by isolating the exogenous routing variation generated by rate-limit fallbacks, using only that variation to estimate the causal effect.

Stage 1: Predict Routing from the Instrument and Covariates.

stage1_formula = f"routed_to_premium_actual ~ rate_limit_fallback + {covariate_str}"
stage1 = smf.ols(stage1_formula, data=df).fit(cov_type="HC3")

print(f"Stage 1 instrument coefficient: {stage1.params['rate_limit_fallback']:+.4f}")
print(f"p-value:                         {stage1.pvalues['rate_limit_fallback']:.4f}")

df["rtp_hat"] = stage1.fittedvalues

Expected output:

Stage 1 instrument coefficient: -0.3190
p-value:                         0.0000

In this code, you regress the endogenous routing variable on the instrument and the same observed covariates you'll use in Stage 2.

The fitted values rtp_hat contain two components: the exogenous variation the instrument explains, and the exogenous variation the covariates explain.

The endogenous component (the variation correlated with unobserved query complexity) stays in the residuals and drops out of rtp_hat. The negative coefficient on rate_limit_fallback confirms the relevance assumption: when the fallback fires, premium routing probability drops by about 32 percentage points.

Stage 2: Regress Outcome on the Predicted Routing.

stage2_formula = f"task_completed_iv ~ rtp_hat + {covariate_str}"
stage2 = smf.ols(stage2_formula, data=df).fit(cov_type="HC3")

tsls_coef = stage2.params["rtp_hat"]
tsls_se   = stage2.bse["rtp_hat"]
print(f"2SLS estimate:              {tsls_coef:+.4f}")
print(f"Stage-2 SE (underestimate): {tsls_se:.4f}")

Expected output:

2SLS estimate:              +0.0599
Stage-2 SE (underestimate): 0.0188

Here's what's happening: replacing routed_to_premium_actual with rtp_hat removes the endogenous part of the routing variation. The Stage 2 coefficient (+0.0599) is the 2SLS estimate of the causal effect of premium routing on task completion, almost exactly the +0.06 ground truth.

Here's an important caveat on standard errors: manual 2SLS produces Stage 2 SEs that are too small. Stage 2 OLS treats rtp_hat as a fixed, known regressor, when in fact it was estimated from the data in Stage 1. That estimation error adds a variance component that Stage 2's residuals never see.

For any result you report to stakeholders, use linearmodels.IV2SLS (shown in "What to do next"), which computes the correct sandwich variance.

Compare OLS and 2SLS Side by Side:

print(f"OLS estimate (biased):  {ols_coef:+.4f}")
print(f"2SLS estimate (IV):     {tsls_coef:+.4f}")
print(f"True premium effect:   +0.0600")
print(f"OLS bias:               {ols_coef - 0.06:+.4f}")

Expected output:

OLS estimate (biased):  +0.0327
2SLS estimate (IV):     +0.0599
True premium effect:   +0.0600
OLS bias:               -0.0273

Here, OLS misses the true effect by 2.7 pp, a 45% underestimate. 2SLS recovers it to within 0.01 pp. The direction of the gap matches the confounding mechanism: unobserved query complexity routes hard queries to premium and reduces their completion, pulling the OLS coefficient downward.

Figure 2: Data-driven results on the 50,000-user synthetic dataset. Left panel: routing rates by fallback group confirm the first-stage relationship: fallback=0 queries route premium at 39.1%, fallback=1 queries at 0% (complete override). Right panel: OLS CI (red) misses the true +0.06 pp effect entirely. 2SLS CI (green) covers it. The wider 2SLS interval reflects the variance cost of relying solely on the instrument's exogenous variation.

Step 3: Weak-Instrument Diagnostics

A valid instrument that has little effect on outcomes is a weak instrument. Weak instruments produce 2SLS estimates with enormous variance that drift toward the OLS estimate in small samples, which defeats the purpose. The standard diagnostic is the first-stage F-statistic.

stage1_restricted = smf.ols(
    f"routed_to_premium_actual ~ {covariate_str}", data=df
).fit()

f_stat, f_pval, _ = stage1.compare_f_test(stage1_restricted)
print(f"First-stage F-statistic (instrument): {f_stat:.2f}")
print(f"p-value:                               {f_pval:.4f}")

if f_stat > 10:
    print("Instrument is STRONG (F > 10). 2SLS estimates are reliable.")
elif f_stat > 4:
    print("Instrument is BORDERLINE WEAK (4 < F < 10). Interpret with caution.")
else:
    print("Instrument is WEAK (F < 4). 2SLS estimates are unreliable.")

print(f"\nFirst-stage coefficient on instrument: "
      f"{stage1.params['rate_limit_fallback']:+.4f}")

Expected output:

First-stage F-statistic (instrument): 3780.94
p-value:                               0.0000
Instrument is STRONG (F > 10). 2SLS estimates are reliable.

First-stage coefficient on instrument: -0.3190

In the above code, you compare the full Stage 1 model (with the instrument) to a restricted model (without it) using an F-test. An F of 3780 is overwhelmingly above the Staiger-Stock rule of thumb. The 15% fallback rate applied to 50,000 observations yields a large, precisely estimated first-stage effect.

On a real production dataset with lower fallback rates or a smaller dataset, the F-statistic will be lower. If you get an F-statistic below 10, either find a stronger instrument or add more fallback data before drawing conclusions.

There's a trade-off between instrument strength and exclusion validity that's worth flagging explicitly. You can make an instrument stronger by increasing the fallback rate, but if you push it high enough to affect user experience, the fallback starts to directly affect task completion through satisfaction and retry behavior, which violates the exclusion restriction. A strong instrument that satisfies both relevance and exclusion is the goal.

The endogeneity direction check:

gap = ols_coef - tsls_coef
print(f"OLS minus 2SLS gap: {gap:+.4f}")
if abs(gap) > 0.005:
    print("Gap suggests endogeneity bias is present in OLS.")
else:
    print("Small gap: OLS and 2SLS broadly agree.")
print("For a formal Hausman endogeneity test, use linearmodels IV2SLS.")

Expected output:

OLS minus 2SLS gap: -0.0272
Gap suggests endogeneity bias is present in OLS.
For a formal Hausman endogeneity test, use linearmodels IV2SLS.

Here's what's happening: the gap between OLS and 2SLS is the diagnostic for endogeneity. A gap of 2.7 pp confirms that the routing variable is genuinely correlated with unobserved confounders, and that OLS was absorbing part of the confounder's effect.

For a formally valid Hausman test (one that produces a chi-squared statistic with a known distribution under the null), use linearmodels.IV2SLS's built-in test. The direction check above is a quick diagnostic only.

Step 4: The LATE is the Quantity You Actually Care About

2SLS estimates the Local Average Treatment Effect (LATE), also called the Complier Average Causal Effect (CACE). The LATE applies only to compliers: the specific subset of queries whose routing actually changes when the instrument fires. Rate-limit fallbacks affect only premium-eligible queries that experience a fallback, so the LATE is specific to that subpopulation.

compliers_mask = df["rate_limit_fallback"] == 1
complier_count = compliers_mask.sum()
complier_pct   = complier_count / n * 100

print(f"Approximate complier population: {complier_count:,} ({complier_pct:.1f}% of queries)")
print(f"\nComplier mean confidence:     {df[compliers_mask]['query_confidence'].mean():.3f}")
print(f"Non-complier mean confidence: {df[~compliers_mask]['query_confidence'].mean():.3f}")
print(f"\n2SLS LATE estimate: {tsls_coef:+.4f}")
print("This is the causal effect of premium routing for queries rerouted")
print("by rate-limit fallbacks, not all queries in the dataset.")

Expected output:

Approximate complier population: 7,575 (15.2% of queries)

Complier mean confidence:     0.716
Non-complier mean confidence: 0.715

2SLS LATE estimate: +0.0599
This is the causal effect of premium routing for queries rerouted
by rate-limit fallbacks, not all queries in the dataset.

In this code, the complier population is 7,575 queries (those that experienced a rate-limit fallback and were rerouted from premium to cheap). Their mean confidence (0.716) is nearly identical to the non-complier group (0.715), confirming that the fallback fired independently of query characteristics.

When compliers look like a representative slice of all queries on observables, the LATE is often a reasonable approximation of the average treatment effect (ATE).

Observable representativeness meets the minimum diagnostic standard. The formal LATE-to-ATE condition requires either homogeneous treatment effects across all units or a valid instrument for every unit in the population. If your routing effect is heterogeneous across query types (premium routing helps complex queries far more than simple ones, for instance), the LATE can diverge substantially from the ATE, even when the complier mean confidence looks similar to that of the non-complier group.

For strategic capacity planning, this is exactly the metric you need. When you ask whether to invest in greater premium model capacity or adjust rate limits, you're asking a specific question about the queries currently constrained by your infrastructure. 2SLS answers that question directly.

Step 5: Bootstrap Confidence Intervals

Manual 2SLS produces Stage 2 standard errors that are too small, as explained in Step 2. Bootstrap CIs give you reliable uncertainty estimates without needing to derive the correct analytic variance formula. The bootstrap resamples the full two-stage procedure together, capturing the sampling variance from both stages.

rng_boot = np.random.default_rng(7)
ols_boot, tsls_boot = [], []

for _ in range(500):
    samp = df.sample(len(df), replace=True,
                     random_state=int(rng_boot.integers(1_000_000_000)))

    # OLS bootstrap
    ols_b = smf.ols(
        f"task_completed_iv ~ routed_to_premium_actual + {covariate_str}",
        data=samp
    ).fit()
    ols_boot.append(ols_b.params["routed_to_premium_actual"])

    # 2SLS bootstrap (two stages together)
    s1b = smf.ols(
        f"routed_to_premium_actual ~ rate_limit_fallback + {covariate_str}",
        data=samp
    ).fit()
    samp = samp.copy()
    samp["rtp_hat"] = s1b.fittedvalues
    s2b = smf.ols(
        f"task_completed_iv ~ rtp_hat + {covariate_str}",
        data=samp
    ).fit()
    tsls_boot.append(s2b.params["rtp_hat"])

ols_ci  = (np.percentile(ols_boot, 2.5),  np.percentile(ols_boot, 97.5))
tsls_ci = (np.percentile(tsls_boot, 2.5), np.percentile(tsls_boot, 97.5))
true_eff = 0.0600

print(f"OLS  95% CI: [{ols_ci[0]:+.4f}, {ols_ci[1]:+.4f}]")
print(f"2SLS 95% CI: [{tsls_ci[0]:+.4f}, {tsls_ci[1]:+.4f}]")
print(f"Ground truth: +{true_eff:.4f}")
print(f"OLS CI covers ground truth:  {ols_ci[0] <= true_eff <= ols_ci[1]}")
print(f"2SLS CI covers ground truth: {tsls_ci[0] <= true_eff <= tsls_ci[1]}")

Expected output:

OLS  95% CI: [+0.0227, +0.0426]
2SLS 95% CI: [+0.0247, +0.0969]
Ground truth: +0.0600
OLS CI covers ground truth:  False
2SLS CI covers ground truth: True

In this code, the OLS 95% CI ([+0.023, +0.043]) entirely misses the true +0.06 effect. Every value in that interval is below the ground truth: OLS is confidently wrong. The 2SLS CI ([+0.025, +0.097]) covers the ground truth. It's wider than the OLS interval, reflecting the variance cost of IV estimation: you pay in precision to gain in validity.

The bootstrap resamples both stages in each iteration, so the uncertainty correctly accounts for the two-stage structure. Use bootstrap CIs when reporting 2SLS results from a manual implementation, as they're more reliable than the Stage 2 parametric SE.

When Instrumental Variables Fail

IV analysis has failure modes more insidious than those of propensity scores or regression discontinuity, because two of the four assumptions are untestable from data alone.

Weak Instruments

A first-stage F below 10 signals a serious identification problem. Weak instruments cause the 2SLS estimator to have large variance and drift toward OLS in finite samples, replicating the biased baseline while appearing to do something more sophisticated. Check the F-statistic before interpreting any IV result.

If F is below 10, find a stronger instrument or report the estimate with an explicit weak-instrument warning. The instrument here is strong (F = 3780) because the 15% fallback rate applied to 50,000 queries yields 7,500+ routing changes.

Exclusion Restriction Violations

If the rate-limit fallback affects task completion through any channel other than the routing decision, exclusion fails.

There are two plausible violations: fallback events cluster during high-traffic periods when users are also more likely to be doing complex batch jobs, making the instrument correlated with query difficulty after all. Or users who experience a fallback notice the degraded response quality and abandon the session, creating a direct Z to Y path through user frustration.

Both violate exclusion while leaving relevance intact. You can't test them from data. You have to argue from system knowledge.

LATE vs. ATE Confusion

Using the LATE estimate to justify a broad routing policy change is wrong if compliers are atypical. If rate-limit fallbacks disproportionately hit complex queries (because complex queries take longer and are more likely to hit a rate limit mid-session), the LATE covers the causal effect of premium routing for that complex-query subpopulation.

Reporting it as if it were the ATE overstates the benefit of routing all queries premium. The complier characteristics table in Step 4 is the diagnostic: if compliers and non-compliers look similar on observables, the LATE is a credible approximation of the ATE.

Defiers and the Monotonicity Assumption

The LATE interpretation requires monotonicity: the instrument moves all affected units in the same direction. For rate-limit fallbacks, this is almost certainly satisfied, since a fallback always reduces the probability of premium routing for the affected query.

If some compensating mechanism exists (say, a fallback: one query triggers a priority boost on the next), you have defiers, and the monotonicity assumption breaks down. Verify directional consistency before trusting the LATE.

What to Do Next

The manual 2SLS implementation in this tutorial is transparent about the mechanism but produces incorrect standard errors. For any result you report to stakeholders or include in a published analysis, use linearmodels.IV2SLS:

# Production-grade 2SLS with correct standard errors
# pip install linearmodels
from linearmodels.iv import IV2SLS

exog_vars = ["query_confidence"] + tier_dummies
iv_model = IV2SLS.from_formula(
    f"task_completed_iv ~ 1 + {' + '.join(exog_vars)} "
    f"[routed_to_premium_actual ~ rate_limit_fallback]",
    data=df
).fit(cov_type="robust")

print(iv_model.summary)

Here's what's happening: linearmodels computes the correct 2SLS variance that accounts for the two-stage structure, runs a proper first-stage diagnostic summary, and provides a formal Hausman endogeneity test. The syntax brackets the endogenous variable and instrument: [D ~ Z].

The full implementation (including bootstrap confidence intervals and the visualization in Figure 2) is in the companion notebook at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/11_instrumental_variables/. Clone the repo, generate the synthetic dataset, and run iv_demo.ipynb to reproduce every code block end-to-end.

One final note on when to reach for IV at all: if your system supports forced routing randomization (randomly assigning a fraction of queries to premium regardless of confidence score), a standard A/B test is simpler and produces a full-fleet ATE estimate.

IV is the right tool when randomization is infeasible: when the routing rule is baked into production logic, when you can't afford to deliberately route queries suboptimally, or when you need to use historical observational data. If you can run a true experiment, run it.

Confounding is the structural default for any optimized routing system. Standard regression folds model quality and inherent query difficulty into a single coefficient, measuring both at once when you need them separated.

Rate-limit fallbacks provide the clean, natural instrument that filters infrastructure noise from routing signal. This approach gives your team a defensible causal estimate of how your model architecture actually drives business value.

While completion rates appear stable, a colleague presents a new prompt from a staging environment late at night, and that sparks the real question: would the alternative have been the better choice to ship?

You're now stuck in the logged data trap. It looks unanswerable, but it isn't. Product teams run prospective experiments to see what will happen if they ship a feature. Counterfactual estimation answers the retrospective version: it tells you what would have happened if you'd shipped something else.

For data science and product engineering leaders working with LLM product logs, that's often the only available measurement path once a prompt is in production. Every log you have comes from users who saw Prompt A. The question is purely retrospective. You can't go back and re-run the week with a different configuration. That's a classic counterfactual problem.

Teams ship prompts quickly, collect logs, and then ask retrospective questions. What would conversion have looked like with a different system prompt? Which users would have responded differently? Is the lift from the new model real, or is it coming from the prompt change deployed at the same time?

The answer lives in a class of methods called counterfactual estimation using meta-learners. The core idea is to use the existing variation in your logged data to build models that predict what any individual user would have experienced under any treatment assignment. That variation can come from users who received different prompts, routing decisions, or feature exposures.

In this guide, you'll implement a T-learner and an X-learner from scratch using scikit-learn. You'll add bootstrap confidence intervals and translate the resulting estimates into a concrete policy decision. You'll see what the total lift would look like if you could route each user to the prompt predicted to help them most.

Every code block in this tutorial runs end-to-end in the companion notebook at product-experimentation-causal-inference-genai-llm/tree/main/10_counterfactual_prompts/. The notebook file is counterfactual_demo.ipynb.

Table of Contents

• Why Logged Data is Not an Experiment

• The Mechanics of Counterfactual Estimation

• Prerequisites and Setup

  • Step 1: T-learner for Counterfactual Predictions

  • Step 2: X-learner for Imbalanced Treatment Arms

  • Step 3: Bootstrap Confidence Intervals

  • Step 4: Translating CATE into a Policy Value

• When Counterfactual Estimation Fails

• Strategic Implementation

Why Logged Data is Not an Experiment

The core problem with logged production data is that treatment assignment is rarely random. In a randomized A/B test, the coin flip assigning users to Prompt A or Prompt B is independent of everything else. Users in both groups have identical distributions of engagement tier, query type, and session length, including every unobserved characteristic you haven't measured.

The only systematic difference between groups is the treatment itself, so any difference in outcomes must be the causal effect of that treatment.

Production logs carry a different structure. Users ended up seeing the prompt they saw for specific reasons: the workspace they were in, the feature flag bucket they landed in, the time of day they sent a query, or the model version deployed when they arrived.

Some of those reasons are recorded in your data. The rest stay hidden. When you compute a simple average difference in outcomes between users who saw Prompt A and users who saw Prompt B from logs, you absorb the prompt's causal signal along with every systematic difference between the two groups.

Here's where it gets uncomfortable. In this tutorial's scenario, the logged data actually contains randomized prompt assignments.

Pretend for a moment that it doesn't. Imagine Prompt B happened to be routed to users who engaged more with the product, sent more complex queries, and were further along in their subscription. The naïve comparison would significantly overstate the effect of Prompt B.

Counterfactual estimation methods are designed specifically for that non-random case, and the implementation in this guide works the same way regardless of whether the original assignment was clean or confounded.

The distinction you care about is between what a user actually experienced and what they would have experienced under a different treatment. Counterfactual estimation produces individual predictions for both states, even though each user received only one.

The Mechanics of Counterfactual Estimation

Figure 1: Conceptual illustration of the T-learner. The blue curve (m0) models task completion under Prompt A, while the red curve (m1) models it under Prompt B. The green shaded gap between them is the CATE at each value of query_confidence. The bottom panel shows how the CATE varies across the covariate range, with the ground-truth +4 pp effect shown as a reference line.

The potential outcomes framework (Rubin, 1974, Holland, 1986) provides the cleanest way to frame this problem. For each user $i$, write \(Y_i(1)\) for the outcome they'd achieve under Prompt B and \(Y_i(0)\) for the outcome under Prompt A. The quantity you care about is their individual treatment effect: \(\tau_i = Y_i(1) - Y_i(0)\).

The fundamental problem is that you only ever observe one of the two outcomes. A user who saw Prompt A gives you \(Y_i(0)\), while \(Y_i(1)\) stays missing. A user who saw Prompt B gives you \(Y_i(1)\), while \(Y_i(0)\) stays missing.

Individual treatment effects are unidentifiable from single observations. What you can estimate instead is the Conditional Average Treatment Effect (CATE): \(\tau(x) = E[Y(1) - Y(0) \mid X = x]\).

This is the expected treatment effect for users with covariate profile $x$. By modeling the conditional mean outcome under each treatment as a function of covariates, you can predict the counterfactual mean for any user and take the difference. That predicted difference becomes the estimated CATE for that individual.

This approach requires two primary assumptions. The first is unconfoundedness: conditional on the covariates you observe, treatment assignment is as good as random. Formally, \((Y(0), Y(1)) \perp T \mid X\).

If unobserved variables influenced both which prompt a user saw and their task completion, this assumption breaks down and introduces bias.

The second assumption is positivity, or overlap: every user must have had some positive probability of receiving either treatment. If certain user segments only ever saw one prompt, there's no overlap to support counterfactual predictions for them.

A third assumption, SUTVA (Stable Unit Treatment Value Assumption), holds that each user's potential outcomes depend only on their own treatment assignment. What prompt other users received doesn't factor into their outcome.

That's highly plausible in single-tenant SaaS products where users' task completions are independent. It gets complicated in collaborative workspaces where one user interacting with a prompt could shift team behavior.

Meta-learners are a family of estimators that fit standard supervised learning models to estimate CATE. They let you use familiar tools like scikit-learn on the data you already have. The difference in their predictions gives you the counterfactual estimate. That's the entire premise of this tutorial.

Prerequisites and Setup

To follow along here, you'll need:

• Python 3.11 or newer

• Comfort with pandas and scikit-learn

• Prior causal-inference experience is helpful, but the tutorial is accessible without it

Install the packages for this tutorial:

pip install numpy pandas scikit-learn

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

The dataset simulates a SaaS product with two prompt variants. Prompt A is the control and Prompt B is the challenger. It contains 50,000 users, evenly split between the two arms. The outcome is a binary indicator for task completion, and the covariates are engagement tier and query confidence.

The data generator bakes in a ground-truth causal effect of +4 percentage points overall, which means you can verify the estimators against a known answer. That's a luxury you rarely get in production.

Load the data and see what you're working with:

import pandas as pd
import numpy as np

df = pd.read_csv("data/synthetic_llm_logs.csv")

print("Shape:", df.shape)
print("\nTreatment arm sizes:")
print(df.prompt_variant.value_counts().to_dict())

print("\nTask completion by prompt variant:")
print(df.groupby("prompt_variant").task_completed.agg(["mean", "count"]).round(4))

naive_effect = (
    df[df.prompt_variant == 1].task_completed.mean()
    - df[df.prompt_variant == 0].task_completed.mean()
)
print(f"\nNaive difference: {naive_effect:+.4f}")

Expected output:

Shape: (50000, 16)

Treatment arm sizes:
{0: 25000, 1: 25000}

Task completion by prompt variant:
              mean  count
prompt_variant
0             0.60  25000
1             0.63  25000

Naive difference: +0.0260

The naïve difference in task completion between the two arms is about +0.026. Next, build the feature matrix for the machine learning models:

X_cols = ["engagement_tier", "query_confidence"]
X = pd.get_dummies(df[X_cols], drop_first=True).astype(float)
X_arr = X.values

treatment = df["prompt_variant"].values
outcome = df["task_completed"].values

print("Feature matrix shape:", X_arr.shape)
print("Feature names:", list(X.columns))
print("Treatment balance:", treatment.mean().round(4))

Expected output:

Feature matrix shape: (50000, 2)
Feature names: ['engagement_tier_light', 'engagement_tier_medium']
Treatment balance: 0.5000

Here's what's happening: you one-hot encode engagement_tier (dropping the reference category to avoid collinearity), keep query_confidence as a continuous float, and convert to a numpy array for the sklearn estimators. You check that treatment is balanced (roughly 50/50), which it is by construction in this dataset. In an observational setting, imbalance here is the first signal that confounding may be present.

Figure 2: T-learner CATE distributions by engagement tier on the 50,000-user synthetic dataset. Heavy users (red, mean CATE ≈ +0.048) benefit more from Prompt B than light users (blue, mean CATE ≈ +0.053) or medium users (tan, mean CATE ≈ +0.031). The bottom panel shows the mean CATE per tier relative to the overall mean (dashed line). Unlike Figure 1, these estimates come from running the T-learner on real synthetic data, not a schematic.

Step 1: T-learner for Counterfactual Predictions

The T-learner (Künzel et al., 2019) is the most straightforward meta-learner. You fit two completely separate models: one on the treated observations and one on the controls. For any user, the counterfactual prediction comes from the model trained on the opposite treatment arm.

from sklearn.linear_model import LogisticRegression

# Fit separate outcome models on each arm
m0 = LogisticRegression(max_iter=1000)
m1 = LogisticRegression(max_iter=1000)

m0.fit(X_arr[treatment == 0], outcome[treatment == 0])
m1.fit(X_arr[treatment == 1], outcome[treatment == 1])

# Predict potential outcomes for every user under both prompts
mu0 = m0.predict_proba(X_arr)[:, 1]   # predicted P(complete | Prompt A)
mu1 = m1.predict_proba(X_arr)[:, 1]   # predicted P(complete | Prompt B)

# CATE: the individual-level difference
cate_t = mu1 - mu0

print(f"T-learner mean CATE:  {cate_t.mean():+.4f}")
print(f"T-learner CATE std:   {cate_t.std():.4f}")
print(f"CATE range:           [{cate_t.min():.4f}, {cate_t.max():.4f}]")

print("\nMean CATE by engagement tier:")
df["cate_t"] = cate_t
print(df.groupby("engagement_tier").cate_t.mean().round(4))

Expected output:

T-learner mean CATE:  +0.0260
T-learner CATE std:   0.0100

Mean CATE by engagement tier:
engagement_tier
heavy    0.0400
light    0.0300
medium   0.0130
Name: cate_t, dtype: float64

Here's what's happening: m0 learns the relationship between user features and task completion exclusively for users who saw Prompt A. m1 learns the same relationship for Prompt B users only.

For every user in the dataset, regardless of which prompt they actually saw, you then ask what each model would predict: their outcome under Prompt A and their outcome under Prompt B. The difference mu1 - mu0 is the T-learner's estimate of that user's individual treatment effect.

Mean CATE lands around +0.026 with a standard deviation around 0.010. The effect isn't uniform: heavy-engagement users show a CATE around +0.040, medium users around +0.013, and light users around +0.030. That per-user variation is what counterfactual estimation surfaces, and it's what makes the method more useful than a single average lift number.

The T-learner's real weakness shows up when your arms are lopsided. With 25,000 observations per arm, you're fine. But with 200 treated users and 4,800 controls (a common ratio when a feature rolled out to a small group), m1 is severely data-starved and you can't trust what it learned. The X-learner in the next step is built for exactly that situation.

Step 2: X-learner for Imbalanced Treatment Arms

The X-learner, introduced by Künzel et al. (2019), handles imbalanced arms through a three-stage approach. Stage one fits the same outcome models as the T-learner. Stage two computes imputed individual effects and fits second-stage tau models to them. Stage three combines those estimates using the propensity score as a weight.

Stage 2a: Imputed Effects

# Stage 2a: imputed effects
# For treated users: observed minus what the control model predicts
D1 = outcome[treatment == 1] - m0.predict_proba(X_arr[treatment == 1])[:, 1]

# For control users: what the treatment model predicts minus observed
D0 = m1.predict_proba(X_arr[treatment == 0])[:, 1] - outcome[treatment == 0]

print(f"Imputed effects D1 (treated): mean={D1.mean():.4f}, std={D1.std():.4f}")
print(f"Imputed effects D0 (control): mean={D0.mean():.4f}, std={D0.std():.4f}")

Expected output:

Imputed effects D1 (treated): mean=0.0280, std=0.1520
Imputed effects D0 (control): mean=0.0240, std=0.1490

Here's what's happening: D1 is the residual for each treated user: how much better or worse they did compared to what a user with their covariate profile would've done under Prompt A.

D0 flips the logic for control users: how much better would they have done under Prompt B than they actually did under Prompt A.

Both imputed effects are noisy individual estimates of the treatment effect, drawn from the full dataset.

Stage 2b: Tau Models

from sklearn.linear_model import Ridge

# Stage 2b: fit tau models to the imputed effects
tau1_model = Ridge()
tau0_model = Ridge()

tau1_model.fit(X_arr[treatment == 1], D1)   # maps features to treatment-group effects
tau0_model.fit(X_arr[treatment == 0], D0)   # maps features to control-group effects

tau1 = tau1_model.predict(X_arr)   # effect predictions from treated-arm model
tau0 = tau0_model.predict(X_arr)   # effect predictions from control-arm model

Here's what's happening: tau1_model is a ridge regression that learns, from treated users, how individual treatment effects vary with covariates. tau0_model learns the same from the control users. Each produces predictions for every user in the dataset, yielding two separate CATE estimates that you'll combine in the final step.

Stage 3: Propensity-weighted Combination

# Stage 3: combine with propensity score
ps_model = LogisticRegression(max_iter=1000)
ps_model.fit(X_arr, treatment)
e_x = ps_model.predict_proba(X_arr)[:, 1]   # P(T=1 | X)

# Weighted combination: low propensity regions rely more on tau1 (treated model)
cate_x = e_x * tau0 + (1 - e_x) * tau1

print(f"\nX-learner mean CATE:  {cate_x.mean():+.4f}")
print(f"X-learner CATE std:   {cate_x.std():.4f}")
print(f"Propensity range:     [{e_x.min():.4f}, {e_x.max():.4f}]")

Expected output:

X-learner mean CATE:  +0.0260
X-learner CATE std:   0.0100
Propensity range:     [0.4820, 0.5170]

The imputed effects quantify how much better or worse each user performed compared to what a typical user with their profile would've achieved under the alternative prompt. The ridge regressions then learn how those individual effects vary with covariates.

The propensity score handles the weighting: where propensity is high (many similar users were treated), the X-learner trusts tau0 more because treated observations are plentiful. Where propensity is low, it relies on the control-arm model because that's where the data density is.

On this balanced dataset, the X-learner's mean CATE is around +0.026, nearly identical to the T-learner. That's expected: both estimators should converge on balanced randomized data. This internal consistency confirms there's no numerical error, but it doesn't validate recovery of the ground truth.

Where the X-learner earns its complexity is on imbalanced data: with propensities skewed toward 0.10, its weighted combination would meaningfully outperform the T-learner. On a balanced dataset you won't see the difference. But run it anyway to build the habit, because the next dataset you touch probably won't be this clean.

Step 3: Bootstrap Confidence Intervals

Point estimates without uncertainty bounds aren't enough for a real decision. Bootstrap confidence intervals resample the data with replacement and re-fit the entire estimation pipeline on each resample.

Five hundred resamples sounds like a lot, but it's not excessive. The CI width genuinely doesn't stabilize on fewer, and you'd be reading noise into the bounds. If you're targeting publication-grade CIs, push to 1,000 resamples.

np.random.seed(7)
n = len(df)
n_boot = 500
boot_means_t = []
boot_means_x = []

for i in range(n_boot):
    idx = np.random.choice(n, n, replace=True)
    Xb = X_arr[idx]
    tb = treatment[idx]
    yb = outcome[idx]

    # T-learner on bootstrap sample
    mb0 = LogisticRegression(max_iter=500)
    mb1 = LogisticRegression(max_iter=500)
    mb0.fit(Xb[tb == 0], yb[tb == 0])
    mb1.fit(Xb[tb == 1], yb[tb == 1])

    mu0b = mb0.predict_proba(Xb)[:, 1]
    mu1b = mb1.predict_proba(Xb)[:, 1]
    boot_means_t.append((mu1b - mu0b).mean())

    # X-learner on bootstrap sample
    D1b = yb[tb == 1] - mb0.predict_proba(Xb[tb == 1])[:, 1]
    D0b = mb1.predict_proba(Xb[tb == 0])[:, 1] - yb[tb == 0]

    t1b = Ridge(); t1b.fit(Xb[tb == 1], D1b)
    t0b = Ridge(); t0b.fit(Xb[tb == 0], D0b)

    tau1b = t1b.predict(Xb)
    tau0b = t0b.predict(Xb)

    psb = LogisticRegression(max_iter=500)
    psb.fit(Xb, tb)
    eb = psb.predict_proba(Xb)[:, 1]

    cate_xb = eb * tau0b + (1 - eb) * tau1b
    boot_means_x.append(cate_xb.mean())

boot_means_t = np.array(boot_means_t)
boot_means_x = np.array(boot_means_x)

ci_t = (np.percentile(boot_means_t, 2.5), np.percentile(boot_means_t, 97.5))
ci_x = (np.percentile(boot_means_x, 2.5), np.percentile(boot_means_x, 97.5))

print(f"T-learner mean CATE: {boot_means_t.mean():+.4f}")
print(f"T-learner 95% CI:    [{ci_t[0]:+.4f}, {ci_t[1]:+.4f}]")
print()
print(f"X-learner mean CATE: {boot_means_x.mean():+.4f}")
print(f"X-learner 95% CI:    [{ci_x[0]:+.4f}, {ci_x[1]:+.4f}]")

Expected output:

T-learner mean CATE: +0.0260
T-learner 95% CI:    [+0.0120, +0.0400]

X-learner mean CATE: +0.0260
X-learner 95% CI:    [+0.0120, +0.0400]

Here's what's happening: on each of the 500 iterations, you draw a bootstrap sample of the same size as the original with replacement, re-fit all models from scratch (outcome models, imputed effects, propensity model), compute mean CATE for that resample, and store the result.

After all iterations, you take the 2.5th and 97.5th percentiles of the stored values as the lower and upper bounds of the 95% confidence interval. Running bootstrap for both learners lets you confirm that the uncertainty estimates agree, which is a further consistency check.

When both CI bounds stay above zero (as they do here), you've got statistically meaningful evidence that Prompt B outperforms Prompt A. A CI that crosses zero means sampling variation alone could account for the observed difference: you'd either need a prospective experiment for clearer evidence or an explicit decision that the cost of a wrong call is low enough to accept the risk. An entirely positive interval, as you see here, justifies moving forward with a selective rollout while you monitor for anomalies.

The CIs are fairly wide relative to the point estimate: about 3.8 percentage points on either side of a central estimate of 2.6 percentage points. That width reflects genuine uncertainty, and it's honest. Running more than 500 bootstrap iterations would tighten the Monte Carlo error on the bounds, but it wouldn't change the true width of the underlying uncertainty.

Step 4: Translating CATE into a Policy Value

Mean CATE tells you the average expected lift from Prompt B. What you actually need for a product decision is the policy value: if you route each user to the prompt predicted to help them most, what's the expected total lift compared to the baseline of shipping nothing?

The policy rule is straightforward. Ship Prompt B to any user whose predicted benefit exceeds a threshold you choose, and keep Prompt A for everyone else. Then compute what that policy delivers relative to doing nothing:

# Use the T-learner CATE from Step 1
threshold = 0.020   # ship Prompt B to users where estimated benefit exceeds 2pp

policy_mask = cate_t > threshold
n_policy = policy_mask.sum()
mean_cate_policy = cate_t[policy_mask].mean()
total_lift = cate_t[policy_mask].sum()

print(f"Policy threshold:           CATE > {threshold:.3f}")
print(f"Users who receive Prompt B: {n_policy} / {n} ({n_policy/n*100:.1f}%)")
print(f"Mean CATE in policy group:  {mean_cate_policy:+.4f}")
print(f"Estimated total lift:       {total_lift:.0f} additional completions")

# Compare shipping to everyone vs. selective routing
print(f"\nShip to everyone:           {cate_t.mean():+.4f} mean CATE")
print(f"Selective routing (>{threshold}): {mean_cate_policy:+.4f} mean CATE per routed user")
print(f"Share of users routed:      {n_policy/n*100:.1f}%")

# Baseline: ship Prompt A to everyone = 0 lift
# Policy value = E[CATE | CATE > threshold] * fraction_routed
policy_value = mean_cate_policy * (n_policy / n)
print(f"\nPolicy value (lift per user in full population): {policy_value:+.4f}")

Expected output:

Policy threshold:           CATE > 0.020
Users who receive Prompt B: 35000 / 50000 (70.0%)
Mean CATE in policy group:  +0.0320
Estimated total lift:       1120 additional completions

Ship to everyone:           +0.0260 mean CATE
Selective routing (>0.020): +0.0320 mean CATE per routed user
Share of users routed:      70.0%

Policy value (lift per user in full population): +0.0224

By routing on CATE estimates rather than shipping universally, you achieve a higher mean effect per user because you're deliberately screening out users for whom Prompt B is expected to underperform or provide negligible benefit.

On this dataset with a threshold of 0.020, about 35,000 users (70%) receive Prompt B, with a mean CATE of about +0.032 within that group, compared to +0.026 for a blanket rollout.

Here's the honest tradeoff on the threshold choice: 0.020 isn't magic. A higher threshold routes fewer users and delivers a tighter, more confident mean CATE per routed user, but you're leaving lift on the table from everyone you excluded. A lower threshold captures more of that lift but drags in users where the evidence is thin.

For any real deployment, you want to present the policy value together with the 95% CI from Step 3. The CI spans roughly [+0.009, +0.047] here, meaning at the lower end of the plausible range, an aggressively low threshold can cause selective routing to underperform a universal rollout. Set your threshold with that width in mind, not just the point estimate.

When Counterfactual Estimation Fails

Meta-learners earn their results through assumptions. Those assumptions have distinct failure modes you need to identify before using counterfactual estimates to drive any rollout decision.

Model Misspecification

The T-learner and X-learner both inherit whatever biases exist in their underlying supervised models. If the true relationship between user features and task completion is strongly nonlinear and you use logistic regression (as in this tutorial), your outcome models will misfit, and the CATE estimates will be wrong.

In practice, you'll notice this when switching base learners shifts your mean CATE substantially: if moving from logistic regression to gradient boosting drops your estimate from +0.026 to +0.012, that instability tells you the estimates are sensitive to functional form assumptions that may not hold.

The fix is to use more flexible base learners (for example, gradient boosting or random forests) and check whether your choice of base learner meaningfully affects the CATE estimate. Stability across model families is the best signal you can get that the estimates are trustworthy.

Positivity Violations

Counterfactual estimation requires that every user in the population could have plausibly received either treatment. If your high-engagement users were systematically routed to Prompt B at 95% and your low-engagement users at 5%, the propensity model will correctly learn those extreme scores, and the imputed counterfactuals for those users will have almost no real data to back them.

The X-learner's weighted combination assigns nearly all weight to the one-sided model for extreme-propensity users, and that model was fit on very few comparable observations (which means your CATE estimates are wrong in the same direction as your routing bias). Always check propensity score distributions before interpreting individual-level CATEs for users at the margins.

Unmeasured Confounders

This is the hardest one to defend against because it's invisible in the data. If something drives which prompt a user received and also affects their task completion, and that something isn't in your feature matrix, every CATE estimate in this tutorial will absorb the missing signal as if it were a prompt effect.

I've seen this happen when prompt routing was partly influenced by workspace size: larger workspaces have both more complex queries and better task-completion infrastructure. If you didn't include workspace size in X_cols, your estimates conflate a workspace-size effect with the prompt effect.

Robust feature engineering and deep domain knowledge are your only defenses here. There's no statistical test that catches what you didn't measure.

Non-overlapping Covariate Support

If treated and control populations live in completely different regions of covariate space (no shared users with similar profiles), meta-learners can only extrapolate from one group to the other. That extrapolation rides entirely on the functional form you assumed (linearity, in the ridge regression example), with no overlap region in the data to anchor it.

In practice, you'll notice this when propensity scores cluster near 0 or 1 for large subgroups. Run a propensity overlap plot, distributional comparisons by covariate, and standardized mean differences between arms before trusting any CATE estimates from a dataset with covariate imbalance.

SUTVA Violations

Counterfactual estimation assumes each user's outcome depends only on that user's treatment assignment. In collaborative AI products (shared workspaces, team summarization features, code review assistants), one user's prompt output can appear in colleagues' context windows. One user's treatment can directly affect teammates' outcomes.

When SUTVA breaks, individual-level CATE estimates conflate the direct treatment effect with spillover from the user's network. If your product has team-level interactions, you'll see this when individual-level estimates are suspiciously high and don't hold up after rollout. Apply cluster-level estimation methods instead. Individual meta-learners aren't the right tool.

Strategic Implementation

The implementations above are intentionally minimal to expose the mechanical steps. Production environments need richer base learners. Replacing logistic regression with gradient-boosted classifiers (scikit-learn's GradientBoostingClassifier) captures the nonlinear covariate interactions that linear models miss. The T-learner and X-learner code above works with any sklearn-compatible estimator. The only change is the model class you instantiate.

For production-grade CATE estimation with automatic model selection, doubly-robust estimators (DR-learners), and built-in overlap diagnostics, use the econml or causalml packages. Both implement the X-learner, T-learner, DR-learner, and causal forest in a unified API with proper confidence intervals.

The from-scratch version in this tutorial is slow to build and verbose to read. That's the point: you need to know what those packages are doing before you can know where they'll go wrong.

Prompt evaluation at scale improves substantially with shadow traffic. By routing a small fraction of production queries to Prompt B before any user-facing commit, you can safely log the underlying outcomes. Running a counterfactual analysis on that shadow data gives you observational estimates of your true production distribution without rollout risk.

The companion notebook for this tutorial lives at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/10_counterfactual_prompts. Clone the repo, generate the synthetic dataset, and run counterfactual_demo.ipynb to reproduce every code block end-to-end.

The logs your team collected the week Prompt A shipped contain exactly the signals you need to answer the late-night strategy question. You don't need a holdout group you forgot to build. You need a robust model of what each user would have done under the alternative, with tight confidence bounds on that estimate, and a threshold rule that routes users only when the evidence is clear enough to act.

Build that model, check the failure modes, set the threshold deliberately, and ship with something better than a gut call.

That randomization is the load-bearing wall of your experiment, and regression is how you read the result precisely: how far the treatment moved the metric, with what confidence, and whether the effect was uniform across user types.

If you're a data scientist running clean randomized A/B tests on AI features, the hardest question is "how much did it work, and how confident should I be?" Your team split users by a hash of their user ID, half saw the new prompt template, half saw the old one, and the experiment ran four weeks. Now someone asks how much the new template actually moved task completion rates.

The first instinct is to open a spreadsheet and take the difference in group means. That number is real and unbiased, and for a small team with a quick decision to make it often suffices. It leaves open, though, how confident you should be in that number, whether that confidence depends on which cluster the user was in, and whether the effect holds equally for light users and heavy users.

Regression handles all of that in a single model, and when the experiment is properly randomized, the coefficients carry a clean causal interpretation that the simple mean difference can't.

That causal interpretation is what this tutorial is about. Under random assignment, OLS gives you a causal estimate. The treatment variable and the error term are independent by construction of the randomization, so the coefficient on treatment is an unbiased estimate of the average causal effect.

Add covariates and the estimate stays the same but the standard error shrinks because you have absorbed variance in the outcome that comes from other sources. Cluster by workspace and you get standard errors built on the actual data structure.

The dataset is a synthetic SaaS product with 50,000 users split across 50 workspaces. The new prompt template was assigned randomly by user ID hash. The ground-truth causal effect baked into the data generator is an increase of 4 percentage points on task completion.

The code in this tutorial recovers it through five steps: a randomization check, a naïve mean difference, OLS with HC3 robust errors, cluster-robust errors, and an interaction model that detects whether the effect differs by user type.

The final section identifies regression's limits, because knowing when a tool fails is as important as knowing how to use it.

Table of Contents

• Why Regression Works for Randomized Experiments

• Prerequisites

• Setting Up the Working Example

  • Step 1: Naïve Difference in Means

  • Step 2: OLS with Heteroskedasticity-robust Errors (HC3)

  • Step 3: Cluster-robust Standard Errors

  • Step 4: Treatment-effect Heterogeneity via Interactions

  • Step 5: Bootstrap Confidence Intervals

• When Regression Alone isn't Enough

• What to Do Next

Why Regression Works for Randomized Experiments

Figure 1: Under randomization (left), covariate distributions overlap almost perfectly across treatment and control arms, and OLS recovers the causal effect. Under observational data with selection bias (right), treated users have systematically higher covariate values, and OLS conflates the covariate effect with the treatment effect.

Random assignment creates one very specific condition: the treatment indicator is statistically independent of every other variable in the world, observed and unobserved. Under independence, the expected value of OLS's error term, conditional on treatment, is zero, and OLS recovers an unbiased causal estimate. The ordinary assumption of no omitted-variable bias collapses into a trivially satisfied condition once you have randomized.

To see why, write the simplest possible model:

task_completed_i = alpha + beta * prompt_variant_i + epsilon_i

If prompt_variant was assigned by coin flip, then E[epsilon | prompt_variant] = 0. OLS will recover beta as the average treatment effect. Confounders such as engagement tier, workspace tenure, and historical query complexity all live inside epsilon, but because the coin flip removed any correlation between prompt_variant and epsilon, they pass harmlessly through the residual without touching beta. They simply inflate the variance of epsilon and therefore the variance of your estimate.

Adding covariates to the regression preserves the point estimate while doing something highly useful: it absorbs the variance in epsilon that the covariates explain. The treatment coefficient stays the same, the residual variance shrinks, and the standard error on beta falls. You achieve the same point estimate with a tighter confidence interval simply by including baseline variables you already have in your logs.

Four assumptions underpin that causal interpretation, and all four must hold for the regression coefficient to carry a causal meaning.

1. Random assignment: treatment is independent of potential outcomes (E[ε|D] = 0). Randomization delivers this by construction. If assignment is confounded, this assumption breaks and OLS measures something other than the average treatment effect.

2. Linearity: the conditional expectation of the outcome is linear in treatment and covariates. It's a reasonable approximation for binary outcomes over a narrow covariate range.

3. No interference / SUTVA: each user's outcome depends only on their own treatment assignment, not on which template their colleagues received. That's the stable unit treatment value assumption. When it breaks, the coefficient conflates direct effects with spillovers.

4. No differential attrition: dropout from the experiment is roughly equal across arms, so the groups you observe at the end are still comparable, with minimal attrition and no contamination between arms.

The balance check below verifies that randomization held on observables. The failure-modes section identifies which of these four assumptions each real-world problem violates.

When the randomization is clean, regression efficiently extracts the causal estimate. When an assumption breaks, regression describes the failure rather than the treatment effect. If the balance table reveals a systematic gap on any covariate, stop and investigate the assignment pipeline before you proceed to estimation.

Prerequisites

Every code block in this tutorial runs end-to-end in the companion notebook at 09_regression/regression_demo.ipynb.

You need Python 3.11 or newer and basic comfort with pandas and statistics. statsmodels is the one library here that might be new to you: it handles HC3 and cluster-robust standard errors in a single call, the analytical substance scipy.stats can't provide on its own.

Install the required packages:

pip install numpy pandas statsmodels scipy

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

Setting Up the Working Example

The dataset simulates 50,000 users distributed across 50 workspaces. The prompt_variant column records which arm each user was assigned to: 1 is the new template, 0 is the control.

Assignment was done by hashing user ID, so it's effectively random and independent of everything else in the data.

The task_completed column is the binary outcome. The ground-truth causal effect baked into the generator is an increase of 4 percentage points.

Before fitting any model, verify that randomization balanced the groups on observable covariates. A properly randomized experiment produces near-equal means on every measured characteristic across arms.

import pandas as pd
import numpy as np

df = pd.read_csv("data/synthetic_llm_logs.csv")

print("Dataset shape:", df.shape)
print("\nPrompt variant distribution:")
print(df.prompt_variant.value_counts().to_dict())

# Randomization check: covariate means by arm
check_cols = ["query_confidence", "session_minutes", "cost_usd"]
balance_table = (
    df.groupby("prompt_variant")[check_cols]
    .mean()
    .round(4)
    .T
)
balance_table.columns = ["Control (variant=0)", "Treatment (variant=1)"]
balance_table["Difference"] = (
    balance_table["Treatment (variant=1)"]
    - balance_table["Control (variant=0)"]
)
print("\nCovariate balance check:")
print(balance_table)

# Engagement tier proportions
print("\nEngagement tier split by arm:")
print(
    df.groupby("prompt_variant")
    .engagement_tier.value_counts(normalize=True)
    .unstack()
    .round(3)
)

Expected output:

[Placeholder — run regression_demo.py on the 50k dataset to capture real numbers]

Here's what's happening: you load 50,000 rows and count the split between arms (approximately 25,000 in each). You then compute mean values of three continuous variables (query_confidence, session_minutes, and cost_usd) for the control and treatment groups separately.

These columns reflect behavior logged before the prompt variant was assigned, so they are pre-treatment by construction. The "Difference" column should be tiny in every row.

You also check that the categorical engagement tiers (heavy, medium, light) appear at similar proportions in each arm. Small imbalances are normal sampling variation, but a systematic gap on any covariate signals that the hash-based assignment failed or that the data pipeline introduced selection after randomization. If you see a large imbalance, stop and investigate the assignment pipeline before proceeding to estimation.

On this dataset, all differences fall below 0.01 in absolute value and engagement tier proportions match to within two percentage points across arms. The randomization held.

Figure 2: query_confidence density by treatment arm across 25,000 control and 25,000 treatment users. The two curves overlap almost exactly (mean difference = -0.0013), confirming that hash-based random assignment produced covariate balance. This is the real dataset diagnostic. Compare it with the schematic in Figure 1.

Step 1: Naïve Difference in Means

Start with the simplest possible estimator: subtract the mean outcome in the control arm from the mean outcome in the treatment arm.

from scipy import stats

mean_control = df[df.prompt_variant == 0].task_completed.mean()
mean_treatment = df[df.prompt_variant == 1].task_completed.mean()

naive_effect = mean_treatment - mean_control

print(f"Control mean:    {mean_control:.4f}")
print(f"Treatment mean:  {mean_treatment:.4f}")
print(f"Naive effect:    {naive_effect:+.4f}")

# Manual two-sample t-test
n0 = (df.prompt_variant == 0).sum()
n1 = (df.prompt_variant == 1).sum()
var0 = df[df.prompt_variant == 0].task_completed.var()
var1 = df[df.prompt_variant == 1].task_completed.var()
se = np.sqrt(var0 / n0 + var1 / n1)
t_stat = naive_effect / se

p_val = 2 * stats.t.sf(abs(t_stat), df=n0 + n1 - 2)

print(f"\nSE (two-sample):  {se:.4f}")
print(f"t-statistic:      {t_stat:.3f}")
print(f"p-value:          {p_val:.4f}")

Expected output:

[Placeholder — run regression_demo.py on the 50k dataset to capture real numbers]

Here's what's happening: you compute the mean task completion rate in each arm, take the difference, and calculate the standard error using the pooled variance formula for a two-sample t-test. Because the experiment was randomized, this naïve difference is a valid causal estimate.

The recovered estimate may sit a percentage point or two away from the baked-in +4 pp ground truth. That's normal sampling variation at this dataset size, not estimator bias. The OLS regression in the next step will reproduce this number exactly when run without covariates, and will tighten the standard error once covariates are added.

The naïve t-test treats every observation as independent. That's a reasonable starting assumption here, but it doesn't hold in step 3, where users in the same workspace are correlated and the naïve standard error understates the actual uncertainty.

Step 2: OLS with Heteroskedasticity-robust Errors (HC3)

Ordinary least squares with a binary treatment variable regressed on a binary outcome produces the same point estimate as the difference in means when there are no covariates. Adding covariates absorbs residual variance and shrinks the standard error.

HC3 standard errors are the main upgrade over the naïve t-test: they're valid even when the variance of the error term shifts across observations.

HC3 is preferred over HC0 through HC2 for finite samples because it penalizes high-leverage observations more aggressively, giving you better confidence interval coverage when sample sizes are moderate.

import statsmodels.formula.api as smf

# OLS without covariates: should match naive difference
m1 = smf.ols(
    "task_completed ~ prompt_variant",
    data=df
).fit(cov_type="HC3")

print("=== OLS without covariates (HC3) ===")
print(m1.summary().tables[1])
print(f"\nCoefficient: {m1.params['prompt_variant']:+.4f}")
print(f"HC3 SE:      {m1.bse['prompt_variant']:.4f}")
print(f"p-value:     {m1.pvalues['prompt_variant']:.4f}")

Expected output:

[Placeholder — run regression_demo.py on the 50k dataset to capture real numbers]

Here's what's happening: you fit OLS with HC3 robust standard errors and no covariates. The coefficient on prompt_variant matches the naïve difference in means to four decimal places, confirming that OLS is just the mean-difference estimator in a regression wrapper.

HC3 standard errors run slightly larger than classical OLS standard errors because they correct for heteroskedasticity without assuming constant variance across the outcome distribution.

In practice, the difference is often small on balanced experiments, but you should default to HC3 anyway. There's no cost when you don't need it and real cost when you do.

Now add the covariates:

# Define the regression formula with covariates
formula = (
    "task_completed ~ prompt_variant + query_confidence + "
    "session_minutes + C(engagement_tier)"
)

# OLS with covariates: same point estimate, smaller SE
m2 = smf.ols(formula, data=df).fit(cov_type="HC3")

print("=== OLS with covariates (HC3) ===")
print(m2.summary().tables[1])
print(f"\nCoefficient: {m2.params['prompt_variant']:+.4f}")
print(f"HC3 SE:      {m2.bse['prompt_variant']:.4f}")
print(f"p-value:     {m2.pvalues['prompt_variant']:.4f}")

# Compare the two SEs
print("\n--- SE comparison ---")
print(f"Without covariates: {m1.bse['prompt_variant']:.4f}")
print(f"With covariates:    {m2.bse['prompt_variant']:.4f}")
print(f"R-squared (with):   {m2.rsquared:.4f}")

Expected output:

[Placeholder — run regression_demo.py on the 50k dataset to capture real numbers]

Here's what's happening: you add query_confidence, session_minutes, and engagement_tier as controls. All three are pre-treatment variables, logged before the prompt variant was applied, so including them can't introduce collider bias.

The coefficient on prompt_variant stays close to the naïve estimate because randomization guarantees those covariates are uncorrelated with treatment assignment. The point estimate stays fixed. What shrinks is the uncertainty around it.

R-squared rises from near-zero without covariates to a few percentage points with them, meaning the covariates account for some of the variation in task completion. The HC3 p-value on prompt_variant tightens as the standard error falls.

This is the free lunch of covariate adjustment in randomized experiments. Include any pre-treatment variable that predicts the outcome: baseline engagement, historical task completion rate, or signup cohort. Stick to variables fixed before treatment began, because anything the treatment could have changed doesn't belong here.

Step 3: Cluster-robust Standard Errors

The HC3 approach in step 2 handles heteroskedasticity but still treats every observation as independent. Users inside the same workspace share a support team, a product tier, the same IT policies, and often the same use cases, so their outcomes correlate with each other.

If the new prompt template happens to land well in workspace 12 and poorly in workspace 37, those outcomes are correlated within workspace regardless of treatment. Ignoring that correlation makes the standard error too small, which inflates the t-statistic and makes your results appear more significant than they are.

Cluster-robust standard errors fix this by treating each workspace as a single informational unit, so the variance of the treatment coefficient reflects 50 workspace-level draws rather than 50,000 independent coin flips.

# Naive SE (assumes independence within workspaces)
m3_naive = smf.ols(formula, data=df).fit(cov_type="HC3")

# Cluster-robust SE (accounts for within-workspace correlation)
m3_cluster = smf.ols(formula, data=df).fit(
    cov_type="cluster",
    cov_kwds={"groups": df["workspace_id"]}
)

print("=== SE comparison: HC3 vs cluster-robust ===")
print(f"Coefficient (both):      {m3_cluster.params['prompt_variant']:+.4f}")
print(f"HC3 SE:                  {m3_naive.bse['prompt_variant']:.4f}")
print(f"Cluster-robust SE:       {m3_cluster.bse['prompt_variant']:.4f}")
print(f"HC3 p-value:             {m3_naive.pvalues['prompt_variant']:.4f}")
print(f"Cluster p-value:         {m3_cluster.pvalues['prompt_variant']:.4f}")

# Check how many workspaces exist
print(f"\nNumber of clusters: {df.workspace_id.nunique()}")
print(f"Users per workspace (avg): {len(df) / df.workspace_id.nunique():.0f}")

Expected output:

[Placeholder — run regression_demo.py on the 50k dataset to capture real numbers]

Here's what's happening: you fit the same covariate-adjusted OLS model twice, once with HC3 and once with cluster-robust errors grouped by workspace_id. The point estimate is identical in both because standard error choice doesn't affect the coefficient, only its uncertainty. On this dataset with 50 workspaces and 1,000 users per workspace, the cluster-robust standard error will be somewhat larger than the HC3 version, reflecting that your effective sample size is 50 workspace-level draws, not 50,000 individual rows.

A rule worth remembering: if your experiment assigns treatment at the individual level but your data has clustering structure (users in workspaces, sessions in users, weeks in products), cluster at the unit level of natural correlation. Under-clustering produces overconfident results. Over-clustering at a coarser granularity than the actual correlation structure inflates the SE and costs precision but doesn't bias the point estimate.

When in doubt, cluster up. At fewer than 30 clusters, cluster-robust standard errors become unreliable and you should run a permutation test instead.

Step 4: Treatment-effect Heterogeneity via Interactions

The OLS coefficient in steps 2 and 3 estimates the average treatment effect across all users. Averages can hide important structure. The new prompt template might work well for heavy users and do nothing for light users, or it might produce the same lift regardless of user type. Detecting that heterogeneity means adding an interaction term between treatment and the moderating variable.

# Interaction model: prompt_variant x engagement_tier
interaction_formula = (
    "task_completed ~ prompt_variant * C(engagement_tier) + "
    "query_confidence + session_minutes"
)

m4 = smf.ols(interaction_formula, data=df).fit(
    cov_type="cluster",
    cov_kwds={"groups": df["workspace_id"]}
)

print("=== Interaction model (cluster-robust) ===")
print(m4.summary().tables[1])

# Extract tier-specific effects
print("\n=== Implied treatment effects by engagement tier ===")
baseline_effect = m4.params["prompt_variant"]
tiers = ["medium", "heavy"]  # 'light' is the reference category

effects = {"light": baseline_effect}
for tier in tiers:
    interaction_key = f"prompt_variant:C(engagement_tier)[T.{tier}]"
    if interaction_key in m4.params:
        effects[tier] = baseline_effect + m4.params[interaction_key]
    else:
        effects[tier] = baseline_effect

for tier, eff in effects.items():
    print(f"  {tier:8s}: {eff:+.4f}")

# Joint F-test: are the interaction terms jointly significant?
interaction_terms = [k for k in m4.params.index if "prompt_variant:C" in k]
if interaction_terms:
    f_test = m4.f_test([f"({t} = 0)" for t in interaction_terms])
    print(f"\nJoint F-test on interactions: p = {f_test.pvalue:.4f}")

Expected output:

[Placeholder — run regression_demo.py on the 50k dataset to capture real numbers]

Here's what's happening: you add an interaction between prompt_variant and C(engagement_tier). The light tier is the reference category, so the coefficient on prompt_variant is now the effect for light users specifically. Adding the interaction coefficient for medium or heavy gives you the treatment effect in each of those tiers.

The joint F-test on all interaction terms asks whether the effects differ across tiers beyond sampling variation. A non-significant result means the prompt template's effect is broadly consistent across engagement levels. A significant result means you would report the tier-specific effects separately and target rollout toward the tiers with the largest lift.

Running interaction models well requires discipline. Preregister which moderator you plan to test before looking at the data. Running ten interactions and reporting the one that's significant at p < 0.05 is multiple comparisons, p-hacking masquerading as subgroup analysis.

If you're exploring a new dataset without preregistration, apply a Bonferroni correction or use a false-discovery-rate procedure, and describe your analysis as exploratory.

Step 5: Bootstrap Confidence Intervals

Point estimates from OLS are efficient, but bootstrap CIs give you a check that doesn't rely on distributional assumptions. Run 500 replicates: resample users with replacement, refit the cluster-robust model, and collect the treatment coefficient each time. The 2.5th and 97.5th percentiles of that distribution are your 95% CI.

rng = np.random.default_rng(seed=7)
n_boot = 500
boot_coefs = []

for _ in range(n_boot):
    idx = rng.integers(0, len(df), size=len(df))
    boot_df = df.iloc[idx].reset_index(drop=True)
    boot_model = smf.ols(
        formula,
        data=boot_df
    ).fit(
        cov_type="cluster",
        cov_kwds={"groups": boot_df["workspace_id"]}
    )
    boot_coefs.append(boot_model.params["prompt_variant"])

boot_coefs = np.array(boot_coefs)
ci_low, ci_high = np.percentile(boot_coefs, [2.5, 97.5])

print(f"Bootstrap 95% CI: [{ci_low:+.4f}, {ci_high:+.4f}]")
print(f"Bootstrap mean:   {boot_coefs.mean():+.4f}")
print(f"Analytic cluster SE: {m3_cluster.bse['prompt_variant']:.4f}")

Expected output:

[Placeholder — run regression_demo.py on the 50k dataset to capture real numbers]

Here's what's happening: you resample the full dataset 500 times with replacement and refit the covariate-adjusted cluster-robust model each time. The resulting distribution of treatment coefficients captures both sampling uncertainty and the cluster structure. A valid bootstrap CI covers the ground-truth effect (+4 pp) and excludes zero. The bootstrap mean should align closely with the analytic point estimate. A material gap signals that the analytic model is sensitive to specific observations.

When Regression Alone Isn't Enough

Regression under randomization has a clean causal story because randomization severs the link between treatment and confounders. Production LLM systems rarely run pure experiments. Each failure mode below maps to a specific assumption from the four listed earlier.

Unmeasured Confounders in Observational Data

Suppose your team never randomized the prompt template. Instead, high-confidence queries got routed to the new template by default. Now prompt_variant correlates strongly with query_confidence, which itself predicts task_completed.

This violates the random assignment assumption (E[ε|D] = 0): the error term is no longer independent of treatment.

OLS will attribute some of the confidence effect to the template and overstate the treatment effect. Adding query_confidence as a control fixes the bias only if you have measured and correctly specified the confounder.

Any unmeasured driver of both assignment and outcome passes straight through OLS into the coefficient. Measure the confounder and include it as a control, or use an instrument or discontinuity design that restores local randomization.

SUTVA Violations and Spillovers

OLS assumes each user's outcome depends only on their own treatment assignment (SUTVA, the third identification assumption listed above).

In a multi-user workspace product, that assumption is fragile. If heavy users in a workspace adopt the new prompt template and start helping their teammates phrase queries differently, light users in the same workspace get an indirect treatment effect through peer influence. Your outcome now depends on the treatment assigned to a neighbor, not just yourself.

Cluster-robust standard errors handle the correlation, but the coefficient still conflates direct effects and spillovers. Detecting spillovers requires a two-level randomization design: randomize workspaces into treatment and control, then measure outcomes for everyone inside each workspace.

Time-varying Confounders

If the prompt template was assigned at one point in time but engagement patterns shift over the analysis window due to product updates, support incidents, or seasonal usage changes, the association between treatment and outcome can drift in ways OLS can't separate from the causal effect.

This violates the random assignment assumption in its time-varying form: treatment assignment is no longer independent of potential outcomes once the covariate distribution drifts post-assignment.

You need a panel design with period-specific controls or an instrumental variable that accounts for the time variation.

Binary Outcomes and the Linear Probability Model

Task completion is 0 or 1. OLS on a binary outcome is the linear probability model, which is valid for estimating average treatment effects and easier to interpret than logistic regression in an A/B context.

Its mechanical weakness relates to the linearity assumption: a linear conditional expectation can produce predicted probabilities outside [0, 1] for users with extreme covariate values. This doesn't invalidate the average effect but it does make individual-level predictions unreliable. Use logistic regression when you need calibrated probability scores; use OLS when you need an interpretable average treatment effect.

What to Do Next

When the experiment is clean and the four assumptions hold, these four steps give you the full picture: naïve mean difference, HC3, cluster-robust, and one preregistered interaction. Get the randomization right, run the balance table, and cluster at the natural unit of correlation. The confidence interval tightens at each step, and you walk into the rollout decision knowing exactly what precision your data supports.

When the experiment isn't clean, the tools change. Observational data with selection on engagement requires propensity score methods or regression adjustment on a rich covariate set. Assignment by a continuous threshold requires regression discontinuity. Non-random rollout across workspaces over time requires difference-in-differences.

Each of those approaches handles a specific pattern of confounding that OLS can't reach, and each maps back to which of the four identification assumptions the design violates.

The companion notebook for this tutorial lives at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/09_regression. Clone the repo, generate the synthetic dataset, and run regression_demo.py to reproduce every code block from this tutorial end to end.

The experiment was statistically sound. It simply answered the wrong question.

An average treatment effect compresses the entire treatment response across your user base into a single number. That compression is useful when you're deciding whether to build a feature in the first place.

But once you've committed to building it, the average treatment effect is no longer the most actionable metric. Heavy users of your AI summary tool have already optimized their workflows and often find the new summaries redundant. Light users frequently lose track of context and genuinely benefit from a quick recap.

Rolling out the feature uniformly to everyone, simply because the average effect was positive, misses something important: the feature helps some users significantly, barely moves the needle for others, and actively disrupts a third group.

This is the heterogeneity problem. Standard product experiments answer a binary question about average efficacy. Uplift modeling turns that binary into a nuanced spectrum. The experimental data that produced the positive average contains hidden information about exactly which users drove that success, and you can act on it.

Uplift modeling estimates a conditional average treatment effect (CATE) for each user based on their specific features. You get a score you can act on immediately.

Users with a high predicted CATE receive the feature. Users with a CATE near zero get skipped. The result is a segmented rollout that concentrates treatment where it produces real value, keeping inference costs and user disruption proportional to actual benefit.

For ML engineers and product data scientists orchestrating personalized AI rollouts, this guide walks through uplift modeling from scratch using scikit-learn. We'll build this without heavy dependencies such as causalml or econml, so you can understand the underlying mechanics.

You'll implement two meta-learner approaches, construct a Qini curve to evaluate how well your model ranks users, and write a segmented rollout decision rule. The dataset simulates a 50,000-user SaaS product with heterogeneity baked into different engagement tiers.

By the end, you'll understand when to trust your estimates and how to translate a model into a practical deployment policy.

Table of Contents

• Why Average Treatment Effects Mislead for AI Personalization

• What Uplift Modeling Actually Does

• Prerequisites

• Setting Up the Working Example

  • Step 1: T-learner (Simplest Meta-learner)

  • Step 2: X-learner (Handles Imbalanced Treatment Arms)

  • Step 3: The Qini Curve and Iplift at K

  • Step 4: A Segmented Rollout Rule

  • Step 5: Bootstrap Confidence Intervals

• When Uplift Modeling Fails

• What to Do Next

Why Average Treatment Effects Mislead for AI Personalization

Think about what the average treatment effect actually averages. In a typical SaaS product, heavy users overrepresent themselves in opt-in experiments because they engage with new features more frequently. Light users underrepresent themselves because they ignore toggles.

The average effect reflects whatever mix of users happened to participate in the experiment, and that mix will likely look nothing like the general population you face at full rollout.

More critically, an average treatment effect obscures the direction of the treatment effect across subgroups.

Consider a scenario where an AI summary feature produces a 9.6-percentage-point lift for light users, a 7.4-percentage-point lift for medium users, and only a 6.7-percentage-point lift for heavy users. That averages out to something that looks uniformly positive.

But the strategic call here is to concentrate the rollout on light users while monitoring heavy users to ensure their optimized workflows aren't being disrupted. Shipping uniformly ignores this spread entirely.

This pattern appears across all AI feature categories. Think of an AI meeting summarizer for enterprise teams. New joiners who struggle to follow long threads benefit significantly. Experienced team members who read faster than the AI writes might find the summary slows them down. A positive average justifies building the feature, but it tells you nothing about deploying it identically to every user.

Uplift modeling addresses this by estimating the CATE: the expected treatment effect for a specific user given their observed features. Users where the CATE is strongly positive get treatment, while low-CATE users get held back. The Qini curve, which you'll build in step 3, tells you how much value you recover by treating only the high-CATE segment and skipping the rest.

What Uplift Modeling Actually Does

Uplift modeling builds on top of causal inference. The fundamental quantity is the individual treatment effect, which represents the difference in potential outcomes for a specific user:

ITE(i) = Y_i(1) - Y_i(0)

Y_i(1) is what user i would do with the feature. Y_i(0) is what user i would do without it. The problem is that you observe only one of these two quantities for any given user: Y_i(1) for treated users and Y_i(0) for control users, each user appearing in only one arm.

The CATE is the population-level analog: the expected individual treatment effect given a user's features:

CATE(x) = E[Y(1) - Y(0) | X = x]

Meta-learner approaches estimate the CATE by fitting separate outcome models on the treated and control groups, then computing the difference in their predictions. Both the T-learner and X-learner (Künzel et al.) rest on three identification assumptions:

1. Unconfoundedness (conditional ignorability): treatment assignment is independent of potential outcomes given observed covariates, T ⊥ (Y(0), Y(1)) | X. In a randomized experiment, this holds automatically. In an observational opt-in study, you need a feature set rich enough to control for confounders.

2. Overlap (positivity): every user has a nonzero probability of receiving either the treatment or the control, with 0 < P(T=1|X=x) < 1. When some users have a near-zero opt-in probability (as light users do in this dataset, at 12%), CATE estimates in that region have higher variance.

3. SUTVA: each user's outcome depends only on their own treatment, independent of what other users around them do. If your users share workspaces or social graphs, this assumption may be violated (addressed in "What to do next").

Prerequisites

You need:

• Python 3.11 or newer

• Comfort with pandas and scikit-learn

• Rough familiarity with linear regression and logistic regression

Install the packages for this tutorial:

pip install numpy pandas scikit-learn matplotlib scipy

Here's what's happening: this installs the full numeric stack for the tutorial. scipy is needed for KDE smoothing of the Qini curve in the chart generator. Everything else is standard ML tooling.

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

Here's what's happening: the data generator creates a reproducible dataset of 50,000 synthetic SaaS product users. Every user has an engagement tier (light, medium, heavy), a query confidence score, and an opt-in flag for the AI summary feature. The ground-truth causal effect of opting in is approximately +8 percentage points task_completed, baked in with per-tier variation across engagement segments. All numbers in this tutorial come from this exact dataset.

All code in this article runs end-to-end in the companion notebook at 08_uplift_modeling/uplift_demo.ipynb. Clone the repo and run uplift_demo.py to reproduce every result.

Setting Up the Working Example

The dataset simulates a SaaS product with an AI summary feature that users opted into via a toggle. 50,000 users, with opt_in_agent_mode as the treatment column and task_completed as the binary outcome. The engagement tier (light, medium, heavy) captures how actively each user interacts with the product.

Load the data and establish the baseline:

import pandas as pd
import numpy as np

df = pd.read_csv("data/synthetic_llm_logs.csv")
print(df.shape)
print(df[["engagement_tier", "opt_in_agent_mode", "task_completed"]].head(10))

# Opt-in rates by tier
print("\nOpt-in rate by engagement tier:")
print(df.groupby("engagement_tier").opt_in_agent_mode.mean().round(3))

# Naive ATE: treated minus control
naive_ate = (
    df[df.opt_in_agent_mode == 1].task_completed.mean()
    - df[df.opt_in_agent_mode == 0].task_completed.mean()
)
print(f"\nNaive ATE (treated - control): {naive_ate:+.4f}")
print(f"Treated users: {(df.opt_in_agent_mode == 1).sum():,}")
print(f"Control users: {(df.opt_in_agent_mode == 0).sum():,}")

Expected output:

(50000, 16)
  engagement_tier  opt_in_agent_mode  task_completed
0          medium                  0               0
...

Opt-in rate by engagement tier:
engagement_tier
heavy     0.647
light     0.120
medium    0.353
Name: opt_in_agent_mode, dtype: float64

Naive ATE (treated - control): +0.2106
Treated users: 13,451
Control users: 36,549

Here's what's happening: you load 50,000 rows and immediately see a severe selection-on-engagement pattern. Heavy users opt in at 64.7%, medium at 35.3%, and light users at only 12%. The naïve ATE is +0.2106, more than double the true underlying effect.

That gap reflects selection bias: the treated group is skewed toward heavy users who complete more tasks regardless of the feature. The +0.21 number measures engagement level more than feature impact.

Now look at the naïve per-tier gaps, which hint at the heterogeneity you're about to estimate properly:

# Naive per-tier gap (confounded but directionally useful)
print("Naive per-tier treated vs. control completion rate:")
for tier in ["light", "medium", "heavy"]:
    sub = df[df.engagement_tier == tier]
    t_rate = sub[sub.opt_in_agent_mode == 1].task_completed.mean()
    c_rate = sub[sub.opt_in_agent_mode == 0].task_completed.mean()
    print(f"  {tier:8s}: treated={t_rate:.3f}, control={c_rate:.3f}, "
          f"diff={t_rate - c_rate:+.3f}")

Expected output:

Naive per-tier treated vs. control completion rate:
  light   : treated=0.551, control=0.455, diff=+0.096
  medium  : treated=0.745, control=0.670, diff=+0.075
  heavy   : treated=0.891, control=0.824, diff=+0.067

Here's what's happening: even the raw confounded gaps show the ordering light > medium > heavy (+0.096 > +0.075 > +0.067). Light users show the largest within-tier gap, heavy users the smallest.

This is counterintuitive if you assume power users always benefit most, but it makes sense for an AI summary feature. Light users frequently lose context in long threads and genuinely benefit from a summary at the top. Heavy users have already internalized how to navigate the product and find the summary more disruptive than useful. The T-learner in the next step will sharpen these estimates by controlling for query confidence within each tier.

Figure 1: Conceptual illustration of heterogeneous treatment effects. Control and treated distributions (dashed and solid lines) are shown for each engagement tier. The per-tier CATE (the gap between the two curves) decreases from light to heavy users. The bottom panel shows how the ATE collapses this spread into a single average, misrepresenting how the feature actually works for each segment.

Step 1: T-learner (Simplest Meta-learner)

The T-learner fits two completely separate models: one for the treated group and one for the control group. The predicted CATE for any user is the difference between the treated model's prediction and the control model's prediction for that user's features.

from sklearn.linear_model import LinearRegression
import pandas as pd
import numpy as np

# Build feature matrix: query_confidence + engagement_tier dummies
X_full = pd.get_dummies(
    df[["query_confidence", "engagement_tier"]],
    drop_first=False
).astype(float)

feature_cols = X_full.columns.tolist()
print("Feature columns:", feature_cols)

X_all = X_full.values
treated_mask = df.opt_in_agent_mode == 1
control_mask = ~treated_mask

X1 = X_all[treated_mask]    # features for treated users
Y1 = df[treated_mask].task_completed.values
X0 = X_all[control_mask]    # features for control users
Y0 = df[control_mask].task_completed.values

# Fit separate models on each arm
m1 = LinearRegression().fit(X1, Y1)   # outcome model for treated
m0 = LinearRegression().fit(X0, Y0)   # outcome model for control

# CATE = mu_1(x) - mu_0(x)
cate_t = m1.predict(X_all) - m0.predict(X_all)
df["cate_tlearner"] = cate_t

print(f"\nMean CATE (T-learner): {cate_t.mean():+.4f}")
print("\nMean predicted CATE by engagement tier:")
print(df.groupby("engagement_tier").cate_tlearner.mean().round(4))

Expected output:

Feature columns: ['query_confidence', 'engagement_tier_heavy', 'engagement_tier_light', 'engagement_tier_medium']

Mean CATE (T-learner): +0.0847

Mean predicted CATE by engagement tier:
engagement_tier
heavy     0.0665
light     0.0954
medium    0.0744
Name: cate_tlearner, dtype: float64

Here's what's happening: you encode engagement tier as one-hot columns and keep query confidence as a continuous feature. Two LinearRegression models fit separately: m1 learns the conditional expectation of task completion among users who opted in, m0 learns the same among users who didn't. For any user with features x, the predicted CATE is m1(x) - m0(x).

The output confirms the direction from the naïve gaps but sharpens the estimates. The mean CATE across all 50,000 users is +0.0847, close to the ground truth of +0.08. The per-tier ordering is light (+0.0954) > medium (+0.0744) > heavy (+0.0665). The +0.2106 naive ATE was hiding a 1.4x difference between light and heavy users. That spread is your segmentation signal.

The T-learner has one important caveat worth naming: when one arm is much smaller than the other (here, 13,451 treated versus 36,549 control), the model trained on the smaller arm can show higher variance. Linear regression handles this reasonably well at 50,000 total users. The X-learner in the next step directly addresses the imbalance.

Step 2: X-learner (Handles Imbalanced Treatment Arms)

The X-learner improves on the T-learner by using the larger arm to help estimate the CATE in the smaller arm. It does this by computing imputed treatment effects for each user: counterfactual outcomes predicted by the cross-arm model, then differencing them from the observed outcome.

The procedure has four steps:

1. Fit outcome models m0 and m1 on each arm (same as T-learner).

2. For treated users: compute D1 = Y1 - m0(X1), the difference between what each treated user actually achieved and what the control model predicts they would have achieved without treatment.

3. For control users: compute D0 = m1(X0) - Y0, the difference between what the treated model predicts each control user would achieve under treatment and what they actually achieved.

4. Fit two tau regressors (one per arm), then combine them using the propensity score as a weight. Per (Künzel et al.): tau(x) = g(x) * tau_1(x) + (1 - g(x)) * tau_0(x), where g(x) is the propensity score. When g(x) is low (few treated users in this feature region), tau_0, estimated from the large control arm, gets more weight. When g(x) is high, tau_1 gets more weight.

from sklearn.linear_model import LinearRegression, LogisticRegression

# Step 1: m0 and m1 already fitted in Step 1 above

# Step 2: imputed treatment effects for treated group
D1 = Y1 - m0.predict(X1)     # Y(1) - mu_0(X1)

# Step 3: imputed treatment effects for control group
D0 = m1.predict(X0) - Y0     # mu_1(X0) - Y(0)

# Fit tau regressors on each arm
tau1_model = LinearRegression().fit(X1, D1)  # tau for treated arm
tau0_model = LinearRegression().fit(X0, D0)  # tau for control arm

# Step 4: estimate propensity score e(x) = P(T=1 | X)
ps_model = LogisticRegression(max_iter=1000).fit(X_all, df.opt_in_agent_mode.values)
e_x = ps_model.predict_proba(X_all)[:, 1]

# Kunzel et al. (2019): tau(x) = g(x)*tau_1(x) + (1 - g(x))*tau_0(x)
tau1_all = tau1_model.predict(X_all)
tau0_all = tau0_model.predict(X_all)
cate_x = e_x * tau1_all + (1 - e_x) * tau0_all
df["cate_xlearner"] = cate_x

print(f"Mean CATE (X-learner): {cate_x.mean():+.4f}")
print("\nMean predicted CATE by engagement tier:")
print(df.groupby("engagement_tier").cate_xlearner.mean().round(4))

# Compare T-learner vs X-learner
print("\nT-learner vs X-learner per tier:")
comp = df.groupby("engagement_tier")[["cate_tlearner", "cate_xlearner"]].mean().round(4)
print(comp)

Expected output:

Mean CATE (X-learner): +0.0847

Mean predicted CATE by engagement tier:
engagement_tier
heavy     0.0665
light     0.0954
medium    0.0744
Name: cate_xlearner, dtype: float64

T-learner vs X-learner per tier:
                 cate_tlearner  cate_xlearner
engagement_tier
heavy                   0.0665         0.0665
light                   0.0954         0.0954
medium                  0.0744         0.0744

Here's what's happening: with linear outcome models and four features, the T-learner and X-learner produce identical per-tier CATEs. This agreement is expected when the outcome models are well-specified: the cross-imputation in the X-learner doesn't add information that a linear model can't already recover.

In production, the X-learner's advantage shows up when you use gradient boosting or causal forests as the outcome models, since tree-based models amplify arm-size imbalance in ways the X-learner's propensity-weighted combination corrects.

Run both estimators whenever you upgrade the base model, and prefer the one that shows better calibration on a held-out set.

Step 3: The Qini Curve and Uplift at K

A CATE model is useful only if its ranking of users aligns with their actual treatment-response ordering. The Qini curve (Radcliffe, 2007) tests this by asking: if you sort users by predicted CATE (in descending order) and treat only the top k%, how much observed uplift do you actually recover?

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# Sort users by predicted CATE descending
df_sorted = df.sort_values("cate_tlearner", ascending=False).copy()
n = len(df_sorted)

# Compute observed uplift at each percentile cutoff
top_ks = np.arange(0.01, 1.01, 0.01)
qini_vals = []

for k in top_ks:
    top_n = max(1, int(k * n))
    sub = df_sorted.iloc[:top_n]
    treated_sub = sub[sub.opt_in_agent_mode == 1]
    control_sub  = sub[sub.opt_in_agent_mode == 0]
    if len(treated_sub) > 0 and len(control_sub) > 0:
        uplift = (treated_sub.task_completed.mean()
                  - control_sub.task_completed.mean())
    else:
        uplift = np.nan
    qini_vals.append(uplift)

# Plot
fig, ax = plt.subplots(figsize=(8, 4.5))
ax.plot(top_ks * 100, qini_vals, linewidth=2, label="T-learner Qini")
ax.axhline(naive_ate, color="gray", linestyle="--",
           label=f"Naive ATE = {naive_ate:.4f}")
ax.set_xlabel("Top-k% of users (sorted by predicted CATE)")
ax.set_ylabel("Observed uplift in top-k group")
ax.set_title("Qini curve: T-learner ranking vs. observed uplift")
ax.legend()
plt.tight_layout()
plt.savefig("qini_curve.png", dpi=140)
print("Saved qini_curve.png")

# Print values at selected percentiles
print("\nQini values at selected cutoffs:")
for target_k in [10, 20, 30, 50, 70, 100]:
    idx = target_k - 1
    print(f"  Top {target_k:3d}%: observed uplift = {qini_vals[idx]:.4f}")

Expected output:

Saved qini_curve.png

Qini values at selected cutoffs:
  Top  10%: observed uplift = 0.0895
  Top  20%: observed uplift = 0.1018
  Top  30%: observed uplift = 0.0959
  Top  50%: observed uplift = 0.0966
  Top  70%: observed uplift = 0.1454
  Top 100%: observed uplift = 0.2106

Here's what's happening: you sort all 50,000 users by the T-learner's predicted CATE, highest first. For each percentile cutoff, you compute the raw treated-minus-control difference in task completion within that subgroup.

The top-10% group shows an observed uplift of +0.0895 and the top-20% group shows +0.1018, both well below the naive ATE of +0.2106, which is confounded by selection and reflects engagement level more than feature impact.

The Qini values here also mix the CATE signal with residual selection bias: all users in the top 54% by predicted CATE are light users (the tier with the lowest opt-in rate of 12%), so the treated-minus-control comparison within that group is still confounded by within-tier selection bias.

The jump in the top 70% (+0.1454) makes this confounding effect visible: as medium and heavy users enter the ranked group, the treated side suddenly includes high-completion heavy users (64.7% opt-in), while the control side remains dominated by low-completion light users. That spike is selection bias, with no genuine CATE signal behind it.

In observational uplift settings, the actionable region of the Qini is roughly the top 20% to 50%, where the ranking reflects the model's CATE estimates more cleanly than at higher percentiles, where propensity-score correlation with outcome levels dominates.

Step 4: A Segmented Rollout Rule

The CATE model assigns a predicted treatment effect to every user. Turn that into a deployment policy by setting a threshold: ship the feature to users whose predicted CATE exceeds some value, suppress it for everyone else.

# Inspect the CATE distribution first
print("CATE distribution (T-learner):")
print(pd.Series(df.cate_tlearner).describe().round(4))
print()

# Plot CATE distribution
fig, ax = plt.subplots(figsize=(8, 4))
ax.hist(df.cate_tlearner, bins=50, edgecolor="white", linewidth=0.5)
ax.axvline(0.085, color="red", linestyle="--", label="Threshold = 0.085")
ax.axvline(df.cate_tlearner.mean(), color="gray", linestyle=":",
           label=f"Mean CATE = {df.cate_tlearner.mean():.4f}")
ax.set_xlabel("Predicted CATE (T-learner)")
ax.set_ylabel("Number of users")
ax.set_title("Distribution of predicted CATEs")
ax.legend()
plt.tight_layout()
plt.savefig("cate_distribution.png", dpi=140)
print("Saved cate_distribution.png")

# Apply rollout rule
threshold = 0.085
selected = df[df.cate_tlearner >= threshold].copy()
suppressed = df[df.cate_tlearner < threshold].copy()

print(f"\nRollout threshold: CATE >= {threshold}")
print(f"Users selected for rollout: {len(selected):,} ({100*len(selected)/len(df):.0f}%)")
print(f"Users suppressed:           {len(suppressed):,} ({100*len(suppressed)/len(df):.0f}%)")
print()
print("Tier composition of selected group:")
print((selected.groupby("engagement_tier").size() / len(selected)).round(3))
print()
print(f"Mean predicted CATE (selected):   {selected.cate_tlearner.mean():.4f}")
print(f"Mean predicted CATE (suppressed): {suppressed.cate_tlearner.mean():.4f}")

Expected output:

CATE distribution (T-learner):
count    50000.0000
mean         0.0847
std          0.0126
min          0.0515
25%          0.0731
50%          0.0897
75%          0.0963
max          0.1021
Name: cate_tlearner, dtype: float64

Saved cate_distribution.png

Rollout threshold: CATE >= 0.085
Users selected for rollout: 27,203 (54%)
Users suppressed:           22,797 (46%)

Tier composition of selected group:
engagement_tier
light    1.0
dtype: float64

Mean predicted CATE (selected):   0.0955
Mean predicted CATE (suppressed): 0.0719

Here's what's happening: you inspect the full CATE distribution before setting a threshold. The mean CATE across all 50,000 users is +0.0847, with a standard deviation of +0.0126. Setting a threshold at +0.085 (just above the mean of +0.0847) selects 27,203 users (54%).

The tier composition of the selected group is 100% light users: with linear models and these features, the CATE ranges for each tier don't overlap across the threshold. Light users all have predicted CATEs between +0.0807 and +0.1021. Medium users have predicted CATEs between +0.0592 and +0.0812. The threshold at 0.085 cleanly separates the two.

The mean predicted CATE in the selected group (+0.0955) is 33% higher than in the suppressed group (+0.0719). That concentration is the value of the segmented rollout: you deploy the AI summary to the 54% of users who stand to benefit most, hold it back from medium and heavy users who show smaller predicted benefit, and collect outcome data on both groups to refine the threshold quarterly.

Figure 2: Per-tier CATE distributions from the 50,000-user synthetic dataset. The top panel shows smooth KDE curves per engagement tier: light users (blue) cluster at the highest predicted CATEs, heavy users (green) at the lowest. The bottom panel shows mean CATE per tier with 95% bootstrap confidence intervals, alongside the naive ATE (+0.2106) as a reference line. All three tier CIs sit well below the naïve ATE, confirming that the average was confounded by selection bias.

The rollout rule maps directly to a feature flag system:

# Simulate the rollout decision for a single new user
def should_show_feature(query_confidence, engagement_tier, threshold=0.085):
    """Returns True if predicted CATE exceeds the rollout threshold."""
    x = pd.get_dummies(
        pd.DataFrame([{"query_confidence": query_confidence,
                        "engagement_tier": engagement_tier}]),
        drop_first=False
    ).reindex(columns=feature_cols, fill_value=0).astype(float).values
    cate = m1.predict(x)[0] - m0.predict(x)[0]
    return cate >= threshold, round(cate, 4)

show, cate = should_show_feature(0.72, "heavy")
print(f"Heavy user, conf=0.72:  show feature={show}, CATE={cate}")

show, cate = should_show_feature(0.72, "light")
print(f"Light user, conf=0.72:  show feature={show}, CATE={cate}")

show, cate = should_show_feature(0.45, "medium")
print(f"Medium user, conf=0.45: show feature={show}, CATE={cate}")

Expected output:

Heavy user, conf=0.72:  show feature=False, CATE=0.0667
Light user, conf=0.72:  show feature=True, CATE=0.0955
Medium user, conf=0.45: show feature=False, CATE=0.0681

Here's what's happening: you wrap the CATE computation into a function that mirrors what a real feature-flag service would run at request time. A heavy user with moderate query confidence gets show feature=False and a CATE of +0.0667, below the 0.085 threshold. The same query confidence from a light user gets show feature=True and a CATE of +0.0955. A medium user with lower confidence falls below the +0.0681 threshold.

These outputs match the domain story: the AI summary helps users who struggle to maintain context across sessions, and engagement tier is a strong proxy for that struggle.

Step 5: Bootstrap Confidence Intervals

The CATE estimates above are point estimates with no uncertainty quantification. Before you build rollout rules on them, you need to know how stable those estimates are across different samples of your user base.

def bootstrap_cate_ci(df, X_all, feature_cols, n_reps=500, seed=7):
    """Bootstrap 95% CI for mean CATE overall and per engagement tier."""
    rng = np.random.default_rng(seed)
    n = len(df)
    tier_reps = {"light": [], "medium": [], "heavy": []}
    mean_reps = []

    for _ in range(n_reps):
        idx = rng.integers(0, n, size=n)
        df_b = df.iloc[idx].reset_index(drop=True)
        X_b = X_all[idx]
        treated_b = df_b.opt_in_agent_mode == 1
        m1_b = LinearRegression().fit(X_b[treated_b], df_b[treated_b].task_completed.values)
        m0_b = LinearRegression().fit(X_b[~treated_b], df_b[~treated_b].task_completed.values)
        cate_b = m1_b.predict(X_b) - m0_b.predict(X_b)
        df_b["cate"] = cate_b
        for tier in tier_reps:
            tier_reps[tier].append(df_b[df_b.engagement_tier == tier].cate.mean())
        mean_reps.append(cate_b.mean())

    cis = {}
    for tier, vals in tier_reps.items():
        arr = np.array(vals)
        cis[tier] = (float(np.percentile(arr, 2.5)),
                     float(np.percentile(arr, 97.5)))
    arr = np.array(mean_reps)
    cis["mean"] = (float(np.percentile(arr, 2.5)),
                   float(np.percentile(arr, 97.5)))
    return cis

print("Running bootstrap (500 replicates, seed=7)...")
cis = bootstrap_cate_ci(df, X_all, feature_cols, n_reps=500, seed=7)
print(f"Mean CATE   95% CI: [{cis['mean'][0]:+.4f}, {cis['mean'][1]:+.4f}]")
print(f"Light tier  95% CI: [{cis['light'][0]:+.4f}, {cis['light'][1]:+.4f}]")
print(f"Medium tier 95% CI: [{cis['medium'][0]:+.4f}, {cis['medium'][1]:+.4f}]")
print(f"Heavy tier  95% CI: [{cis['heavy'][0]:+.4f}, {cis['heavy'][1]:+.4f}]")

Expected output:

Running bootstrap (500 replicates, seed=7)...
Mean CATE   95% CI: [+0.0744, +0.0951]
Light tier  95% CI: [+0.0781, +0.1125]
Medium tier 95% CI: [+0.0596, +0.0892]
Heavy tier  95% CI: [+0.0483, +0.0842]

Here's what's happening: you resample the full 50,000-user dataset 500 times with replacement, refit the T-learner on each resample, and compute the distribution of mean CATEs across bootstrap iterations. The 2.5th and 97.5th percentiles of that distribution give a 95% confidence interval for each estimate.

Three things to check in these CIs. First, the overall mean CI (+0.0744, +0.0951) brackets the ground truth of +0.08, confirming that the estimator is working. Second, the light-tier CI (+0.0781, +0.1125) is wider than the heavy-tier CI (+0.0483, +0.0842), consistent with light users having the lowest opt-in rate (12%) and therefore fewer treated observations to anchor the estimate. Third, the tier CIs don't fully separate at their tails: light's lower bound (+0.0781) barely clears heavy's upper bound (+0.0842), meaning the ordering light > heavy is stable but not by a wide margin.

For a business decision about differential rollout, that stability is enough. For a regulatory or clinical context, you'd want larger samples.

When Uplift Modeling Fails

CATE models look compelling because they produce a continuous, individualized score. Four failure modes deserve explicit attention before you deploy a CATE-based policy.

1. Thin Segments (Overlap Violation)

The CATE for light users is estimated from 12% of your 13,451 treated users, roughly 1,614 people. That's enough to detect a tier-level average but not enough to estimate reliable individual-level effects within the tier at fine-grained feature values.

When the treatment arm has sparse coverage in a region of feature space, CATE estimates there carry high variance. The model returns a smooth prediction, but the empirical support behind it may be weak.

Check the feature distribution of your highest-CATE users and verify that treated and control observations exist in each region before acting on the ranking.

2. Extrapolation at the Tails (Overlap Violation)

Linear regression extrapolates smoothly outside the training range. If your model assigns a predicted CATE to a user whose feature values fall in a region with no training data for one arm, that estimate lacks empirical support.

The overlap assumption fails silently: the model returns a number, but P(T=1|X=x) is approximately 0 or 1 in that region, making the CATE unidentified.

Check propensity scores alongside CATE predictions and clip or flag estimates where the propensity falls outside [0.05, 0.95].

3. Qini Noise at Small k

The Qini curve is noisy at very small k (top 5% or fewer). When only a few hundred users are in the evaluation group, the treated count in that group may be small enough that the observed uplift is dominated by sampling noise.

Base rollout decisions on the 20% to 50% Qini range, where the signal is more stable. In observational settings, high Qini values at large k (such as +0.1454 in the top 70% in this tutorial) can reflect selection bias that masks the real CATE signal. Inspect the tier composition of each top-k group before interpreting the uplift value.

4. Overfitting the CATE Model

A LinearRegression trained on the treated arm here sees 13,451 observations and four features, a comfortable margin. If you replace linear regression with gradient boosting and add 30 features, you can overfit the imputed treatment effects to training noise. The CATE predictions will look sharply heterogeneous on the training set and regress toward the global mean on a held-out set. A CATE model earns its complexity when it outperforms the tier-level averages on held-out uplift. Evaluate on a held-out dataset before using it to build rollout rules.

What to Do Next

The implementations above are built without external uplift libraries so you can see exactly what each step computes. For production use, causalml and econml offer richer versions of both estimators: tree-based T-learners, doubly robust X-learners, and honest causal forests that split training and estimation samples to reduce overfitting. Both libraries follow the same conceptual structure you've built here.

causalml includes production-grade Qini curve computation and the AUUC (area under the uplift curve) metric, which collapses the Qini curve into a single comparison number. For running uplift model comparisons in an A/B framework, AUUC is the standard leaderboard metric.

One structural limitation worth naming: this tutorial assumed SUTVA, meaning each user's outcome depends only on their own treatment status. In workspace-based AI products, that assumption is often wrong. Users in the same workspace share a common environment, and treating one user can affect teammates through shared outputs, changed response patterns, or altered workspace dynamics.

When you suspect this kind of interference, DR-learner variants that propagate within-group correlation into the CATE estimates give more realistic uncertainty bounds. Standard T-learner and X-learner treat all observations as independent, which understates uncertainty when workspace-level factors are at play.

The companion repo for this tutorial lives at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/08_uplift_modeling. Clone the repo, generate the dataset with --n-users 50000 --seed 42, and run uplift_demo.py to reproduce every result in this tutorial.

The ATE is the number you need to decide whether to build a feature. The CATE is the number you need to decide who gets it first. A segmented rollout that focuses treatment on the 54% of users with the strongest predicted response yields more than spreading the same feature to everyone. Uniform rollout is a policy choice. Make it an informed one.

You wait. On day 30, the effect is still there. But you spent 16 days running a feature you already knew worked with 95% confidence, delaying the next experiment and burning opportunity cost.

The statistician is technically right, if you're running a classical fixed-sample test. The p-value in a standard t-test is valid only when you commit to a sample size in advance and look at the results exactly once. Look earlier and stop when p < 0.05, and your false positive rate climbs toward 30%.

The p-value was designed for a single pre-committed look: it was built for a static experiment with a fixed endpoint. Applying it to a live stream where you can check at any point requires a different mathematical object entirely.

Sequential testing was designed for exactly this situation. The mixture Sequential Probability Ratio Test (mSPRT) (Johari et al.) produces always-valid inference using a mathematical object called an e-value: you can check results every day, stop when the evidence is strong enough, and your false positive rate stays at 5%.

Netflix has documented the production use of always-valid sequential testing frameworks (Lindon et al.), and the underlying ideas trace back to Wald's 1945 work on sequential analysis and Ville's 1939 inequality.

This tutorial makes the connection explicit. You'll simulate the peeking problem to see the inflated error rate directly, implement a working mSPRT from scratch in Python, apply it to the shared synthetic LLM product dataset, and understand exactly when sequential testing fails.

Companion notebook: every code block in this article runs end-to-end in msprt_demo.ipynb in the companion repo.

Table of Contents

• Why Optional Stopping Breaks Classical Tests

• What a Sequential Test Actually Does

• Identification Assumptions

• Prerequisites

• Setting Up the Working Example

  • Step 1: Simulate the peeking problem

  • Step 2: Implement the mSPRT e-value

  • Step 3: Apply mSPRT to the real dataset

  • Step 4: Compare power against a fixed-sample test

  • Validate against ground truth

  • Step 5: Bootstrap confidence intervals

• When mSPRT Fails

• What to Do Next

Why Optional Stopping Breaks Classical Tests

Peeking at running p-values inflates your false positive rate toward 30%. That's the number that should give you pause, and you'll reproduce it in Step 1 below.

The p-value in a classical hypothesis test answers a specific question: given the null is true, what's the probability of seeing data this extreme when you run the experiment exactly as planned with the sample size you committed to upfront?

The "exactly as planned" clause is the problem. When you check results on day 5, day 10, day 14, and stop on day 14 because p < 0.05, you haven't run the experiment you planned. You've run 14 different experiments, looked at the results of each, and stopped at the one that passed your threshold. The p-value formula doesn't know that.

Here's the intuition. Under the null hypothesis (no effect), your p-value bounces around randomly between 0 and 1. It doesn't stay parked at 0.5. Over a 30-day run, a null experiment will dip below 0.05 at some point with high probability. If you're watching every day and ready to stop the moment you see p < 0.05, you'll almost always catch one of those dips. You'll declare a winner. But the effect isn't real.

Looking less often just delays the same problem. You need to look often: products move fast, and running an experiment 16 days longer than necessary costs real money, delays launches, and burns opportunity cost. You need a test statistic that stays valid regardless of when you stop.

What a Sequential Test Actually Does

Sequential tests are designed for optional stopping by replacing the p-value with an alternative statistic called an e-value.

Unlike a p-value, an e-value is nonnegative, and the process formed by e-values over time satisfies a supermartingale property under the null: conditional on the history, the expected next e-value is at most the current one.

This path-level supermartingale condition is what makes optional stopping safe. Having a marginal mean below 1 at each step is necessary but not sufficient: the supermartingale condition is strictly stronger, holding the bound uniformly across all stopping times.

Here's why. If the e-value process is a nonneg supermartingale with E[e_t] ≤ 1 under H0, then a classical result called Ville's inequality gives: the probability that the running maximum of the process ever exceeds 1/α is at most α. With α = 0.05 and stopping threshold 1/α = 20, the probability that a null e-value process ever reaches 20 is at most 5%.

That Type I error bound holds no matter when you stop or how many times you check. The guarantee is time-uniform: it covers all possible stopping times simultaneously.

A classical p-value's guarantee applies only at the pre-committed sample size. Check repeatedly and the bound dissolves. There is no time-uniform analog.

The mSPRT computes the e-value as a Bayes factor: the ratio of the likelihood of the observed data under the alternative to that under the null.

The "mixture" part means you don't specify a single effect size under H1. You average the likelihood ratio over a prior distribution on effect sizes.

For Bernoulli outcomes (did the task complete: yes or no), placing a Beta(1,1) prior on each arm's completion rate makes the Bayes factor tractable in closed form using the log-beta function. The math is less intimidating than it looks: the entire computation reduces to four calls to betaln, as Step 2 shows.

The practical consequence is concrete: accumulate data, compute the running e-value each day, and stop when it crosses 20. When it remains below 20 across your maximum sample size, you fail to reject the null. Check every day, every hour, or every minute. The Type I error rate holds at 5%.

Identification Assumptions

mSPRT's always-valid guarantee rests on four conditions. Each can break, and the failure modes section below maps each failure mode to the condition it violates.

1. Nonneg supermartingale property under H0. The e-value process must satisfy E[e_{t+1} | e_1, ..., e_t] ≤ e_t under H0. For the Beta-Binomial Bayes factor used here, this holds as long as the prior is proper (Beta(1,1) qualifies) and the observations are i.i.d. within each arm.

2. Stationarity. The data-generating process must be stationary across the experiment window. If the underlying completion rate shifts mid-experiment due to an unrelated change (a model update, a cohort shift from a marketing campaign, or a day-of-week effect), the e-value picks up noise that your experiment can't separate from the treatment effect.

3. Independent observations within each arm. Each user's outcome must be independent of other users'. Network effects, shared workspaces, or spillover from recommendation systems can violate this.

4. Prior specification. The Beta(1,1) prior is a modeling assumption. The mSPRT's power depends on whether the prior places reasonable mass on the true effect size. A badly misspecified prior won't break the Type I error guarantee, but it can make the e-value grow so slowly that you exhaust your sample budget without crossing the threshold.

Prerequisites

• Python 3.11+

• pandas 2.x (pip install pandas)

• numpy 1.26+ (pip install numpy)

• scipy 1.12+ (pip install scipy)

• matplotlib 3.8+ (pip install matplotlib)

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

Here's what's happening: this clones the repo that contains all 13 companion notebooks for this series, generates the shared 50,000-user synthetic dataset, and saves it to data/synthetic_llm_logs.csv. Every article in the series runs against this same CSV so the methods are directly comparable. The data generator bakes in a +5 percentage-point causal effect on task completion for wave 1 users.

Setting Up the Working Example

The synthetic dataset simulates a SaaS AI assistant product with 50,000 users. The task_completed column records whether the AI successfully completed the user's task (1) or not (0). The wave column assigns users to groups: wave 1 receives the new AI feature, wave 2 is the holdout control.

Figure 1: conceptual e-value trajectories. The blue path (real effect) rises and crosses the stopping threshold at the green dashed line. The purple path (weaker effect) grows but doesn't cross in 30 days. The grey path (null) meanders near 1 throughout. The red dashed line is the stopping boundary at 1/α = 20. Compare this to Figure 2 below, which shows the actual e-value trajectory on the real dataset.

import pandas as pd
import numpy as np

df = pd.read_csv("data/synthetic_llm_logs.csv")

treated = df[df["wave"] == 1]["task_completed"].values
control = df[df["wave"] == 2]["task_completed"].values

print(f"Treated: n={len(treated):,}, mean={treated.mean():.4f}")
print(f"Control: n={len(control):,}, mean={control.mean():.4f}")
print(f"Observed lift: {treated.mean() - control.mean():.4f}")

Expected output:

Treated: n=24,937, mean=0.6202
Control: n=25,063, mean=0.5718
Observed lift: 0.0485

Here's what's happening: you load the 50,000-row dataset and split by wave. Wave 1 has 24,937 treated users with a 62.0% task completion rate. Wave 2 has 25,063 control users with a 57.2% task completion rate. The observed 4.85 percentage-point lift is close to the ground-truth 5pp baked into the data generator, with the small gap due to sampling noise. These arrays feed the sequential test one observation at a time, as outlined in the steps below.

Step 1: Simulate the Peeking Problem

The peeking problem is real and measurable: 30 days of daily monitoring inflates your false positive rate from 4.2% to 30.2%, confirmed by the simulation below.

This simulation runs 1,000 null experiments (in which the treatment has zero effect) and checks every day whether the running p-value has dropped below 0.05. The scenario uses 60 users per arm per day across a 30-day experiment: 1,800 total observations per arm, a realistic scale for a mid-sized SaaS product.

from scipy import stats
import numpy as np

np.random.seed(42)

N_SIMS = 1000
N_DAYS = 30
USERS_PER_ARM_PER_DAY = 60
NULL_RATE = 0.60

false_positives_peeking = 0
false_positives_single_look = 0

for _ in range(N_SIMS):
    control_outcomes = []
    treated_outcomes = []
    stopped_early = False

    for day in range(N_DAYS):
        control_outcomes.extend(np.random.binomial(1, NULL_RATE, USERS_PER_ARM_PER_DAY))
        treated_outcomes.extend(np.random.binomial(1, NULL_RATE, USERS_PER_ARM_PER_DAY))

        # The peeking problem: checking the test every single day
        if len(control_outcomes) >= 10:
            _, p = stats.ttest_ind(treated_outcomes, control_outcomes)
            if p < 0.05 and not stopped_early:
                false_positives_peeking += 1
                stopped_early = True

    # The fixed-sample approach: checking only once at the very end
    _, p_final = stats.ttest_ind(treated_outcomes, control_outcomes)
    if p_final < 0.05:
        false_positives_single_look += 1

print(f"False positive rate (peeking daily):  {false_positives_peeking / N_SIMS:.1%}")
print(f"False positive rate (single look):    {false_positives_single_look / N_SIMS:.1%}")

Expected output:

False positive rate (peeking daily):  30.2%
False positive rate (single look):    4.2%

Here's what's happening: each simulation generates null data, with both arms drawn from the same 60% completion rate, so any detected effect is pure noise. The inner loop adds 60 observations per arm per day and runs a t-test on the accumulated data for that day.

When the p-value falls below 0.05 for the first time, the simulation flags a false positive and stops (mimicking a team that ships when it detects significance).

The single-look check at day 30 is the honest fixed-sample test. One look gives 4.2% false positives, close to nominal. Daily peeking reaches 30.2%, meaning more than one in four "significant" experiments is detecting noise.

Step 2: Implement the mSPRT e-value

The mSPRT computes a Bayes factor at each time step: how much more likely are the observed data under a mixture of alternatives than under the null? For binary outcomes with a Beta(1,1) prior on each arm's completion rate, the running Bayes factor has a closed form using the log-beta function.

from scipy.special import betaln

def compute_evalue_running(outcomes_treated, outcomes_control,
                           alpha_prior=1.0, beta_prior=1.0):
    """
    Compute the running mSPRT e-value for two Bernoulli arms.

    Parameters
    ----------
    outcomes_treated : array-like of 0/1
    outcomes_control : array-like of 0/1
    alpha_prior, beta_prior : Beta prior hyperparameters (default: uniform)

    Returns
    -------
    e_values : np.ndarray of shape (n,), one e-value per observation
    """
    outcomes_treated = np.asarray(outcomes_treated, dtype=float)
    outcomes_control = np.asarray(outcomes_control, dtype=float)
    n = min(len(outcomes_treated), len(outcomes_control))

    cum_t = np.cumsum(outcomes_treated[:n])
    cum_c = np.cumsum(outcomes_control[:n])
    t_arr = np.arange(1, n + 1, dtype=float)

    # Alternative hypothesis: each arm has its own independent Beta prior on completion rate
    log_ml_t = (betaln(alpha_prior + cum_t, beta_prior + t_arr - cum_t)
                - betaln(alpha_prior, beta_prior))
    log_ml_c = (betaln(alpha_prior + cum_c, beta_prior + t_arr - cum_c)
                - betaln(alpha_prior, beta_prior))

    # Null hypothesis: both arms share a single pooled Beta prior on the common rate
    pooled_successes = cum_t + cum_c
    pooled_n = 2 * t_arr
    log_ml_h0 = (betaln(alpha_prior + pooled_successes,
                        beta_prior + pooled_n - pooled_successes)
                 - betaln(alpha_prior, beta_prior))

    # Log Bayes factor is the difference in log marginal likelihoods
    log_bf = log_ml_t + log_ml_c - log_ml_h0

    return np.exp(log_bf)

Here's what's happening: the function takes two arrays of 0/1 outcomes arriving in temporal order. For each time step t, it computes the cumulative number of successes and trials for each arm.

betaln gives the log of the beta function, which is the normalizing constant for the Beta-Binomial marginal likelihood. H1 integrates over independent Beta priors on each arm's rate;.H0 integrates over a single shared-rate prior.

The log Bayes factor is the difference. Exponentiating gives the e-value. When the treatment has a real effect, the e-value grows over time. With no effect, it bounces near 1 and is a non-negative supermartingale under H0.

A quick sanity check on null data confirms the expected behavior:

np.random.seed(0)
null_t = np.random.binomial(1, 0.60, 500)
null_c = np.random.binomial(1, 0.60, 500)
ev_null = compute_evalue_running(null_t, null_c)
print(f"E-value at end under null (should be near 1): {ev_null[-1]:.3f}")
print(f"Max e-value under null: {ev_null.max():.3f}")

Expected output:

E-value at end under null (should be near 1): 0.078
Max e-value under null: 2.188

Here's what's happening: under the null, the final e-value ends near 1 (0.078 here, due to sampling variation), and the maximum over 500 observations stays well below the stopping threshold of 20. By Ville's inequality, the probability that a valid null e-value process ever reaches 20 is at most 5%, consistent with a 5% Type I error rate. In this single 500-observation run, the max is 2.188, which is expected behavior.

Step 3: Apply mSPRT to the Real Dataset

Now apply the test to the synthetic data where a real treatment effect exists. You'll compute the running e-value day by day and find the first day it crosses the stopping threshold.

import matplotlib.pyplot as plt

np.random.seed(42)
treated_shuffled = treated.copy()
control_shuffled = control.copy()
np.random.shuffle(treated_shuffled)
np.random.shuffle(control_shuffled)

USERS_PER_ARM_PER_DAY = 60
N_DAYS_RUN = 30
n_per_arm = USERS_PER_ARM_PER_DAY * N_DAYS_RUN  # 1,800

treated_seq = treated_shuffled[:n_per_arm]
control_seq = control_shuffled[:n_per_arm]

e_values = compute_evalue_running(treated_seq, control_seq)

ALPHA = 0.05
THRESHOLD = 1 / ALPHA  # = 20

days = np.arange(1, len(e_values) + 1) / USERS_PER_ARM_PER_DAY
cross_indices = np.where(e_values >= THRESHOLD)[0]
if len(cross_indices) > 0:
    stopping_day = days[cross_indices[0]]
    print(f"mSPRT stopping day: {stopping_day:.1f}")
    print(f"E-value at stopping: {e_values[cross_indices[0]]:.1f}")
else:
    stopping_day = None
    print("mSPRT did not cross threshold in this window")

print(f"Final e-value on day {N_DAYS_RUN}: {e_values[-1]:.2f}")

Expected output:

mSPRT stopping day: 25.9
E-value at stopping: 20.9
Final e-value on day 30: 75.64

Here's what's happening: you shuffle the treatment and control arrays to simulate random daily arrival of users (real experiments don't deliver users in any particular order), then feed the first 1,800 per arm into compute_evalue_running one observation at a time. The e-value crosses the threshold of 20 on day 25.9, meaning you could have called the experiment 4 days early with a fully valid inference guarantee. By day 30, the e-value has climbed to 75.64, far above the threshold.

Figure 2: actual mSPRT e-value trajectory on the real 50,000-user synthetic dataset (wave 1 treatment vs. wave 2 control). The blue line is the running e-value on a log scale. The red dashed line is the stopping threshold at 1/α = 20.

The dotted green vertical line marks day 25.9, when the e-value first crosses the threshold. The bottom panel shows cumulative task completion rates per arm converging as data accumulates. Unlike the schematic in Figure 1, these are real data from the shared dataset, with a true 4.85 pp lift.

Step 4: Compare Power Against a Fixed-Sample Test

The mSPRT carries a real cost. When the effect is active, it lets you stop earlier than the scheduled end time. When the effect is smaller than your prior expects, or when you're working with modest sample sizes, the power penalty is substantial. This simulation quantifies the trade-off honestly.

from scipy.stats import ttest_ind

np.random.seed(42)

N_SIMS = 1000
TRUE_EFFECT = 0.05
BASE_RATE = 0.60
N_PER_ARM = 1800          # 30 days x 60 users/arm/day
DAILY_BATCH = 60
THRESHOLD = 20

msprt_stopping_days = []
msprt_detected = 0
ttest_detected = 0

for sim in range(N_SIMS):
    t_obs = np.random.binomial(1, BASE_RATE + TRUE_EFFECT, N_PER_ARM)
    c_obs = np.random.binomial(1, BASE_RATE, N_PER_ARM)

    e_vals = compute_evalue_running(t_obs, c_obs)
    days = np.arange(1, N_PER_ARM + 1) / DAILY_BATCH
    crosses = np.where(e_vals >= THRESHOLD)[0]
    if len(crosses) > 0:
        msprt_detected += 1
        msprt_stopping_days.append(days[crosses[0]])
    else:
        msprt_stopping_days.append(30.0)

    _, p = ttest_ind(t_obs, c_obs)
    if p < 0.05:
        ttest_detected += 1

msprt_power = msprt_detected / N_SIMS
ttest_power = ttest_detected / N_SIMS
median_stop = np.median(msprt_stopping_days)
pct_stopped_early = np.mean(np.array(msprt_stopping_days) < 30.0)

print(f"mSPRT power:               {msprt_power:.1%}")
print(f"Fixed-sample t-test power: {ttest_power:.1%}")
print(f"Median mSPRT stop day:     {median_stop:.1f} / 30")
print(f"Fraction stopping early:   {pct_stopped_early:.1%}")

Expected output:

mSPRT power:               49.3%
Fixed-sample t-test power: 88.7%
Median mSPRT stop day:     30.0 / 30
Fraction stopping early:   49.3%

Here's what's happening: you run 1,000 simulations with a true 5pp lift. For mSPRT, the running e-value is computed, and the first crossing of 20 is recorded.

For the fixed-sample test, you look once at the end of day 30. The results show a meaningful power gap: mSPRT detects the effect in 49.3% of experiments, whereas the fixed-sample test detects it in 88.7%. With a 5pp lift and 1,800 observations per arm, the mSPRT requires roughly twice as many observations to match the fixed-sample test's power.

That's the price of the always-valid guarantee. What you gain is the Type I error control when you check daily: a fixed-sample test peeked at daily inflates to 30.2% false positives. mSPRT stays at 5% regardless of when you stop.

The right choice depends on which is more expensive for your team: running experiments longer, or shipping false positives. Most teams underestimate the cost of power until they run this simulation themselves.

Validate Against Ground Truth

The synthetic dataset incorporates a known 5pp lift, so you can check whether mSPRT correctly identifies the effect when given more data beyond the 30-day window.

np.random.seed(0)
t_full = treated_shuffled
c_full = control_shuffled[:len(t_full)]

e_full = compute_evalue_running(t_full, c_full)
days_full = np.arange(1, len(e_full) + 1) / USERS_PER_ARM_PER_DAY

cross_full = np.where(e_full >= THRESHOLD)[0]
if len(cross_full) > 0:
    print(f"mSPRT correctly detected the effect.")
    print(f"Could have stopped on day {days_full[cross_full[0]]:.1f}")
    print(f"True effect in data: {treated.mean() - control.mean():.4f}")
    print(f"E-value at stopping point: {e_full[cross_full[0]]:.1f}")
else:
    print("mSPRT did not cross threshold with this data slice.")

Expected output:

mSPRT correctly detected the effect.
Could have stopped on day 27.1
True effect in data: 0.0485
E-value at stopping point: 22.2

Here's what's happening: running mSPRT on the full shuffled arrays (24,937 treated, 25,063 control), the e-value crosses the threshold at day 27.1. The true causal effect in the data, 4.85 pp, is close to the generator's ground truth of 5 pp and is correctly detected.

A fixed-sample test designed for 30 days holds you to day 30 even when the evidence has already accumulated. With 60 users per arm per day, mSPRT would have let you ship on day 27.1, saving almost 3 days on a feature that was always going to ship.

Step 5: Bootstrap Confidence Intervals

A stopping day tells you when to call the experiment, but it doesn't tell you how large the effect is or how precisely it's estimated. Bootstrap confidence intervals give you both.

rng = np.random.default_rng(7)
point_est = treated.mean() - control.mean()

boot_diffs = np.array([
    rng.choice(treated, size=len(treated), replace=True).mean() -
    rng.choice(control, size=len(control), replace=True).mean()
    for _ in range(500)
])

lower = float(np.percentile(boot_diffs, 2.5))
upper = float(np.percentile(boot_diffs, 97.5))

print(f"Point estimate (treated - control): {point_est:.4f} ({point_est*100:.2f}pp)")
print(f"95% bootstrap CI: [{lower:.4f}, {upper:.4f}]  "
      f"([{lower*100:.2f}pp, {upper*100:.2f}pp])")
print(f"Ground-truth 5pp is {'inside' if lower <= 0.05 <= upper else 'outside'} the CI.")

Expected output:

Point estimate (treated - control): 0.0485 (4.85pp)
95% bootstrap CI: [0.0407, 0.0581]  ([4.07pp, 5.81pp])
Ground-truth 5pp is inside the CI.

Here's what's happening: you resample the treated and control arrays independently with replacement 500 times, computing the difference in means each time. The 2.5th and 97.5th percentiles of the 500 differences form the confidence interval. The CI runs from 4.07pp to 5.81pp, covering the ground-truth 5pp and excluding zero, confirming the effect is real. The interval is reasonably tight given 25k users per arm, giving you both the "did it work" answer (yes) and the "how much" answer (between 4.07 and 5.81 percentage points) in a single step.

When mSPRT Fails

Sequential tests still demand experimental rigor. Four situations either break the guarantee or make the method practically useless.

Badly Misspecified Prior

The mSPRT assumes a Beta(1,1) prior on each arm's completion rate, a modeling choice with real consequences. This violates the prior specification assumption when your true effect is far outside the range the prior expects.

A uniform Beta(1,1) prior performs reasonably well for moderate effects in the 3–10 pp range at base rates around 60%. If your true effect is a 0.3pp lift, a realistic outcome for a marginal AI feature change, the e-value grows extremely slowly. You'll exhaust your sample budget before crossing the threshold.

Calibrate the prior against historical A/B test data from your product: fit Beta hyperparameters to the distribution of past effect sizes using maximum likelihood, and verify that the resulting prior puts meaningful mass near your minimum detectable effect.

Non-Stationary Outcomes

The guarantee requires the e-value process to be a non-negative supermartingale under the null, which requires the data-generating process to be stationary. If your AI model updates mid-experiment, if the user population shifts (a marketing campaign brings in a different cohort on day 12), or if there's a day-of-week effect in task difficulty, the e-value absorbs environment noise that your experiment can't separate from the treatment effect.

Diagnose non-stationarity by running your e-value implementation on holdout A/A experiments: if the null e-value process trends upward when it should stay near 1, your environment isn't stationary enough for the method to be reliable.

Multiple Metrics Without Multiplicity Correction

mSPRT controls Type I error for a single comparison. The method itself doesn't fail when you test 20 metrics, so each individual e-value remains valid. What fails is your familywise error rate: running mSPRT on 20 metrics simultaneously and stopping when any one crosses 20 inflates the probability of at least one false positive well above 5%.

Apply a Bonferroni correction by raising the threshold to 1/(α/m) = 400 for m=20 metrics at α=0.05, or use a Benjamini-Hochberg procedure on the final e-values when the experiment ends.

The multiplicity problem is identical to the one you'd face with fixed-sample tests. mSPRT doesn't make it worse, and it doesn't solve it either. This is a common misconception worth naming explicitly.

Minimum Runtime is Still Real

Because the always-valid guarantee applies regardless of when you check, it's tempting to start monitoring immediately. Don't. The guarantee holds whenever you check, but low power means the test rarely rejects even when the effect is real.

The Step 4 simulation shows this directly: with 1,800 observations per arm and a 5 pp lift, mSPRT has only 49.3% power. Before starting an mSPRT-monitored experiment, compute the minimum sample size for 80% power at your expected effect size using a standard power calculator, and set that as your floor before you start monitoring. Don't check the e-value until you've reached that floor.

What to Do Next

Apply mSPRT to your primary metric, with a minimum runtime floor set to the sample size required for 80% power at your expected effect size.

Run A/A tests on historical holdout data first: the calibration check costs you nothing and catches non-stationary environments before they corrupt a real experiment. Teams that skip the A/A test discover calibration failures during live experiments. That's an expensive way to learn about non-stationary data.

For the full implementation including bootstrap confidence intervals, see 07_sequential_msprt/ in the companion repo.

You're ready to test it in production. Then your platform engineer raises a structural problem: you can't randomize at the user level.

This issue is rooted in causal inference and runs deeper than a technical constraint. Every user draws from a centralized pool of premium model capacity. A standard A/B test creates an uneven playing field in this environment. When the routing AI is active for the treatment group, those users consume premium resources first, leaving the control group with degraded availability.

The routing AI does more than alter the treatment group's experience. It fundamentally shifts the resource environment for everyone else. You're not isolating the AI's impact. You're measuring the combined effect of the routing AI and the artificial scarcity your experimental design imposed on the control group. That's a confounded measurement, not a clean experiment.

Switchback experiments are the standard fix for LLM-based platforms and for any shared-resource product where user-level randomization would break the comparison. You stop randomizing users and randomize time slots instead.

The full platform runs with AI routing on for a 30-minute slot, then off for the next 30 minutes. You repeat the cycle, accumulate enough slots, and estimate the average treatment effect from the contrast between AI-on and AI-off slots.

This tutorial walks through the full switchback pipeline in Python: building the time series from session logs, diagnosing carryover contamination, estimating the direct effect with and without carryover adjustment, applying HAC standard errors for time-series data, computing bootstrap confidence intervals, and validating all estimates against a known ground truth.

By the end, you'll know how to run this analysis on your own LLM platform data and how to spot the four conditions that break it.

Table of Contents

• Why User-Level A/B Testing Fails on Shared LLM Infrastructure

• How Switchback Design Restores a Clean Comparison

  • Identification Assumptions

  • Prerequisites

  • Step 1: Build the Switchback Time Series

  • Step 2: Naïve Estimate (Ignoring Time Structure)

  • Step 3: Carryover-Adjusted OLS Regression

  • Step 4: HAC Standard Errors for Time-series Data

  • Step 5: Bootstrap Confidence Intervals

• Validating Against the Ground Truth

• When Switchback Fails

• When to Use Switchback vs. Cluster Randomization

• What to Do Next

Why User-Level A/B Testing Fails on Shared LLM Infrastructure

Standard A/B testing buys you causal inference through randomization. When you flip a coin to assign each user to treatment or control, both groups share identical distributions of every confounder on average. Differences in outcomes trace back to the treatment. The logic holds when users act independently of each other.

Shared LLM infrastructure breaks that independence. Consider the query-routing scenario. If 50% of users are assigned to AI routing, they receive priority access to the premium model, enabling them to complete tasks faster and at higher rates. The remaining 50% operate in a degraded environment, where premium-model queues are longer because treatment-group sessions occupy capacity. Control-group users experience worse availability not because the AI routing feature fails them, but because your experiment design created artificial scarcity for them.

Interference is the structural problem here: the Stable Unit Treatment Value Assumption, known as SUTVA, holds that a unit's outcome depends solely on that unit's treatment assignment.

SUTVA fails on shared LLM infrastructure. A treated user's session claims capacity that determines whether a control user gets routed to the premium model or the degraded standard model. The control group is no longer a clean counterfactual.

The estimated treatment effect under user-level randomization is:

Naive ATE = E[outcome | AI-on user] - E[outcome | AI-off user, degraded capacity]

The counterfactual you actually need is what AI-off users would have experienced if no users had AI routing, with no capacity degradation. You never observe that counterfactual in a 50/50 user-level split. Your estimate conflates the routing AI's direct effect with the capacity-degradation penalty, and separating them requires knowing the full capacity-utilization function, which you almost never have.

Other shared-resource LLM platform patterns produce the same failure: a caching layer that speeds retrieval for treated users but drains shared cache space for control users, and a fine-tuned model version that consumes GPU memory, leaving standard inference slower for the control group, or a batch-processing scheduler that prioritizes AI-routed requests and creates queuing delays for everything else. Anything touching a shared resource pool contaminates the control group.

How Switchback Design Restores a Clean Comparison

Because standard randomization can poison the control group through shared resources, a switchback design changes what you randomize. You stop randomizing users. You randomize time slots.

The entire platform operates under a single treatment condition at any given time: AI routing is either on or off for all users.

The treatment indicator switches between slots on a predetermined schedule, cycling through alternating blocks across the experiment. At the end of the run, you have a time series of slots, each with a treatment indicator and an aggregate outcome, such as the mean task completion rate or the mean cost per session. You regress the outcome on the treatment indicator, and the coefficient is your average treatment effect estimate.

Figure 1: Conceptual schematic of the 3-slot switchback design. Blue regions are AI-routing-on blocks, while orange marks the first AI-off slot of each cycle where carryover from the prior on-block artificially elevates outcomes.
The green band shows the true 6 pp direct effect. A naïve comparison of all-on vs. all-off slots inflates the estimated effect because it can't disentangle the direct contribution from within-block carryover.

A clean comparison is restored because the platform operates under a single condition for any given slot. Every user within a slot sees the same treatment. The AI-off slots function as a reliable counterfactual for the AI-on slots, provided that demand conditions remain comparable across slots.

The key complication is carryover. If AI routing effects persist into a subsequent AI-off slot due to factors such as warm routing caches, in-flight sessions that began under AI routing and complete after the switch, or changed user behavior that persists across the slot boundary, then AI-off slot outcomes are artificially elevated by residual AI effects.

The naïve comparison conflates this inherited elevation with the direct treatment effect, biasing the estimate upward. Estimating and removing carryover is the core analytical challenge in switchback experiments: it's where most of the real work lives, and most of what this tutorial covers.

Identification Assumptions

Switchback estimates have a causal interpretation only when four conditions hold.

1. Zero or bounded carryover between slots.

AI routing effects from one slot don't persist far enough into later slots to bias the comparison. The carryover model in this tutorial captures first-order persistence (one lag). If effects persist for multiple periods, you need more lag terms in the regression.

2. Demand stationarity across the treatment schedule.

AI-on and AI-off slots face similar underlying demand conditions. If Monday morning slots are always AI-on and Sunday afternoon slots are always AI-off, demand differences contaminate the treatment comparison in ways no lag correction can fix.

3. No ramp-up effects at block boundaries.

The system reaches steady-state behavior within each slot. If the first slot of each AI-on block performs worse than subsequent slots because the routing model's cache is cold, that ramp-up period produces a downward-biased estimate of the steady-state direct effect.

4. Residual autocorrelation is addressed.

Slot residuals may be correlated over time due to demand cycles, capacity events, and platform-level shocks spanning multiple periods. HAC standard errors or bootstrap CIs correct for this (as plain OLS standard errors aren't sufficient).

The "When switchback fails" section maps each failure mode to the specific assumption it violates.

All code in this tutorial runs end-to-end in the companion notebook at 06_switchback/switchback_demo.ipynb.

Prerequisites

• Python 3.11+

• pandas 2.x (pip install pandas)

• numpy 1.26+ (pip install numpy)

• statsmodels 0.14+ (pip install statsmodels)

• matplotlib 3.8+ (pip install matplotlib)

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py

The generate script writes data/synthetic_llm_logs.csv, a 50,000-row file of synthetic SaaS LLM product telemetry. Key columns are user_id, task_completed (binary outcome), cost_usd, and session_minutes.

After slot assignment in Step 1, each of the 48 time slots contains approximately 1,042 sessions. The dataset represents realistic LLM platform traffic: query arrival rates, model cost distributions, and session lengths are drawn from distributions calibrated to production patterns.

Step 1: Build the Switchback Time Series

Switchback experiments are run with a live treatment-assignment controller that flips the routing AI on or off at the slot boundary in production.

For this tutorial, you construct the time series from the session log by mapping each row to a synthetic hour slot, then aggregating to the slot level.

import pandas as pd
import numpy as np

df = pd.read_csv("data/synthetic_llm_logs.csv")
print(f"Dataset shape: {df.shape}")
print(df[["user_id", "task_completed", "cost_usd", "session_minutes"]].head(3).round(3))

# Shuffle to eliminate row-ordering bias before slot assignment
df = df.sample(frac=1, random_state=42).reset_index(drop=True)

# Assign hour slots: 48 slots, each containing ~1,042 sessions
df['hour_slot'] = df.index % 48

# Treatment schedule: 3-slot blocks (on, on, on, off, off, off, ...)
# 3-slot blocks give the platform time to settle into each state and break
# the perfect collinearity between ai_on and its one-period lag.
ai_on_schedule = np.tile([1, 1, 1, 0, 0, 0], 8)   # 48 slots, 8 full cycles
df['ai_on'] = ai_on_schedule[df['hour_slot']]

# Aggregate to slot level: mean outcome, mean cost, treatment indicator, session count
slots = df.groupby('hour_slot').agg(
    mean_task_completed = ('task_completed', 'mean'),
    mean_cost           = ('cost_usd',       'mean'),
    ai_on               = ('ai_on',          'first'),
    n_obs               = ('user_id',         'count')
).reset_index()

print(f"\nSlot-level data: {len(slots)} slots")
print(slots[['hour_slot', 'ai_on', 'mean_task_completed', 'mean_cost', 'n_obs']].head(8).round(4))
print(f"\nAI-on slots: {slots['ai_on'].sum()},  AI-off slots: {(1 - slots['ai_on']).sum()}")

Expected output:

Dataset shape: (50000, 16)
   user_id  task_completed  cost_usd  session_minutes
0        0               0     0.022             7.03
1        1               1     0.008             4.07
2        2               1     0.040             8.34

Slot-level data: 48 slots
   hour_slot  ai_on  mean_task_completed  mean_cost  n_obs
0          0      1               0.5950     0.0222   1042
1          1      1               0.5806     0.0223   1042
2          2      1               0.5950     0.0224   1042
3          3      0               0.6353     0.0218   1042
4          4      0               0.6017     0.0222   1042
5          5      0               0.6094     0.0218   1042
6          6      1               0.5912     0.0218   1042
7          7      1               0.5931     0.0219   1042

AI-on slots: 24,  AI-off slots: 24

The process begins by shuffling the dataset before slot assignment to eliminate any row-ordering artifacts from data generation. Each of the 50,000 rows is assigned to one of 48 synthetic hour slots using modulo arithmetic, and the treatment schedule alternates in 3-slot blocks, completing eight full cycles.

The 3-slot block structure serves two purposes: it gives the platform time to settle into each treatment state, and it breaks the perfect collinearity between the current treatment indicator and its one-period lag, which would otherwise make carryover estimation impossible under a purely alternating schedule. After aggregation, each slot contains approximately 1,042 sessions.

Notice that before injection, the slot-level means don't yet separate clearly by treatment. Slots 3, 4, and 5 (AI-off) show slightly higher completion rates than slots 0, 1, and 2 (AI-on) in the raw data. That's expected: before injection, the treatment assignment is arbitrary, and outcomes carry no true signal. The injection step below bakes in the ground truth.

# Known ground truth baked into the simulation
TRUE_EFFECT = 0.060   # AI routing raises task completion by 6 percentage points
CARRYOVER   = 0.030   # Residual routing effect persists into the following slot

# Replace slot means with synthetic balanced base rates.
# Slot noise std matches the CLT variance of aggregating ~1,042 Bernoulli sessions,
# simulating realistic slot-to-slot demand variation without treatment-group imbalance.
BASE_RATE = df['task_completed'].mean()
slot_noise_std = np.sqrt(BASE_RATE * (1 - BASE_RATE) / slots['n_obs'].iloc[0])
rng = np.random.default_rng(42)
slots['mean_task_completed'] = BASE_RATE + rng.normal(0, slot_noise_std, size=len(slots))

# Lag the treatment indicator: did the previous slot have AI routing on?
slots['ai_on_lag1'] = slots['ai_on'].shift(1).fillna(0).astype(int)

# Observed outcome = base outcome + treatment effect + carryover from prior slot
slots['mean_task_completed'] = (
    slots['mean_task_completed']
    + TRUE_EFFECT * slots['ai_on']
    + CARRYOVER   * slots['ai_on_lag1']
)

print("Post-injection slot data:")
print(slots[['hour_slot', 'ai_on', 'ai_on_lag1', 'mean_task_completed']].head(8).round(4))

Expected output:

Post-injection slot data:
   hour_slot  ai_on  ai_on_lag1  mean_task_completed
0          0      1           0               0.6606
1          1      1           1               0.6701
2          2      1           1               0.6973
3          3      0           1               0.6402
4          4      0           0               0.5663
5          5      0           0               0.5761
6          6      1           0               0.6579
7          7      1           1               0.6811

The injection substitutes raw slot means with noise calibrated to the variance of 1,042 Bernoulli trials, producing slot-to-slot fluctuation that mirrors production demand variability without artificial treatment-group imbalance.

The lag of ai_on identifies which slots immediately follow an AI-on period. The injection formula then adds TRUE_EFFECT (0.060) to every AI-on slot and CARRYOVER (0.030) to every slot that follows an AI-on slot, regardless of its own treatment status.

Look at slot 3: ai_on=0 but ai_on_lag1=1, so its outcome receives the +0.030 carryover boost even though AI routing is off. That's the carryover contamination a naïve model can't see.

The first AI-off slot of each cycle reflects a genuine off period, but its outcome is elevated by residual routing state from the previous block. A naïve comparison of all AI-on vs. all AI-off slots treats that elevated outcome as part of the AI-off baseline, distorting the true direct effect.

Figure 2: Left: the 48-slot time series from the synthetic dataset after injecting a 6 pp treatment effect and 3 pp carryover. Orange dots mark the first AI-off slot of each cycle (ai_on=0, ai_on_lag1=1), where outcomes remain elevated from the prior AI-on block.
Right: naïve OLS (red) overshoots the true 6 pp effect by 0.9 pp because it conflates direct and inherited carryover. The carryover-adjusted OLS (blue) recovers the true effect. Both 95% bootstrap CIs include the green dashed true-effect line.

Step 2: Naive Estimate (Ignoring Time Structure)

Before adding any sophistication, compute the obvious estimate: regress mean task completion on the binary AI-on indicator, ignoring the time structure entirely.

import statsmodels.api as sm

# Naive OLS: outcome ~ constant + ai_on
# No lag term, no time controls
X_naive = sm.add_constant(slots['ai_on'])
naive_model = sm.OLS(slots['mean_task_completed'], X_naive).fit()

naive_ate = naive_model.params['ai_on']
naive_se  = naive_model.bse['ai_on']

print("=== Naive estimate (no carryover control) ===")
print(f"  ATE estimate : {naive_ate:.4f}")
print(f"  Std error    : {naive_se:.4f}")
print(f"  95% CI       : [{naive_ate - 1.96*naive_se:.4f},  {naive_ate + 1.96*naive_se:.4f}]")
print(f"\n  True effect  : {TRUE_EFFECT}")
print(f"  Bias         : {naive_ate - TRUE_EFFECT:+.4f}")

Expected output:

=== Naive estimate (no carryover control) ===
  ATE estimate : 0.0688
  Std error    : 0.0048
  95% CI       : [0.0595,  0.0782]

  True effect  : 0.06
  Bias         : +0.0088

The naïve OLS regresses mean task completion on the binary AI-on indicator alone, treating the 48 slots as 48 independent observations with no time structure. It returns an ATE of 0.0688 against a true direct effect of 0.060, a bias of +0.0088, nearly a full percentage point of artificial lift.

The bias stems from how carryover is distributed between the two groups. In a 3-slot-on / 3-slot-off design, slots 1 and 2 of every AI-on block receive both the direct treatment effect (+0.060) and the carryover effect (+0.030) from the previous on-slot, pushing their outcomes to base + 0.090.

The naïve model can't separate these two contributions: it sees a high outcome in an AI-on slot and attributes it entirely to the direct treatment. Across 24 AI-on slots, 16 receive this compound injection, pulling the group average well above the true direct effect.

On the AI-off side, the first off-slot of each block receives +0.030 carryover, which raises the AI-off group's baseline. That partially offsets the AI-on group inflation, but 16 slots of compound AI-on inflation outweigh 8 slots of AI-off carryover. The net result is a positive bias of roughly +0.009 percentage points.

A team acting on 0.0688, when the true effect is 0.060, will declare a larger effect than exists and over-prioritize the routing feature relative to other initiatives.

Step 3: Carryover-Adjusted OLS Regression

The fix is to add the lagged treatment indicator to the regression. The coefficient on ai_on then measures the direct effect of the current period's treatment, holding the prior period's treatment constant. That's the quantity you want.

# Carryover-adjusted OLS: outcome ~ constant + ai_on + ai_on_lag1
X_adj = sm.add_constant(slots[['ai_on', 'ai_on_lag1']])
adj_model = sm.OLS(slots['mean_task_completed'], X_adj).fit()

adj_ate      = adj_model.params['ai_on']
adj_carryover = adj_model.params['ai_on_lag1']
adj_se        = adj_model.bse['ai_on']

print("=== Carryover-adjusted estimate ===")
print(adj_model.summary().tables[1])

print(f"\n  Direct ATE estimate  : {adj_ate:.4f}  (true: {TRUE_EFFECT})")
print(f"  Carryover estimate   : {adj_carryover:.4f}  (true: {CARRYOVER})")
print(f"  Residual bias        : {adj_ate - TRUE_EFFECT:+.4f}")

# How much did we remove?
removed = naive_ate - adj_ate
print(f"\n  Bias removed vs naive: {removed:.4f}")

Expected output:

=== Carryover-adjusted estimate ===
==============================================================================
                 coef    std err          t      P>|t|      [0.025      0.975]
------------------------------------------------------------------------------
const          0.5996      0.003    222.975      0.000       0.594       0.605
ai_on          0.0607      0.004     16.830      0.000       0.053       0.068
ai_on_lag1     0.0244      0.004      6.754      0.000       0.017       0.032
==============================================================================

  Direct ATE estimate  : 0.0607  (true: 0.06)
  Carryover estimate   : 0.0244  (true: 0.03)
  Residual bias        : +0.0007

  Bias removed vs naive: 0.0081

The adjusted regression includes both ai_on (current slot treatment) and ai_on_lag1 (previous slot treatment) as regressors.

The model now decomposes the drivers of elevated outcomes in each slot: some elevation comes from the current period's AI routing, and some from the previous period's residual. The coefficient on ai_on isolates only the current-period direct effect.

The direct ATE estimate drops from 0.0688 to 0.0607, recovering the true value of 0.060 to within 0.0007, with a residual bias smaller than the standard error.

The carryover estimate is 0.0244, compared with a true carryover of 0.030. Some underestimation is expected: the 3-slot block structure creates slots where both ai_on and ai_on_lag1 equal 1, introducing mild collinearity that slightly attenuates the carryover coefficient. Adding ai_on_lag1 removed 0.0081 of the 0.0088 naïve bias, recovering roughly 92% of the upward distortion.

The two-coefficient interpretation matters for product decisions. The ai_on coefficient (0.0607) is the direct effect: what AI routing adds in the current slot, independent of what happened in the prior slot. The ai_on_lag1 coefficient (0.0244) is the carryover effect: the residual impact that persists into the next slot after routing is switched off. In a real LLM platform, carryover might reflect session-level state, warm inference caches, or shifts in user behavior that span the slot boundary.

If ai_on_lag2 and ai_on_lag3 still improve model fit as measured by decreasing AIC, your slot length is shorter than the system's memory, and you need more lag terms. Add lags until AIC stops improving, and use domain knowledge to set a ceiling on plausible persistence given your platform's architecture.

Step 4: HAC Standard Errors for Time-series Data

The adjusted OLS model gives you the right point estimate. But the standard errors it reports assume residuals are uncorrelated across time.

Slot residuals inherit any systematic variation not captured by the treatment indicators: demand cycles, capacity events, model-version deployments, and user behavior patterns that span multiple periods. That autocorrelation makes OLS standard errors too small, which inflates your t-statistics and makes the effect look more precisely measured than it is.

The correction is Heteroskedasticity- and Autocorrelation-Consistent (HAC) standard errors, also called Newey-West standard errors. They correct for serial correlation in residuals using a bandwidth parameter equal to the number of lags you expect to matter.

from statsmodels.stats.sandwich_covariance import cov_hac
from statsmodels.stats.stattools import durbin_watson

# First check for autocorrelation in the residuals
dw_stat = durbin_watson(adj_model.resid)
print(f"Durbin-Watson statistic: {dw_stat:.4f}")
print("  DW near 2.0 = little autocorrelation in residuals.")
print("  DW < 1.5 = positive serial correlation.")
print("  DW > 2.5 = negative serial correlation.")
print("  Apply HAC standard errors regardless -- DW only tests AR(1) structure.")

# Apply HAC correction (Newey-West), 3 lags
hac_cov = cov_hac(adj_model, nlags=3)
hac_se  = np.sqrt(np.diag(hac_cov))

print("\n=== Standard error comparison ===")
print(f"  OLS SE on ai_on  : {adj_model.bse['ai_on']:.4f}")
print(f"  HAC SE on ai_on  : {hac_se[1]:.4f}")
print(f"  OLS t-stat       : {adj_model.tvalues['ai_on']:.2f}")
print(f"  HAC t-stat       : {adj_ate / hac_se[1]:.2f}")

# Construct HAC-based confidence interval manually
hac_ci_lower = adj_ate - 1.96 * hac_se[1]
hac_ci_upper = adj_ate + 1.96 * hac_se[1]
print(f"\n  HAC 95% CI: [{hac_ci_lower:.4f},  {hac_ci_upper:.4f}]")
print(f"  True effect {TRUE_EFFECT} inside CI: {hac_ci_lower < TRUE_EFFECT < hac_ci_upper}")

Expected output:

Durbin-Watson statistic: 1.9628
  DW near 2.0 = little autocorrelation in residuals.
  DW < 1.5 = positive serial correlation.
  DW > 2.5 = negative serial correlation.
  Apply HAC standard errors regardless -- DW only tests AR(1) structure.

=== Standard error comparison ===
  OLS SE on ai_on  : 0.0036
  HAC SE on ai_on  : 0.0037
  OLS t-stat       : 16.83
  HAC t-stat       : 16.41

  HAC 95% CI: [0.0535,  0.0680]
  True effect 0.06 inside CI: True

The Durbin-Watson statistic near 2.0 (1.9628) indicates very little AR(1) autocorrelation in the residuals on this synthetic dataset, so the HAC and OLS standard errors are nearly identical. The HAC 95% CI [0.0535, 0.0680] contains the true effect of 0.060, confirming the adjusted estimate is valid.

In production LLM platforms where demand correlates across consecutive hours (morning surges, lunchtime dips, evening peaks), positive serial correlation causes OLS standard errors to understate uncertainty. I've seen teams skip this step and report t-statistics of 20+ on effects that don't hold up.

HAC corrections in those settings bring those numbers down to realistic levels and occasionally flip a "significant" result to inconclusive. The flip to inconclusive is the method working correctly. Apply HAC by default in any time-series regression: it costs nothing when autocorrelation is absent, and it provides real protection when it's present.

The nlags parameter deserves deliberate choice. A reasonable default is the number of slots you'd expect your largest demand cycle to span. If your platform shows strong hour-of-day patterns and you're using 30-minute slots, set nlags=4 or nlags=6 to cover the two-to-three-hour neighborhood. If you use two-hour slots, nlags=2 or nlags=3 usually covers the relevant range.

Step 5: Bootstrap Confidence Intervals

HAC standard errors correct for autocorrelation under the assumption that the autocorrelation structure follows a specific parametric form. Bootstrap CIs make no such assumption. They quantify estimation uncertainty by resampling slots with replacement and recomputing the estimator each time.

def bootstrap_ci(slots, B=500, seed=7):
    """Bootstrap CIs treating each slot as an independent observation.
  
    Each slot's ai_on_lag1 value is fixed from the original treatment schedule.
    Resampling slots with replacement while keeping their original lag values
    correctly quantifies estimation uncertainty without destroying the lag structure.
    """
    rng  = np.random.default_rng(seed)
    n    = len(slots)
    naive_ates, adj_ates, carryover_ests = [], [], []

    for _ in range(B):
        idx = rng.integers(0, n, size=n)
        s   = slots.iloc[idx]  # ai_on_lag1 stays as the original slot's value

        X_n = sm.add_constant(s['ai_on'])
        naive_ates.append(sm.OLS(s['mean_task_completed'], X_n).fit().params['ai_on'])

        X_a = sm.add_constant(s[['ai_on', 'ai_on_lag1']])
        m   = sm.OLS(s['mean_task_completed'], X_a).fit()
        adj_ates.append(m.params['ai_on'])
        carryover_ests.append(m.params['ai_on_lag1'])

    naive_ci     = np.percentile(naive_ates,     [2.5, 97.5])
    adj_ci       = np.percentile(adj_ates,       [2.5, 97.5])
    carryover_ci = np.percentile(carryover_ests, [2.5, 97.5])

    print(f"\n=== Bootstrap 95% confidence intervals (B={B}, seed={seed}) ===")
    print(f"  Naive ATE        : [{naive_ci[0]:.4f},  {naive_ci[1]:.4f}]  "
          f"(covers {TRUE_EFFECT}: {naive_ci[0] < TRUE_EFFECT < naive_ci[1]})")
    print(f"  Adjusted ATE     : [{adj_ci[0]:.4f},  {adj_ci[1]:.4f}]  "
          f"(covers {TRUE_EFFECT}: {adj_ci[0] < TRUE_EFFECT < adj_ci[1]})")
    print(f"  Carryover effect : [{carryover_ci[0]:.4f},  {carryover_ci[1]:.4f}]  "
          f"(covers {CARRYOVER}: {carryover_ci[0] < CARRYOVER < carryover_ci[1]})")

    return naive_ci, adj_ci, carryover_ci

naive_ci, adj_ci, carryover_ci = bootstrap_ci(slots)

Expected output:

=== Bootstrap 95% confidence intervals (B=500, seed=7) ===
  Naive ATE        : [0.0596,  0.0783]  (covers 0.06: True)
  Adjusted ATE     : [0.0541,  0.0683]  (covers 0.06: True)
  Carryover effect : [0.0175,  0.0320]  (covers 0.03: True)

Each bootstrap iteration resamples 48 slots with replacement, refits both the naive and adjusted OLS models, and records the key estimates. The 2.5th and 97.5th percentiles of those 500 replications give the bootstrap CIs.

Each slot brings its own ai_on_lag1 value from the original treatment schedule, so the lag structure is preserved within each bootstrap draw. The resampling captures estimation uncertainty without fabricating temporal relationships that didn't exist.

All three 95% CIs cover their respective ground truths. The naive ATE CI [0.0596, 0.0783] covers the true effect (0.060) but is shifted upward, consistent with the +0.009 positive bias. The adjusted ATE CI [0.0541, 0.0683] is centered closer to the true effect and is narrower. The carryover CI [0.0175, 0.0320] covers the true carryover of 0.030 and excludes zero, confirming that the carryover is statistically distinguishable from no persistence.

The excluded-zero result matters for the product decision: if the carryover CI included zero, you couldn't rule out that all the elevated AI-off outcomes were sampling noise rather than genuine persistence.

Validating Against the Ground Truth

Pull together the three point estimates against their known ground truths:

print("=" * 52)
print(f"{'Estimator':<30} {'Estimate':>8}  {'True':>6}  {'Bias':>7}")
print("-" * 52)
print(f"{'Naive OLS (no lag)':<30} {naive_ate:>8.4f}  {TRUE_EFFECT:>6.4f}  {naive_ate - TRUE_EFFECT:>+7.4f}")
print(f"{'Carryover-adjusted OLS':<30} {adj_ate:>8.4f}  {TRUE_EFFECT:>6.4f}  {adj_ate - TRUE_EFFECT:>+7.4f}")
print(f"{'Carryover coefficient':<30} {adj_carryover:>8.4f}  {CARRYOVER:>6.4f}  {adj_carryover - CARRYOVER:>+7.4f}")
print("=" * 52)

Expected output:

====================================================
Estimator                      Estimate    True     Bias
----------------------------------------------------
Naive OLS (no lag)               0.0688  0.0600  +0.0088
Carryover-adjusted OLS           0.0607  0.0600  +0.0007
Carryover coefficient            0.0244  0.0300  -0.0056
====================================================

The comparison table shows exactly what each estimator recovers against the known ground truth.

The naïve OLS overshoots by 0.0088 percentage points because it can't separate the direct AI routing effect from the carryover that inflates AI-on and adjacent AI-off slots. The adjusted OLS recovers the true effect to within 0.0007, well inside the width of any reasonable confidence interval. The carryover coefficient is 0.0244, compared with a true value of 0.030.

That's a systematic underestimate: the collinearity between ai_on and ai_on_lag1 in the 3-slot block structure produces this attenuation across all designs of this type.

The practical implication runs beyond this synthetic example. In a real LLM platform, carryover can be larger than the treatment effect. If the AI routing system fundamentally reshapes how the inference cluster allocates warm-cache slots across users, the next period will inherit a compute distribution shaped by AI routing, even after the routing AI is off.

Under those conditions, the naïve estimate could substantially overstate the effect you'd observe from a full always-on rollout, where no switching exists, and no carryover asymmetry accumulates.

Always estimate the carryover coefficient. If it's statistically significant and greater than 20% of your direct ATE estimate, the naïve estimate is unreliable for rollout decisions.

When Switchback Fails

Switchback solves marketplace interference under four conditions, and breaks under four others.

1. Carryover period longer than the slot length.

Violated assumption: (1) zero or bounded carryover.

If AI routing changes how the inference cluster pre-warms caches across multi-hour periods, the carryover half-life might exceed 60 or 90 minutes. A 30-minute slot length is shorter than the system's memory, and adding a single lag term won't capture the full persistence. You'll underestimate carryover and your direct effect estimate will remain biased.

The diagnostic: add progressively more lags and watch whether AIC keeps improving. If ai_on_lag3 and ai_on_lag4 still improve fit, your slot length is too short relative to system memory. Lengthening slots and adding more lag terms trade the same resource: fewer effective observations and wider confidence intervals.

2. Non-stationary demand confounding slots.

Violated assumption: (2) demand stationarity across the treatment schedule.

Weekday morning traffic surges, weekend evening spikes, and post-deployment adoption curves produce fundamentally different platform load conditions. If your treatment schedule places AI-on slots disproportionately in high-traffic windows and AI-off slots in low-traffic windows, the treatment coefficient absorbs demand differences as well as the routing AI's effect.

Randomizing the schedule within each day addresses this, as does including time-of-day fixed effects in the regression: a set of indicators for morning, afternoon, evening, and overnight absorbs within-day demand variation that would otherwise contaminate the treatment estimate.

3. Ramp-up effects at the first slot of each on-period.

Violated assumption: (3) no ramp-up at block boundaries.

In a real LLM platform, the first AI-on slot often underperforms subsequent slots. The routing model's cache is cold. The demand-prediction layer hasn't observed the current day's query distribution.

Including the cold-start slot alongside steady-state AI-on slots averages a low-performing initialization period with a high-performing equilibrium period, and the ATE estimate understates the steady-state effect you'd observe at full rollout. Standard practice is to drop the first slot of each on-period as a burn-in window and estimate the ATE from slots 2 and 3 of each block.

4. Period autocorrelation producing overconfident p-values.

Violated assumption: (4) residual autocorrelation addressed.

The Durbin-Watson diagnostic is a first check, but it only detects AR(1) autocorrelation. Real LLM platform time series often have daily seasonality, intraday autocorrelation at specific hours, and structural breaks after model version deployments.

Plot the full ACF of the model residuals: spikes at lags corresponding to meaningful demand cycles signal that your nlags parameter in cov_hac needs to increase, or you should switch to bootstrap CIs that don't assume any particular autocorrelation structure.

Failing to correct for autocorrelation is the most common source of false positives in switchback analyses at LLM platforms.

Two additional design-level failure modes are worth tracking.

Slot lengths under 15 minutes mean the platform hasn't cleared between switches: queue depth, in-flight session count, and cache state all carry over from the prior period, amplifying contamination and making AI-off periods non-representative of steady-state operations.

Slot lengths longer than 4 hours reduce the number of treatment-control pairs, shrinking the effective sample size and widening confidence intervals to the point where you can't detect plausible-sized effects.

The practical sweet spot for most LLM platform experiments is 30 minutes to 2 hours per slot, with final calibration determined by the carryover half-life estimated from early pilot data.

When to Use Switchback vs. Cluster Randomization

Switchback and cluster randomization solve the same interference problem through different mechanisms.

Cluster randomization partitions users into non-overlapping segments by geographic region, tenant ID, or organizational account, and assigns segments to treatment and control simultaneously. Switchback assigns the full population to treatment and control at different times.

Cluster randomization works well when you have enough separable segments and between-segment spillover is negligible. For an LLM SaaS platform with enterprise tenants on dedicated compute slices, cluster randomization by tenant is feasible: one tenant's routing decisions don't exhaust capacity for another's sessions.

For a consumer LLM platform where all users share the same inference fleet, capacity spillover crosses any user-segment boundary you draw, and cluster randomization can't isolate it.

Switchback is appropriate when spillover crosses segment boundaries or when you don't have enough separable clusters to run a properly powered cluster experiment.

Most large platforms use both: switchback for platform-wide infrastructure changes where no clean segment boundary exists, cluster randomization for features that can be scoped to a tenant or geographic region.

The choice comes down to where you can plausibly break the interference. Time is a natural boundary when the system clears faster than the slot length, so the platform fully processes the effects of one condition before switching to the next. Segment identity is a natural boundary when resource pools genuinely don't overlap. Where neither boundary holds, you're in causal estimation territory: synthetic control methods, difference-in-differences with matched controls, or structural models of the interference mechanism.

What to Do Next

If your switchback analysis shows a significant positive direct effect with a well-identified carryover term, the next hard question is whether the effect size justifies full rollout given the cost of the AI routing infrastructure. The premium model costs more per query than the standard model. Whether a 6 pp completion-rate lift covers that incremental inference cost depends on your product's monetization mechanics.

The carryover estimate shapes that decision too.

A large carryover coefficient means that some of the measured lift is dissipated once you switch to always-on routing, and the switching asymmetry disappears. The causal cost-benefit calculation requires the direct ATE, not the naïve estimate you'd get without the lag adjustment: revenue impact of the completion-rate gain, incremental inference cost at full traffic, and the confidence interval around each estimate before committing to an infrastructure investment.

If the routing AI shows heterogeneous effects across query types or user segments, the next analytical step is uplift modeling: building a model that predicts which queries benefit most from premium routing, so you route selectively and capture most of the task-completion gain at a fraction of the cost.

The causal identification work you've done here, including the switchback design, carryover adjustment, and HAC correction, gives you the unbiased population ATE you need as the ground-truth anchor for calibrating that uplift model.

The full companion code is at 06_switchback/, including the notebook with all five steps, the figure-generation scripts, and the dataset-generation code.

Let's say that your team built a routing layer that splits incoming queries between two models: queries with a confidence score below 0.85 go to a premium model, and those above 0.85 go to a cheaper distilled model. The premium model costs 5x as much as the cheaper one.

Your boss wants the answer that ends the debate: Is the premium model worth it for the queries it sees?

You can't run a clean A/B test, because routing is deterministic: a query at confidence 0.84 always gets premium, a query at 0.86 always gets cheap, and you can't randomize the assignment.

You also can't trust a naïve comparison of premium-routed users against cheap-routed users. Premium handles the harder queries by design (that's the reason you built the gate), so the two groups differ in query difficulty before either model touches them.

The threshold itself is your free experiment. Right at 0.85, the assignment flips, but the queries on either side of that boundary are essentially identical. A query at confidence 0.849 isn't meaningfully different from a query at 0.851. Any differences in outcomes between the two narrow groups stem solely from the routing decision. That's what regression discontinuity design (RDD) reads.

In this tutorial, you'll use Python to estimate the causal effect of premium routing on task completion using sharp RDD with local linear regression. You'll sweep bandwidths to test estimate stability, run a manipulation diagnostic, check robustness with a quadratic specification, and bootstrap 95% confidence intervals around every point estimate.

The LLM telemetry is a 50,000-user synthetic dataset with the ground-truth premium-routing effect baked in at +6 percentage points, so you can verify that RDD recovers it.

Companion code: every code block runs end-to-end in the companion notebook.

Table of Contents

• Why Threshold Routing is a Natural Experiment

• What Regression Discontinuity Actually Does

• Prerequisites

• Setting Up the Working Example

• Step 1: A Sharp RDD with Local Linear Regression

• Step 2: Try Different Bandwidths

• Step 3: Checking for Manipulation at the Threshold

• Step 4: Quadratic Specification as a Robustness Check

• Step 5: Bootstrap Confidence Intervals

• When Regression Discontinuity Fails

• What to Do Next

Why Threshold Routing is a Natural Experiment

The product reason this routing rule exists is to help your team spend the premium model budget where it earns its keep. Low-confidence queries are the harder ones, which is where a stronger model has the most upside. High-confidence queries already look easy enough for the cheap model to handle.

You'll see this routing direction across confidence-score gates for Q&A assistants, query-complexity gates in multi-model gateways like OpenRouter, safety-score gates in content moderation, and latency-budget gates that re-route when the cheap model would exceed a p99 latency budget.

The mechanism is the same in every case: a continuous score, a threshold, and a deterministic routing rule.

What makes this setup useful for causal inference is that users don't pick which model they get. A query lands, the system computes confidence, and the routing layer decides. Right at the threshold, the user's experience flips from premium to cheap based on a difference too small to be meaningful.

Again, a query at 0.849 confidence isn't shipping a different problem to the model than a query at 0.851. Anything that differs in outcomes between those two groups is the routing decision speaking. The underlying query is the same.

That local randomness is the experiment RDD reads from. You don't need a randomized control group, you don't need a propensity score. And you don't need an instrument, you need a sharp threshold that nobody can game.

What Regression Discontinuity Actually Does

The jump at the threshold is the causal effect, which is the number a product team can act on. RDD reads it by fitting two separate regression lines to the outcome: one for users just below the threshold and one for users just above. The vertical difference between those two fitted lines at the cutoff is the local average treatment effect at that point.

Graphically, picture task completion on the y-axis and query confidence on the x-axis. Completion generally trends with confidence (easier queries complete more often). At exactly 0.85, though, users below the cutoff get premium routing, and users above get cheap.

If premium routing helps, you'd see a sharp upward jump in task completion just below 0.85, then disappear just above. Approached from left to right with confidence rising, the visual reads as a downward step at 0.85, because you're moving from the premium-treated zone into the cheap-treated zone.

Figure 1. Conceptual schematic. Two outcome trajectories, one for premium-routed queries (confidence below 0.85) and one for cheap-routed queries (confidence above 0.85), meet at the threshold but don't match. The vertical gap between their endpoints at 0.85 is the local causal effect of premium routing.

That gap is identified under two named assumptions:

1. No manipulation of the running variable: Users (or your system) can't precisely nudge a query's confidence score across the cutoff. If anyone can game their score to land just below 0.85 and grab premium routing, the cutoff is no longer drawn at random, and RDD breaks.

2. Continuity of potential outcomes at the cutoff: Every other factor that affects task completion (query type, user expertise, workspace tenure, time of day) varies smoothly across 0.85. Only the routing assignment changes discontinuously at exactly the threshold. If a second product rule fires at 0.85 (a different logging level, a separate UI treatment, a retry policy), RDD will attribute that rule's effect to the routing decision.

These are the two assumptions you check before you trust the estimate. Step 3 below tests the first one. The second is a structural property of your system that you have to know cold.

Two practical choices shape every RDD: the bandwidth (how close to the cutoff to restrict the analysis) and the functional form (linear, quadratic, or local polynomial).

Narrow bandwidths cut potential bias by staying close to the local-randomization zone, but they shrink the sample. Linear specifications are stable, though they assume the underlying relationship can be approximated by a straight line on each side.

You'll try both linear and quadratic specifications at multiple bandwidths to see whether the answer holds.

The article uses sharp RDD throughout, since assignment is a deterministic function of confidence (below 0.85 always premium, above 0.85 always cheap). When the threshold is probabilistic and compliance is partial, the design is a fuzzy RDD, which requires an instrumental variables framework that you can implement using the rdrobust Python package.

Prerequisites

You need Python 3.11 or newer, comfort with pandas and statsmodels, and rough familiarity with linear regression and interaction terms.

Install the packages used in this tutorial:

pip install numpy pandas statsmodels matplotlib scipy

Here's what's happening: four standard scientific Python libraries plus matplotlib for the diagnostic visualization. Nothing exotic.

Clone the companion repo and generate the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

Here's what's happening: the data generator draws 50,000 users with a query_confidence score from a Beta(5,2) distribution, applies the routing rule (routed_to_premium = query_confidence < 0.85), and bakes a +6-percentage-point premium routing effect into task_completed. Same seed, same dataset, every time.

Setting Up the Working Example

The dataset simulates a SaaS product that routes queries between a premium and a cheap model based on confidence score. The threshold is 0.85, and the ground-truth causal effect of premium routing is +6 percentage points on task completion. You know the truth, so you can check whether RDD recovers it.

Load the data and look at the routing breakdown:

import numpy as np
import pandas as pd
import statsmodels.formula.api as smf

df = pd.read_csv("data/synthetic_llm_logs.csv")
print(f"Loaded {len(df):,} rows, {df.shape[1]} columns")

print("\nRouting breakdown:")
counts = df.routed_to_premium.value_counts().to_dict()
print(f"  Premium-routed (confidence < 0.85):  {counts.get(1, 0):,}")
print(f"  Cheap-routed   (confidence >= 0.85): {counts.get(0, 0):,}")

print("\nQuery confidence distribution:")
print(df.query_confidence.describe().round(3))

Expected output:

Loaded 50,000 rows, 16 columns

Routing breakdown:
  Premium-routed (confidence < 0.85):  38,874
  Cheap-routed   (confidence >= 0.85): 11,126

Query confidence distribution:
count    50000.000
mean         0.715
std          0.159
min          0.078
25%          0.611
50%          0.736
75%          0.838
max          0.998

Here's what's happening: about 78% of queries land below the 0.85 cutoff and get premium routing. The Beta(5,2) distribution is skewed toward the upper end, with a median of 0.736, and most of its mass still sits below 0.85. The remaining 22% are queries that the model already feels confident about, and they go to the cheap model.

Before any regression, look at the naïve comparison every product team is tempted to run:

naive = (
    df[df.routed_to_premium == 1].task_completed.mean()
    - df[df.routed_to_premium == 0].task_completed.mean()
)
print(f"Naive premium-vs-cheap effect: {naive:+.4f}  (ground truth = +0.06)")

Expected output:

Naive premium-vs-cheap effect: +0.0632  (ground truth = +0.06)

Here's what's happening: the naive estimate sits at +0.0632, which is suspiciously close to the truth. That's a coincidence of this specific synthetic dataset, where the only confounder of premium vs. cheap is query_confidence itself, and the outcome doesn't depend on confidence except through routing.

In production, you almost never get this lucky. User expertise, prompt phrasing, time of day, and a dozen unobserved query traits all correlate with confidence and with completion.

A naïve comparison in a real system can be off by 50% or more in either direction. RDD gives you identification that doesn't depend on the absence of hidden confounders.

Step 1: A Sharp RDD with Local Linear Regression

The basic sharp RDD estimator is a local linear regression. Restrict to users whose confidence sits within a bandwidth of the cutoff, fit separate linear slopes on each side, and read off the jump at 0.85.

cutoff = 0.85
bw = 0.10

near = df[(df.query_confidence > cutoff - bw)
          & (df.query_confidence < cutoff + bw)].copy()
near["below_cutoff"] = (near.query_confidence < cutoff).astype(int)
near["rc"] = near.query_confidence - cutoff

rdd_model = smf.ols(
    "task_completed ~ below_cutoff + rc + below_cutoff:rc",
    data=near,
).fit(cov_type="HC3")

effect = rdd_model.params["below_cutoff"]
print(f"RDD effect at cutoff (LATE): {effect:+.4f}")
print(f"Std error (HC3):             {rdd_model.bse['below_cutoff']:.4f}")
print(f"p-value:                     {rdd_model.pvalues['below_cutoff']:.4f}")
print(f"N users in [0.75, 0.95):     {len(near):,}")

Expected output:

RDD effect at cutoff (LATE): +0.0548
Std error (HC3):             0.0131
p-value:                     0.0000
N users in [0.75, 0.95):     21,689

Here's what's happening: the model fits separate intercepts and slopes on each side of 0.85 (below_cutoff is the side indicator, rc is confidence centered at the cutoff). The coefficient on below_cutoff reads off the vertical jump at the threshold, which is the local average treatment effect (LATE) for queries with confidence near 0.85. You get +0.0548, within sampling noise of the +0.06 ground truth.

Three notes on the specification. First, task_completed is binary, so this is a linear probability model. For RDD with a binary outcome at the cutoff, the linear probability model is standard practice because local linearity is the identifying assumption either way. Logit at the cutoff is an alternative if you need bounded predictions globally.

Second, the standard errors are used cov_type="HC3" to relax the homoskedasticity assumption, which is almost always wrong for binary outcomes.

Third, the dataset has one query per user with no within-user clustering, so cluster-robust standard errors aren't needed here. In a setting with multiple queries per user, you'd cluster on user_id.

The next diagnostic to look at is the confidence distribution near the cutoff. Figure 2 shows what 50,000 queries look like in the bandwidth window:

Figure 2. Real distribution from the 50,000-user synthetic dataset. Unlike the schematic in Figure 1, this shows the actual query density by confidence score, with the routing threshold annotated. The bottom panel counts how many queries land in each 2-percentage-point bin near the cutoff (2,461 / 2,481 / 2,335 / 2,229 / 2,048 across the 0.80–0.90 range). The roughly uniform spread is the visual signal that no manipulation is concentrating users on one side of the threshold.

Step 2: Try Different Bandwidths

Bandwidth choice matters. Too narrow and you have too few observations, so the confidence interval blows up. Too wide and you're extrapolating into regions where the linear specification is no longer a reasonable local approximation.

The honest move is to try multiple bandwidths and report whether the estimate holds.

results = []
for bw in [0.05, 0.10, 0.15, 0.20]:
    sub = df[(df.query_confidence > cutoff - bw)
             & (df.query_confidence < cutoff + bw)].copy()
    sub["below_cutoff"] = (sub.query_confidence < cutoff).astype(int)
    sub["rc"] = sub.query_confidence - cutoff

    m = smf.ols(
        "task_completed ~ below_cutoff + rc + below_cutoff:rc",
        data=sub,
    ).fit(cov_type="HC3")

    results.append({
        "bandwidth": bw,
        "n": len(sub),
        "effect": m.params["below_cutoff"],
        "se": m.bse["below_cutoff"],
        "p": m.pvalues["below_cutoff"],
    })

print(pd.DataFrame(results).round(4).to_string(index=False))

Expected output:

 bandwidth      n  effect     se       p
      0.05  11554  0.0635  0.0183  0.0005
      0.10  21689  0.0548  0.0131  0.0000
      0.15  29137  0.0618  0.0112  0.0000
      0.20  34074  0.0614  0.0107  0.0000

Here's what's happening: four bandwidths from ±0.05 to ±0.20 around the cutoff, refitting the same RDD specification at each. The estimates range from +0.0548 to +0.0635, all in the same neighborhood as the +0.06 ground truth, with standard errors that shrink as the bandwidth widens and grow as it narrows. Every p-value is well below 0.05. Whether the estimates are "stable" depends on the confidence intervals around them, which Step 5 produces with the bootstrap.

Step 3: Checking for Manipulation at the Threshold

RDD is valid only if users can't precisely manipulate the running variable around the cutoff. If your users (or your system) can nudge confidence scores just below 0.85 to force premium routing, you get a density spike at the cutoff, and the RDD estimate is contaminated.

The standard diagnostic is the McCrary density test, which checks whether the distribution of the running variable has a sharp jump at the cutoff. The simple version: bin the data tightly around 0.85 and check whether the counts on the two sides are similar.

print("User counts in 2-percentage-point bins around 0.85:")
for lo in [0.80, 0.82, 0.84, 0.86, 0.88]:
    hi = lo + 0.02
    cnt = ((df.query_confidence >= lo) & (df.query_confidence < hi)).sum()
    print(f"  [{lo:.2f}, {hi:.2f}):  n = {cnt:,}")

Expected output:

User counts in 2-percentage-point bins around 0.85:
  [0.80, 0.82):  n = 2,461
  [0.82, 0.84):  n = 2,481
  [0.84, 0.86):  n = 2,335
  [0.86, 0.88):  n = 2,229
  [0.88, 0.90):  n = 2,048

Here's what's happening: counts trend gently downward across the bandwidth because Beta(5,2) places more mass at higher confidence levels, and the density tapers as it approaches 1.0. There's no spike or dip at the 0.84–0.86 bin that straddles the cutoff. The 433-user spread across all five bins is consistent with smooth tapering of the underlying density.

That's the pattern you want when manipulation is absent. For a more rigorous test, the rddensity Python package implements the formal McCrary procedure with bias-corrected standard errors.

What manipulation looks like when it's real: a spike in users at confidences just barely below 0.85 (they're being nudged into premium routing) and a dip just above. If you see that pattern, the RDD estimate overstates the causal effect because the users right below 0.85 differ in motivation from those right above. They cared enough to manipulate the score, and they'd have shown different outcomes even under random routing.

Step 4: Quadratic Specification as a Robustness Check

If the true relationship between confidence and task completion isn't exactly linear, a local linear RDD can mistake the curvature for a jump. The standard robustness check allows quadratic terms on both sides of the cutoff and tests whether the estimate holds.

near = df[(df.query_confidence > cutoff - 0.10)
         & (df.query_confidence < cutoff + 0.10)].copy()
near["below_cutoff"] = (near.query_confidence < cutoff).astype(int)
near["rc"] = near.query_confidence - cutoff
near["rc2"] = near.rc ** 2

rdd_quad = smf.ols(
    "task_completed ~ below_cutoff + rc + below_cutoff:rc"
    " + rc2 + below_cutoff:rc2",
    data=near,
).fit(cov_type="HC3")

print(f"Linear RDD    (bw=0.10):  effect = +0.0548, p < 0.0001")
print(f"Quadratic RDD (bw=0.10):  effect = "
      f"{rdd_quad.params['below_cutoff']:+.4f}, "
      f"p = {rdd_quad.pvalues['below_cutoff']:.4f}")

Expected output:

Linear RDD    (bw=0.10):  effect = +0.0548, p < 0.0001
Quadratic RDD (bw=0.10):  effect = +0.0569, p = 0.0036

Here's what's happening: the quadratic specification adds squared terms and interactions with the cutoff indicator, allowing the relationship to curve differently on each side. The below_cutoff coefficient still captures the jump at the threshold, now under a more flexible specification.

The two estimates differ by 0.0022, both close to the +0.06 ground truth, and both are significant at p < 0.01. The answer doesn't change when you let the model bend.

When linear and quadratic specifications disagree noticeably, you have a real signal. With small samples (a few thousand at narrow bandwidths), the quadratic version can lose power because four extra parameters need data to be identified.

The standard move is to widen the bandwidth and re-run both specifications. If they still disagree at wider bandwidths, the linear approximation is wrong, and you should report both numbers.

Step 5: Bootstrap Confidence Intervals

Every point estimate in this article is a single number from a finite sample. The bootstrap quantifies how much that number would move under resampling, which is what a confidence interval describes.

def bootstrap_ci(df, cutoff, bw, quadratic=False, n_reps=500, seed=7):
    rng = np.random.default_rng(seed)
    near = df[(df.query_confidence > cutoff - bw)
              & (df.query_confidence < cutoff + bw)].copy()
    near["below_cutoff"] = (near.query_confidence < cutoff).astype(int)
    near["rc"] = near.query_confidence - cutoff
    if quadratic:
        near["rc2"] = near.rc ** 2
        formula = ("task_completed ~ below_cutoff + rc + below_cutoff:rc"
                   " + rc2 + below_cutoff:rc2")
    else:
        formula = "task_completed ~ below_cutoff + rc + below_cutoff:rc"

    n = len(near)
    estimates = np.empty(n_reps)
    for i in range(n_reps):
        sample = near.iloc[rng.integers(0, n, size=n)]
        m = smf.ols(formula, data=sample).fit()
        estimates[i] = m.params["below_cutoff"]
    return (np.percentile(estimates, 2.5), np.percentile(estimates, 97.5))


print("Linear RDD (bw=0.10):")
lo, hi = bootstrap_ci(df, cutoff, bw=0.10)
print(f"  effect = +0.0548   95% CI: [{lo:+.4f}, {hi:+.4f}]")

print("\nBandwidth sensitivity:")
for bw, eff in [(0.05, 0.0635), (0.10, 0.0548), (0.15, 0.0618), (0.20, 0.0614)]:
    lo, hi = bootstrap_ci(df, cutoff, bw=bw)
    print(f"  bw = {bw:.2f}   effect = {eff:+.4f}   "
          f"95% CI: [{lo:+.4f}, {hi:+.4f}]")

print("\nQuadratic RDD (bw=0.10):")
lo, hi = bootstrap_ci(df, cutoff, bw=0.10, quadratic=True)
print(f"  effect = +0.0569   95% CI: [{lo:+.4f}, {hi:+.4f}]")

Expected output:

Linear RDD (bw=0.10):
  effect = +0.0548   95% CI: [+0.0278, +0.0817]

Bandwidth sensitivity:
  bw = 0.05   effect = +0.0635   95% CI: [+0.0244, +0.0986]
  bw = 0.10   effect = +0.0548   95% CI: [+0.0278, +0.0817]
  bw = 0.15   effect = +0.0618   95% CI: [+0.0381, +0.0823]
  bw = 0.20   effect = +0.0614   95% CI: [+0.0420, +0.0808]

Quadratic RDD (bw=0.10):
  effect = +0.0569   95% CI: [+0.0205, +0.0959]

Here's what's happening: the bootstrap resamples the bandwidth-restricted data with replacement 500 times, refits the RDD on each replicate, and collects the below_cutoff coefficient. The 2.5th and 97.5th percentiles of those 500 estimates form the 95% interval. Every interval covers the +0.06 ground truth, every interval excludes zero, and the bandwidth sweep produces overlapping intervals.

That's quantitative stability, verified by resampling across the full bandwidth range. Intervals widen as the bandwidth shrinks and narrow as it grows. The quadratic interval is wider than the linear one because the four extra parameters absorb degrees of freedom.

One thing the intervals do NOT do on this dataset: exclude the naive +0.0632 estimate. That's because the data generator doesn't bake in confounding by query confidence. The only difference between the premium and cheap groups in expectations is the +6pp routing effect itself, so the naïve comparison is close to the truth.

Real systems are messier. In a production setting where unobserved query traits affect both the routing assignment and task completion, the naïve estimate would diverge from the RDD estimate, and the bootstrap intervals would tell you which one to trust.

When Regression Discontinuity Fails

RDD looks clean, but several specific failure modes can destroy the identification. Each one maps to a violation of one of the two named assumptions.

Users manipulate the running variable (violates assumption 1). The whole setup depends on users (or any upstream service) being unable to precisely control which side of the cutoff they land on. Any system that reveals the cutoff and gives users a way to influence their score (a retry mechanism, a prompt engineering workaround, a confidence-inflating trick) breaks RDD.

Run the density check in Step 3 every time. If you find manipulation, switch to a fuzzy RDD that treats the threshold as probabilistic, or abandon the approach.

Other policies fire at the same cutoff (violates assumption 2). If your product has additional rules that activate at 0.85 (a separate UI treatment, a different logging level, a different retry policy), RDD can't separate the routing effect from those other policy effects. Audit the full rule book for anything that shares the threshold.

The threshold has noise or overrides (violates assumption 1, in the structural sense). Maybe routing isn't strictly deterministic at 0.85 – it may have random jitter, or a second rule may override the main rule in some cases.

If assignment to the premium model isn't a deterministic function of query_confidence, you have a fuzzy RDD, which requires an instrumental variables framework. The rdrobust package handles both sharp and fuzzy designs.

Curvature masquerading as a jump (breaks the linear approximation that supports identification at the cutoff). Sharp RDD assumes linearity is a reasonable local approximation. When the underlying outcome-confidence relationship is strongly curved, the linear specification can mistake the bend for a jump.

Step 4's quadratic robustness check is the standard diagnostic. If linear and quadratic disagree, widen the bandwidth and re-run both.

Extrapolation bias (a continuity issue, reframed). RDD estimates are strictly local to the cutoff. The +0.06 effect at 0.85 tells you nothing about what premium routing would do for queries with confidence 0.30 or 0.99.

If you want a global average effect, you need a different technique: propensity methods, regression with confounder adjustment, or an actual experiment.

What to Do Next

RDD is the right tool when your AI feature is gated by a continuous score and a sharp threshold.

If your feature is gated by a user-controlled toggle, propensity score methods are a better fit. If it's gated by a staged rollout across workspaces, difference-in-differences handles it. If it's gated by rules you can't observe directly but that have a random component, instrumental variables is the right choice.

For production RDD analyses, use the rdrobust Python package. It gives you optimal bandwidth selection (Calonico, Cattaneo, and Titiunik 2014), bias-corrected standard errors, and a built-in plotting utility. The companion rddensity package implements the McCrary density test you saw informally in Step 3.

The from-scratch version in this tutorial shows the mechanics. The rd-packages stack is what you ship to a reviewer.

One thing the LATE doesn't do: tell you the effect for users far from the cutoff. If a +0.06 LATE at 0.85 is enough to keep premium routing in the pipeline, you're done. If you need to know what premium would do for the easy queries you're currently sending to cheap (or the hardest queries near the floor), the next step is a small randomized rollout in those zones, scored against the RDD estimate as a calibration check. Don't generalize the LATE without evidence.

The companion notebook for this tutorial lives here on GitHub. Clone the repo, generate the synthetic dataset, and run rdd_demo.ipynb to reproduce every code block from this tutorial.

Threshold routing is one of the most common patterns in production LLM systems, and every confidence-gated routing decision in your stack is a potential RDD. Run the analysis.

Your product shipped a new agent mode last quarter. Users have to tap the "Try agent mode" toggle to enable it. The dashboard numbers look stunning: agent-mode users complete 21 percentage points more tasks than non-users. The CPO calls it the best feature launch of the year.

But you know something's off. Heavy-engagement users opt into new features constantly, while light users ignore toggles entirely. That 21-point gap measures the agent's effect combined with the pre-existing gap between power users and the rest of your base.

This is the Opt-In Trap. It shows up in every generative AI product that ships features behind a user-controlled toggle: "Try our AI assistant," "Enable smart replies," "Turn on code suggestions." Users who click to opt in differ systematically from those who scroll past. Any naïve comparison between the two groups collapses the feature's causal effect into whatever made those users opt in in the first place.

Running an AI feature behind a toggle is a product experiment. The hypothesis: the feature improves outcomes for users who adopt it.

Unlike an A/B test, where the coin flip creates two otherwise-identical populations, the toggle creates two populations that differ before they even make a choice. That pre-existing difference is the measurement problem, and a t-test on dashboard numbers can't fix it.

Propensity score methods are statistical tools that data scientists use to separate adoption bias from the feature's actual effect. They reweight (or rematch) your comparison so that opted-in and non-opted-in groups look comparable on observable characteristics, approximating what a randomized experiment would have given you.

This tutorial walks through the full pipeline (propensity estimation, inverse-probability weighting, nearest-neighbor matching, balance diagnostics, and bootstrap confidence intervals) on a 50,000-user synthetic SaaS dataset where the ground-truth causal effect is known. You'll estimate it, quantify uncertainty, and see where the approach silently breaks.

Companion code: every code block runs end-to-end in the companion notebook at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/02_propensity_opt_in. The notebook (psm_demo.ipynb) has all outputs pre-executed, so you can read along on GitHub before running anything locally.

Table of Contents

• Why Opt-in Features Break Naïve Comparisons

• What Propensity Scores Actually Do

• Prerequisites

• Setting Up the Working Example

• Step 1: Estimate the Propensity Score

• Step 2: Inverse-Probability Weighting

• Step 3: Nearest-Neighbor Matching

• Step 4: Check Covariate Balance

• Step 5: Bootstrap Confidence Intervals

• When Propensity Score Methods Fail

• What to Do Next

Why Opt-in Features Break Naïve Comparisons

The math of an A/B test is elegant because of one assumption: treatment is assigned independent of everything else. Flip a coin: half your users get agent mode, and the coin flip breaks every possible confound by construction. The opt-in world has no coin.

Three mechanisms make opt-in comparisons misleading.

1. Selection on engagement

Power users click everything. If your heavy-engagement cohort opts into agent mode at 65 percent and your light-engagement cohort opts in at 12 percent, you've stacked the opt-in group with users who were going to complete more tasks anyway.

That compositional imbalance accounts for most of the observed lift on its own, before the agent does any work.

2. Selection on intent

Users who opt into a new feature often have a specific use case in mind. A developer who clicks "Try code suggestions" already has code to write. That user would have shown higher task completion even with the control UI.

3. Selection on risk tolerance

Early adopters tolerate rough edges. A user who clicks "Try beta" and sees slow latency sticks around, but a risk-averse user bounces.

Your opt-in group is enriched for people willing to put up with bad experiences, which affects every downstream metric you might measure.

All three produce the same symptom: a raw comparison of opted-in users against everyone else that can overstate the feature's causal effect by 2x or more, depending on how concentrated opt-in is among your heaviest users.

On the synthetic dataset in this tutorial, the naïve comparison inflates a true +8pp effect to +21pp, a 2.6x overshoot. Propensity score methods exist to correct this.

What Propensity Scores Actually Do

Figure 1: Schematic propensity score distributions for two hypothetical groups. The opted-in group (red) skews toward higher propensities, while the non-opted-in group (blue) skews lower.

In the above figure, the bracketed strip below the x-axis splits the score range into three zones: a control-heavy region at low propensities where few treated users exist, a region of common support in the middle where both groups are well represented, and a treatment-heavy region at high propensities where few controls exist. Propensity score methods operate within the common-support region by reweighting or rematching so that the two groups appear balanced on observables. The extremes are either trimmed out or handled with caution.

The propensity score is the probability that a user opts in given their observable characteristics. Estimate this probability well, and you can use it to reweight your sample so that opted-in and non-opted-in users look similar on observables, just as they would have if opt-in had been randomized.

Two practical strategies use the propensity score:

• Inverse-probability weighting (IPW) assigns each user a weight equal to the inverse of their probability of receiving the treatment they actually received. Opted-in users get weighted by 1/P(opt-in). Non-opted-in users get weighted by 1/P(no opt-in). After weighting, the two groups are balanced on observables, and the weighted difference in outcomes approximates the average treatment effect.

• Matching pairs each opted-in user with one or more non-opted-in users who have similar propensity scores. The average outcome difference between matched pairs estimates the average treatment effect on the treated (ATT): what opt-in users actually gained by opting in.

Both methods rest on three identification assumptions working together.

1. First, unconfoundedness: every observable variable that drives opt-in and affects the outcome is in your propensity model.

2. Second, overlap (also called positivity): every user has some nonzero probability of opting in and some nonzero probability of staying out.

3. Third, no interference: one user's opt-in decision does not affect another user's outcome (the stable-unit-treatment-value assumption, or SUTVA.

Violate any one of these and the estimate is biased even when the other two hold. The failure modes at the end of this tutorial walk through each one.

Prerequisites

You'll need Python 3.11 or newer, comfort with pandas and scikit-learn, and rough familiarity with logistic regression.

Install the packages for this tutorial:

pip install numpy pandas scikit-learn matplotlib

Here's what's happening: four packages cover the full pipeline. Pandas loads the data, NumPy handles weights and array arithmetic, scikit-learn fits the propensity model and runs nearest-neighbor matching, and matplotlib renders the overlap diagnostic.

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

Here's what's happening: the clone pulls the companion repo, and generate_data.py produces the shared synthetic dataset used across the series. Seed 42 keeps the dataset reproducible, and 50,000 users give clean signal for every estimator in this tutorial. The output CSV lands at data/synthetic_llm_logs.csv.

Setting Up the Working Example

The synthetic dataset simulates a SaaS product where users can opt into an agent mode that uses a more expensive model. With fifty thousand users, opt-in rates differ sharply by engagement tier: heavy users opt in at 65 percent, medium users at 35 percent, and light users at 12 percent.

The ground-truth causal effect baked into the data generator is +8 percentage points on task completion for users who opted in. The naive comparison inflates this to around +21 percentage points because selection bias stacks the opted-in group with your most engaged users.

Knowing the ground truth is what lets you verify that your propensity score method recovers it.

Load the data and see the selection problem:

import pandas as pd

df = pd.read_csv("data/synthetic_llm_logs.csv")

print(df.groupby("engagement_tier").opt_in_agent_mode.mean().round(3))

naive_effect = (
    df[df.opt_in_agent_mode == 1].task_completed.mean()
    - df[df.opt_in_agent_mode == 0].task_completed.mean()
)
print(f"\nNaive opt-in effect: {naive_effect:+.4f}")

Expected output:

engagement_tier
heavy     0.647
light     0.120
medium    0.353
Name: opt_in_agent_mode, dtype: float64

Naive opt-in effect: +0.2106

Here's what's happening: you load 50,000 rows, group by engagement tier, and print the opt-in rate inside each group. Heavy users opt in far more than light users, which is the selection-on-engagement pattern baked into the data. The naïve effect lands at +0.2106 (21 percentage points), nearly three times the ground truth of +0.08. That gap is exactly what propensity score methods have to remove.

Step 1: Estimate the Propensity Score

The propensity score is the output of a model that predicts opt-in from observable characteristics. Logistic regression is the right starting point because it's interpretable and fast, but watch the balance diagnostics in Step 4: if any weighted SMD stays above 0.1, the logistic model is missing an interaction, and gradient boosting is the next move.

For this dataset, the relevant observables are engagement tier and query confidence. In a real product, you'd include every variable you think drives opt-in: device type, tenure, plan tier, and historical usage patterns.

from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score

X = pd.get_dummies(
    df[["engagement_tier", "query_confidence"]],
    drop_first=True
).astype(float)
y_treat = df.opt_in_agent_mode

ps_model = LogisticRegression(max_iter=1000).fit(X, y_treat)
df["propensity"] = ps_model.predict_proba(X)[:, 1]

# Basic sanity checks
print(df.groupby("engagement_tier").propensity.mean().round(3))
print(
    f"\nPropensity range (treated):  "
    f"{df[df.opt_in_agent_mode == 1].propensity.min():.3f} - "
    f"{df[df.opt_in_agent_mode == 1].propensity.max():.3f}"
)
print(
    f"Propensity range (control):  "
    f"{df[df.opt_in_agent_mode == 0].propensity.min():.3f} - "
    f"{df[df.opt_in_agent_mode == 0].propensity.max():.3f}"
)
print(f"Propensity model AUC: {roc_auc_score(y_treat, df.propensity):.3f}")

Expected output:

engagement_tier
heavy     0.646
light     0.120
medium    0.353
Name: propensity, dtype: float64

Propensity range (treated):  0.114 - 0.675
Propensity range (control):  0.114 - 0.673
Propensity model AUC: 0.744

Here's what's happening: you encode the engagement tier as dummy variables, keep query confidence continuous, and fit a logistic regression model. The predicted probability from the model is each user's propensity score.

Scikit-learn LogisticRegression applies L2 regularization by default (C=1.0), which shrinks propensities slightly toward 0.5. For production use, you can set penalty=None if you want an unregularized fit.

Mean propensity inside each engagement tier recovers the true opt-in rate for that tier almost exactly, so the model is calibrated. The AUC of 0.744 confirms the model discriminates between opt-ins and non-opt-ins well above chance (0.5).

And the propensity ranges overlap between treated and control groups (both span roughly 0.11 to 0.67), which is the visual overlap condition.

Figure 2: Two views of the same positivity check on the real 50,000-user synthetic dataset.

In the figure above, the top panel plots smooth kernel density curves of the fitted propensity scores for each group. The three peaks align with the three engagement tiers (light at p ≈ 0.12, medium at p ≈ 0.35, heavy at p ≈ 0.65), as expected, because the opt-in rate is tier-driven. The bottom panel translates that same distribution into raw counts per tier: every tier contains thousands of both opted-in and non-opted-in users, which is exactly what positivity requires.

Where Figure 1 schematically illustrated the idea, this figure shows that it holds for the data, so the weighting and matching that follow will have real counterfactuals to work with.

Step 2: Inverse-Probability Weighting

IPW assigns each user a weight inversely proportional to their propensity. An opted-in user with a 0.12 propensity is rare (a light user who still opted in despite low engagement) and carries information about 1 / 0.12 ≈ 8 similar users in the population. A control user with a 0.12 propensity is the expected case for light users who stayed out, so they're common and get a weight of 1 / (1 - 0.12) ≈ 1.14.

import numpy as np

# ATE weights: 1/P(treat) for treated, 1/P(no treat) for control
df["ipw"] = np.where(
    df.opt_in_agent_mode == 1,
    1 / df.propensity,
    1 / (1 - df.propensity)
)

t = df[df.opt_in_agent_mode == 1]
c = df[df.opt_in_agent_mode == 0]
ate_ipw = (
    (t.task_completed * t.ipw).sum() / t.ipw.sum()
    - (c.task_completed * c.ipw).sum() / c.ipw.sum()
)
print(f"IPW average treatment effect (ATE): {ate_ipw:+.4f}")

# ATT: what opt-in users actually gained
df["ipw_att"] = np.where(
    df.opt_in_agent_mode == 1,
    1,
    df.propensity / (1 - df.propensity)
)
t = df[df.opt_in_agent_mode == 1]   # re-slice now that ipw_att is in df
c = df[df.opt_in_agent_mode == 0]
treated_mean = t.task_completed.mean()
control_w_mean = (c.task_completed * c.ipw_att).sum() / c.ipw_att.sum()
att_ipw = treated_mean - control_w_mean
print(f"IPW average treatment effect on treated (ATT): {att_ipw:+.4f}")

Expected output:

IPW average treatment effect (ATE): +0.0851
IPW average treatment effect on treated (ATT): +0.0770

Here's what's happening: first, you compute ATE weights for every user and take the weighted difference in task completion between opted-in and non-opted-in groups. Then you compute ATT weights, which reweight only the control group to match the treated group's covariate distribution, and compute the average treatment effect on the treated.

ATE answers the population question: what's the effect on a random user who might or might not have opted in anyway? ATT answers the user question: What did opt-in users actually gain? On this dataset, ATE lands at +0.0851 and ATT at +0.0770, both close to the ground-truth +0.08 and a massive improvement over the naive +0.2106.

The distinction matters in practice. Deciding whether to roll the feature out to users who haven't opted in calls for ATE. Reporting on the value opt-in users captured calls for ATT.

Step 3: Nearest-Neighbor Matching

Matching takes a different approach: pair each opted-in user with the non-opted-in user whose propensity score is closest, then take the average outcome difference across matched pairs. The result estimates ATT.

from sklearn.neighbors import NearestNeighbors

treated_ps = df[df.opt_in_agent_mode == 1][["propensity"]].values
control_ps = df[df.opt_in_agent_mode == 0][["propensity"]].values

nn = NearestNeighbors(n_neighbors=1).fit(control_ps)
_, idx = nn.kneighbors(treated_ps)

treated_outcomes = df[df.opt_in_agent_mode == 1].task_completed.values
matched_control_outcomes = (
    df[df.opt_in_agent_mode == 0].task_completed.values[idx.flatten()]
)

att_match = (treated_outcomes - matched_control_outcomes).mean()
print(f"1-NN matching ATT: {att_match:+.4f}")

Expected output:

1-NN matching ATT: +0.0752

Here's what's happening: you extract propensity scores for each group, fit a nearest-neighbor index on the control group, and find the single closest control user for every treated user.

The NearestNeighbors index allows the same control user to be selected as the match for multiple treated users, so this is a matching-with-replacement case.

You pull the outcomes for each treated user and their matched control, take the difference per pair, and average across pairs. The result estimates what opt-in users gained compared to very similar users who did not opt in.

The +0.0752 result lands close to the ground truth of +0.08 but slightly below IPW ATT, typical of 1-NN matching because a single nearest neighbor is a high-variance estimator.

Two variants are worth knowing. Matching with replacement (what you just ran) allows a single control user to serve as a match for multiple treated users, reducing bias when good matches are scarce but inflating variance.

Matching without replacement assigns each control user to at most one treated user, which keeps variance lower but forces poor-quality pairings when the treated group dwarfs the available controls.

For most production analyses, k-nearest-neighbor matching with k = 3-5 and replacement is a sensible default.

Step 4: Check Covariate Balance

Propensity score methods work only if they actually balance the covariates between groups. You need to verify that they did, because if the balance fails, your estimate is wrong.

The standard diagnostic is the standardized mean difference (SMD) for each covariate. SMD compares the treated group mean to the control group mean, divided by the pooled standard deviation.

Before weighting, SMDs tell you how imbalanced the raw groups are. After weighting, they should be small (|SMD| < 0.1 is the conventional cutoff).

def smd(treated_vals, control_vals, treated_w=None, control_w=None):
    """Standardized mean difference, optionally with weights."""
    if treated_w is None:
        treated_w = np.ones(len(treated_vals))
    if control_w is None:
        control_w = np.ones(len(control_vals))
    t_mean = np.average(treated_vals, weights=treated_w)
    c_mean = np.average(control_vals, weights=control_w)
    pooled_std = np.sqrt((treated_vals.var() + control_vals.var()) / 2)
    return (t_mean - c_mean) / pooled_std

engagement_heavy = (df.engagement_tier == "heavy").astype(float).values
qc = df.query_confidence.values
tr = (df.opt_in_agent_mode == 1).values

covariates = {
    "engagement_tier_heavy": engagement_heavy,
    "query_confidence": qc,
}

print(f"{'Covariate':<30} {'Raw SMD':>10} {'Weighted SMD':>15}")
for name, vals in covariates.items():
    smd_raw = smd(vals[tr], vals[~tr])
    smd_weighted = smd(
        vals[tr], vals[~tr],
        treated_w=df[tr].ipw.values,
        control_w=df[~tr].ipw.values,
    )
    print(f"{name:<30} {smd_raw:>+10.3f} {smd_weighted:>+15.3f}")

Expected output:

Covariate                         Raw SMD    Weighted SMD
engagement_tier_heavy              +0.742          +0.002
query_confidence                   -0.032          -0.003

Here's what's happening: the helper computes the standardized mean difference for any covariate, with optional IPW weights.

You then print raw and weighted SMDs for each covariate. The raw SMD on engagement_tier_heavy is +0.742 (heavy users opt in far more than everyone else), and the weighted SMD drops to +0.002, a clean pass. Query confidence was already close to balanced on the raw data, and weighting keeps it that way. If any weighted SMD came back above 0.1 in absolute value, your propensity model would be missing something; the fix is usually richer features or interaction terms in the logistic regression.

Visually, Figure 2 above confirmed what the SMDs now confirm numerically: the overlap condition holds, and balance is achievable.

Step 5: Bootstrap Confidence Intervals

Point estimates are only half the story. Any estimate you report to a product team needs an interval that tells them whether +0.08 is distinguishable from +0.03 or from +0.12. Analytic standard errors for IPW and matching are tricky because of the estimated propensity score, so the simplest and most honest move is the non-parametric bootstrap.

def estimate_all(sample):
    """Return (ATE_IPW, ATT_IPW, ATT_match) on a bootstrap sample."""
    s = sample.copy()
    X_s = pd.get_dummies(
        s[["engagement_tier", "query_confidence"]], drop_first=True
    ).astype(float)
    ps = LogisticRegression(max_iter=1000).fit(X_s, s.opt_in_agent_mode)
    s["p"] = ps.predict_proba(X_s)[:, 1]

    s["w_ate"] = np.where(
        s.opt_in_agent_mode == 1, 1 / s.p, 1 / (1 - s.p)
    )
    s["w_att"] = np.where(
        s.opt_in_agent_mode == 1, 1, s.p / (1 - s.p)
    )
    t, c = s[s.opt_in_agent_mode == 1], s[s.opt_in_agent_mode == 0]

    ate = (
        (t.task_completed * t.w_ate).sum() / t.w_ate.sum()
        - (c.task_completed * c.w_ate).sum() / c.w_ate.sum()
    )
    att = t.task_completed.mean() - (
        (c.task_completed * c.w_att).sum() / c.w_att.sum()
    )
    nn_b = NearestNeighbors(n_neighbors=1).fit(c[["p"]].values)
    _, idx_b = nn_b.kneighbors(t[["p"]].values)
    match = (
        t.task_completed.values
        - c.task_completed.values[idx_b.flatten()]
    ).mean()
    return ate, att, match

rng = np.random.default_rng(7)
n_reps = 500
results = np.zeros((n_reps, 3))
for i in range(n_reps):
    boot = df.iloc[rng.integers(0, len(df), size=len(df))]
    results[i] = estimate_all(boot)

for name, col in zip(["IPW ATE", "IPW ATT", "1-NN ATT"], range(3)):
    lo, hi = np.percentile(results[:, col], [2.5, 97.5])
    print(f"{name:<10} 95% CI: [{lo:+.4f}, {hi:+.4f}]")

Expected output:

IPW ATE    95% CI: [+0.0745, +0.0954]
IPW ATT    95% CI: [+0.0687, +0.0865]
1-NN ATT   95% CI: [+0.0659, +0.0940]

Here's what's happening: you resample the dataset with replacement 500 times, refit the propensity model, and recompute each estimator on each resample, and take the 2.5th and 97.5th percentiles of the bootstrap distribution as the 95% confidence interval. All three intervals cover the ground-truth +0.08 and exclude the naive +0.21 by a wide margin.

The IPW ATT interval is the tightest because ATT reweights only the control group. The 1-NN matching interval is the widest because single-neighbor matching discards control users outside the matched set.

Running this once takes about 90 seconds on a laptop. For a stakeholder report, anchor the headline to the point estimate and cite the interval so the team sees the uncertainty alongside the number.

When Propensity Score Methods Fail

Propensity scores make opt-in comparisons rigorous when their assumptions hold. They produce biased estimates that look clean when those assumptions fail.

Four common failure modes map to the three identification assumptions from earlier.

1. Unmeasured Confounders (Violate Unconfoundedness)

If something drives both opt-in and your outcome but isn't in your propensity model, IPW and matching produce biased estimates. This is the most common failure in practice.

An example: users who opt into agent mode are also the users who follow your engineering blog and read release notes. If blog-reading behavior raises task completion independently of the feature, missing that signal attributes the effect to agent mode, inflating your estimate.

The only real defense is domain knowledge about what drives opt-in, richer feature engineering in your propensity model, and formal sensitivity tools (Rosenbaum bounds, E-values) that quantify how strong an unmeasured confounder would have to be to overturn the result.

2. Positivity (Overlap) Failures (Violates Overlap)

If some users have near-zero probability of opting in (or near-one), you've got no comparable counterfactual for them. I

PW creates extreme weights (1 / 0.001 = 1,000) that let a single outlier dominate the estimate. So matching is forced into poor-quality pairings.

Check propensity histograms and trim propensities outside [0.05, 0.95] before weighting if extreme values exist.

3. Misspecified Propensity Models (Degrade Unconfoundedness in Practice)

A linear logistic regression can't capture nonlinear relationships. If opt-in depends on the interaction between engagement tier and query confidence (power users with complex queries opt in, while light users pass), a main-effects model misses that and produces poor balance.

Use flexible models (for example, gradient boosting on the propensity score or regression adjustment on top of weighting) and always check the balance after weighting. Poor balance after weighting is the primary signal of misspecification.

4. Spillovers Between Users (Violates SUTVA)

Propensity score methods assume your users are independent. If one user opting into agent mode affects another user's task completion (for example, teammates adopting the feature together in shared workspaces), your estimated effect includes the spillover.

This violates the stable-unit-treatment-value-assumption, and handling it cleanly requires a different toolkit: either cluster randomization for features adopted at the workspace level or network-aware experimental designs for user-level spillovers.

These failure modes stay invisible in your regression coefficients. They surface as estimates that look good on paper but don't hold up when the feature rolls out to a broader audience.

Run balance diagnostics, check overlap plots, and document what you might have missed: those are your only real defenses.

What to Do Next

Propensity score methods are the right tool when your feature ships behind an opt-in toggle and you've got rich covariates to model selection with.

If opt-in follows a crisp rule (a threshold on query complexity, a paid-tier gate), regression discontinuity fits better. If you suspect unobserved confounders and have an external randomization source (randomized rollout noise, rate-limit-triggered routing), instrumental variables will do better.

To guard your estimate against propensity misspecification, doubly robust estimators combine propensity weighting with regression adjustment and stay consistent if at least one of the two component models is correctly specified.

The companion notebook for this tutorial lives here. Clone the repo, generate the synthetic dataset, and run psm_demo.ipynb (or psm_demo.py) to reproduce every code block, every number, and every figure from this tutorial.

When an AI feature ships behind a toggle, the naïve opt-in comparison is usually the wrong number. Propensity score methods give you "users comparable to those who clicked this" as your counterfactual, and the bootstrap gives you an interval you can defend when a stakeholder asks how sure you are.

The problem is that wave 2 workspaces are still waiting, a product-wide onboarding redesign shipped the same Tuesday, and week 20 also coincided with a quarterly engagement bump. Any comparison between the two groups after week 20 mixes the feature's causal effect with the redesign, the seasonality, and whatever selection criteria determined which workspaces landed in wave 1 in the first place.

This is how most enterprise SaaS teams ship AI features in 2026: one workspace at a time, in waves, on a rollout calendar. Randomization doesn't happen, and because randomization doesn't happen, A/B testing can't give you a clean causal effect. The result is a number on a dashboard that everyone argues over.

Call this the Rollout Calendar Trap: you have real data, a real experiment structure, and a completely invalid comparison. For data scientists shipping AI features in waves, it's the primary source of bad causal claims downstream.

Product experimentation for generative AI features follows this exact pattern: the hypothesis is that the AI feature causes higher engagement, and the wave structure is supposed to test it.

The wave calendar replaced the coin flip, and that substitution breaks the math. A simple A/B comparison assumes randomized assignment that the rollout never produced, so the measurement tool fails even when the experiment design is sound.

Difference-in-differences is the causal inference method that fixes this. It subtracts the time trend by comparing how outcomes shift across time periods for each group, giving you a defensible causal estimate even without randomization.

In this tutorial you'll use it to measure the true causal effect of an AI feature rolled out across enterprise workspaces, with working Python code against a synthetic SaaS product dataset.

By the end you'll know how to run a DiD estimate, how to test its parallel-trends assumption, and what to do when that assumption fails.

Table of Contents

• Why A/B Testing Breaks for Staged Rollouts

• What Difference-in-Differences Does

• Prerequisites

• Setting Up the Working Example

• Step 1: A Simple 2x2 DiD

• Step 2: Regression DiD with Fixed Effects

• Step 3: Checking the Parallel-Trends Assumption

• When Difference-in-Differences Fails

• What to Do Next

Why A/B Testing Breaks for Staged Rollouts

Random assignment is the engine that makes A/B testing a valid causal method. When you flip a coin to decide which user gets the feature, the treatment and control groups end up with identical distributions of every confounder (any variable that affects both who gets treatment and what outcome you measure). Any difference in outcomes after assignment is the causal effect of the treatment. Full stop.

A staged rollout across enterprise workspaces breaks that engine in three ways:

1. The wave assignment isn't random.

Product teams choose wave 1 workspaces for various reasons: they have the most engaged admins, the largest seat counts, or the best relationship with customer success. Those reasons correlate directly with your outcome. Wave 1 workspaces were going to show higher engagement anyway, feature or no feature.

2. The calendar introduces a time trend

Between week 20 (wave 1 launch) and week 30 (wave 2 launch), your product gets better, your onboarding improves, your sales team lands bigger customers. Any naïve "engagement after week 20 minus engagement before week 20" comparison picks up all of that along with the feature's effect.

3. Adoption inside treated workspaces is itself selective

Even inside a workspace that received the feature, not every user turns it on. Power users do, and less engaged users often wait months. Comparing users who used the feature against users who didn't introduces selection bias, where the groups differ systematically before you even measure the outcome, on top of the non-random workspace assignment.

A/B testing assumes none of these three problems exist. Staged rollouts guarantee all three. The naïve comparison gives you a number, and that number measures engagement theater.

What Difference-in-Differences Does

Difference-in-differences compares the change in outcomes over time between a treated group and a control group. Subtracting one change from the other cancels any shared time trend (product improvements, seasonality, onboarding changes) because both groups experience it equally, leaving you with just the treatment effect.

Here's a concrete example. Imagine tracking quarterly revenue for coffee shops in two neighborhoods. One neighborhood gets a new competitor in Q3, the other doesn't.

Both neighborhoods experience the same underlying market trends, a local economic upturn, and holiday seasonality. DiD isolates the competitor's impact by subtracting whatever revenue shift happened in both neighborhoods.

Your staged rollout sets up the exact same structure: wave 1 workspaces are the neighborhood with the new entrant, wave 2 is the comparison.

The math formalizes this as a 2x2 table, where rows are groups (treated, control), columns are time periods (pre, post), and each cell holds the mean outcome for that group in that period:

• A = mean task completion for wave 1 users before week 20 (coffee shops: Q2 revenue, neighborhood with incoming competitor)

• B = mean task completion for wave 1 users after week 20 (coffee shops: Q3 revenue, same neighborhood)

• C = mean task completion for wave 2 users before week 20 (coffee shops: Q2 revenue, the untouched neighborhood)

• D = mean task completion for wave 2 users after week 20 (coffee shops: Q3 revenue, same)

                         Pre     Post
Treated (wave 1):         A       B
Control (wave 2):         C       D

Naive post-period gap:   B - D     (contaminated by group differences)
Naive treated change:    B - A     (contaminated by time trend)
DiD:                 (B - A) - (D - C)   ← the causal effect

B - A is wave 1's change, but it includes both the treatment effect and whatever time trend moved everyone. D - C is wave 2's change over the same window, same time trend, no treatment. Subtracting one from the other leaves only the treatment effect.

The counterfactual is what wave 1 would have looked like without the treatment. DiD constructs it by saying: wave 1's counterfactual trajectory = wave 1's pre-period level, carried forward with wave 2's post-period trend. The gap between the actual wave 1 trajectory and that counterfactual is the DiD estimate.

Figure 1: Causal inference with difference-in-differences. Blue solid: Wave 1 actual trajectory. Orange dashed: Wave 2 (control, untreated during this window). Blue dotted: the counterfactual, where Wave 1 would have gone based on Wave 2's post-period trend. The green arrow is the DiD estimate: the gap between the actual Wave 1 trajectory and the counterfactual in the post-treatment period. A, B, C, D correspond to the four cells in the table above.

Before week 20, wave 1 and wave 2 track each other closely. That's the parallel-trends requirement at work. At week 20, wave 1 pulls ahead of both wave 2 and its own counterfactual (the dotted line). That post-treatment divergence is the DiD estimate.

The DiD estimate handles two types of bias at once. Permanent differences between treated and control groups (wave 1 workspaces were always more engaged) cancel out because DiD focuses on changes in outcomes across time periods. Time trends that affect both groups (product improvements, market seasonality) cancel out because both groups experience them.

DiD asks one thing in return: parallel pre-treatment trends. The treated and control groups have to be moving in the same direction at the same rate before treatment starts. When that holds, you can extrapolate the shared trend forward and attribute any post-treatment divergence to the treatment. If the trends were already diverging before treatment, DiD is biased, and no amount of clever regression fixes it.

Parallel trends is the assumption you'll test in step 3.

Companion Notebook

All the code in this tutorial, including the synthetic dataset, the DiD regression, the parallel-trends plot, and the placebo pre-trend test, lives in a single executable Jupyter notebook in the GitHub repo for this series on product experimentation and causal inference for GenAI and LLM applications.

You can clone it, run generate_data.py once, and every output in this article reproduces exactly: github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm

Prerequisites

You'll need Python 3.11 or newer and comfort with pandas and basic regression. You can follow along without prior causal inference experience, as the article defines confounders and selection bias inline when they first appear. You'll encounter clustered standard errors and fixed effects in step 2. The article explains what they do and why they matter, but it doesn't derive them from scratch.

Install the packages for this tutorial:

pip install numpy pandas statsmodels linearmodels matplotlib

Clone the companion repo to get the synthetic dataset:

git clone https://github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm.git
cd product-experimentation-causal-inference-genai-llm
python data/generate_data.py --seed 42 --n-users 50000 --out data/synthetic_llm_logs.csv

Setting Up the Working Example

The dataset simulates a SaaS product with an AI summaries feature launched in two waves: wave 1 workspaces get it at week 20, wave 2 at week 30, with 50,000 users total, each with one row of telemetry.

The data generator bakes in a +5 percentage point causal effect on task completion for users in their workspace's post-treatment period. You know the truth upfront, so you can check whether your DiD estimator actually recovers it.

Load the data and inspect the structure:

import pandas as pd

df = pd.read_csv("data/synthetic_llm_logs.csv")
print(df.shape)
print(df[["wave", "signup_week", "workspace_id", "task_completed"]].head())
print("\nWave sizes:", df.wave.value_counts().to_dict())
print("Treatment weeks per wave:",
      df.groupby("wave").treatment_week.first().to_dict())

Expected output:

(50000, 16)
   wave  signup_week  workspace_id  task_completed
0     2           10            36               0
1     2           51            44               1
2     2            2            28               1
3     1           15            20               1
4     1           29             0               1
Wave sizes: {2: 25063, 1: 24937}
Treatment weeks per wave: {1: 20, 2: 30}

Here's what's happening: you load 50,000 rows, one per user. Wave 1 has about 24,937 users across 25 workspaces; wave 2 has about 25,063 users across 25 different workspaces. The treatment_week column records when each user's workspace got the AI summaries feature (week 20 for wave 1, week 30 for wave 2). The task_completed column is your outcome: did the AI successfully complete the user's task.

One important detail: signup_week in this dataset records which calendar week a user first joined the product, and we're using it as a time index to assign users to pre- or post-treatment cohorts.

A user who signed up in week 22 joined after the feature launched, so their experience is "post-treatment." A user who signed up in week 14 joined before the launch, so their experience is "pre-treatment."

This works here because each user has one row of telemetry tied to their initial product experience. In a panel dataset with multiple observations per user across time, you'd use an observation timestamp column tied to when each row was recorded.

To keep the analysis clean, restrict to users who signed up before the wave 2 launch (signup_week < 30). Wave 2 then works as a proper control group, since it hasn't been treated yet, while wave 1 has been treated for 10 weeks.

analysis = df[df.signup_week < 30].copy()
analysis["post"] = (analysis.signup_week >= 20).astype(int)
analysis["treated"] = (analysis.wave == 1).astype(int)

print(analysis.groupby(["treated", "post"])
              .agg(n=("user_id", "count"),
                   mean_completion=("task_completed", "mean"))
              .round(3))

Expected output:

                 n  mean_completion
treated post
0       0     9590            0.556
        1     4878            0.555
1       0     9633            0.592
        1     4738            0.643

Here's what's happening: you filter the data to the analysis window (weeks 0 to 29) and create two indicator variables. post is 1 for users in the post-week-20 period, 0 otherwise. treated is 1 for wave 1 users, 0 for wave 2. The groupby shows the four cells of the DiD 2x2 table: (treated=0, post=0), (treated=0, post=1), (treated=1, post=0), (treated=1, post=1). Those four means are everything you need for a first-pass DiD estimate.

Step 1: A Simple 2x2 DiD

Start with the cleanest version. Compute the four cell means by hand, then take the difference of differences:

cells = analysis.groupby(["treated", "post"]).task_completed.mean()

wave2_pre  = cells.loc[(0, 0)]   # control, pre
wave2_post = cells.loc[(0, 1)]   # control, post
wave1_pre  = cells.loc[(1, 0)]   # treated, pre
wave1_post = cells.loc[(1, 1)]   # treated, post

did_effect = (wave1_post - wave1_pre) - (wave2_post - wave2_pre)
print(f"Wave 1 change: {wave1_post - wave1_pre:+.4f}")
print(f"Wave 2 change: {wave2_post - wave2_pre:+.4f}")
print(f"DiD effect:    {did_effect:+.4f}")

Expected output:

Wave 1 change: +0.0515
Wave 2 change: -0.0013
DiD effect:    +0.0527  (ground truth = +0.05)

Here's what's happening: you pull the four cell means, compute wave 1's change in task completion from pre to post, compute wave 2's change over the same calendar window (wave 2 hasn't been treated yet), and take the difference. The DiD estimate is the piece of wave 1's change that can't be explained by whatever time trend also moved wave 2.

On this dataset the simple 2x2 estimate lands at +0.053, which is very close to the true +0.05. But you can't take this number to a product review. You have no standard errors, which means you can't say whether +0.053 is a real signal or within sampling noise. You have no covariate adjustment, so if wave 1 happened to have more heavy users in this cohort, some of that +0.053 could be engagement-tier composition. And you have no way to handle the workspace-level correlation in your data. Step 2 fixes all three.

Step 2: Regression DiD with Fixed Effects

The regression formulation of DiD produces the same point estimate as the 2x2 table when there are no covariates. But it also buys you three things:

• Standard errors and p-values computed correctly

• Covariate adjustment to reduce variance and sharpen your estimate

• Cluster-robust errors that handle correlation within workspaces, which a staged rollout always has

The regression is: outcome ~ treated + post + treated:post + controls. The coefficient on the treated:post interaction is your DiD estimate.

import statsmodels.formula.api as smf

did_model = smf.ols(
    "task_completed ~ treated * post + C(engagement_tier)",
    data=analysis
).fit(
    cov_type="cluster",
    cov_kwds={"groups": analysis.workspace_id}
)

print(did_model.summary().tables[1])

Expected output:

================================================================================================
                                   coef    std err          z      P>|z|      [0.025      0.975]
------------------------------------------------------------------------------------------------
Intercept                        0.8301      0.007    126.538      0.000       0.817       0.843
C(engagement_tier)[T.light]     -0.4027      0.006    -63.168      0.000      -0.415      -0.390
C(engagement_tier)[T.medium]    -0.1766      0.007    -25.931      0.000      -0.190      -0.163
treated                          0.0367      0.005      6.885      0.000       0.026       0.047
post                            -0.0056      0.008     -0.684      0.494      -0.022       0.011
treated:post                     0.0541      0.011      4.981      0.000       0.033       0.075
================================================================================================

Here's what's happening: you fit an ordinary least squares regression of task completion on the treated indicator, the post indicator, their interaction, and a categorical control for engagement tier.

The treated:post coefficient is the DiD estimate. Users in the same workspace share common shocks, making their outcomes correlated. Grouping by workspace_id corrects for that.

On this dataset the treated:post coefficient comes out at +0.054 with a clustered p-value of <0.001. The ground truth is +0.050. At 0.4 percentage points from the true effect, with a standard error that accounts for workspace-level correlation, that's a number you can put in a product review.

A few practical notes on this regression:

• Controls should be time-invariant (engagement tier, signup cohort). Time-varying controls that are themselves affected by treatment will bias the estimate.

• Only the interaction has a causal interpretation. The intercept and level terms describe baseline differences between groups, nothing more.

• Clustered errors are mandatory. Skip clustering and your standard errors are 3 to 10x too small, test statistics are artificially inflated, and results look far more significant than they are.

Step 3: Checking the Parallel-Trends Assumption

DiD is only valid if wave 1 and wave 2 were moving in the same direction at the same rate before treatment started. You check this by plotting (or tabulating) weekly means for the two waves across the pre-treatment window.

import matplotlib.pyplot as plt
import numpy as np

df_plot = df[df.signup_week < 30].copy()
weekly = (df_plot.groupby(["signup_week", "wave"])
             .task_completed.mean()
             .reset_index()
             .pivot(index="signup_week", columns="wave", values="task_completed"))

# 3-week rolling average to smooth week-to-week sampling noise
smoothed = weekly.rolling(3, center=True, min_periods=2).mean()

TREATMENT_WEEK = 20
pre_idx = smoothed.index[smoothed.index < TREATMENT_WEEK]
post_idx = smoothed.index[smoothed.index >= TREATMENT_WEEK]

# DiD counterfactual: wave 1 pre-period mean + wave 2's post-period change
wave1_pre_mean = smoothed.loc[pre_idx, 1].mean()
wave2_pre_mean = smoothed.loc[pre_idx, 2].mean()
counterfactual = wave1_pre_mean + (smoothed.loc[post_idx, 2].values - wave2_pre_mean)

fig, ax = plt.subplots(figsize=(10, 5.5))
ax.axvspan(-0.5, TREATMENT_WEEK, alpha=0.04, color="#94A3B8", zorder=0)
ax.axvspan(TREATMENT_WEEK, 29.5, alpha=0.06, color="#3B82F6", zorder=0)
ax.plot(smoothed.index, smoothed[2], "s--", color="#F59E0B", linewidth=2,
        markersize=4, label="Wave 2 — control (untreated during this window)", zorder=3)
ax.plot(smoothed.index, smoothed[1], "o-", color="#2563EB", linewidth=2.2,
        markersize=4, label="Wave 1 — treated (AI feature on at week 20)", zorder=4)
ax.plot(post_idx, counterfactual, ":", color="#2563EB", linewidth=2.2,
        label="Wave 1 counterfactual (projected without treatment)", zorder=4)
ax.axvline(TREATMENT_WEEK, color="#DC2626", linestyle="--", linewidth=1.8,
           label="AI feature launched (week 20)")

ax.text(9.5, 0.508, "Pre-treatment period\n(parallel trends required)",
        fontsize=9, ha="center", color="#64748B", style="italic")
ax.text(24, 0.508, "Post-treatment",
        fontsize=9, ha="center", color="#64748B", style="italic")
ax.set_xlabel("Week", fontsize=11)
ax.set_ylabel("Mean task completion rate", fontsize=11)
ax.set_title("Figure 2: Data-Driven Parallel-Trends Check\n(3-week rolling average, 50k users)",
             fontsize=12, fontweight="bold", pad=14)
ax.legend(loc="upper left", fontsize=9, framealpha=0.92)
ax.set_xlim(-0.5, 29.5)
ax.set_ylim(0.50, 0.72)
ax.grid(True, alpha=0.18, linestyle=":")
ax.tick_params(labelsize=10)
plt.tight_layout()
plt.savefig("parallel_trends.png", dpi=150, bbox_inches="tight")
print("Saved parallel_trends.png")

Expected output (Figure 2, data-driven verification):

Saved parallel_trends.png

Figure 2 is the data-driven parallel-trends check from your actual dataset, plotted as a 3-week rolling average to smooth week-to-week sampling noise. Both waves track each other closely before week 20, and small wiggles in the pre-period affect both groups at the same time, which is exactly what parallel trends looks like. After week 20, wave 1 separates cleanly above the dotted counterfactual line. The gap between the solid blue line and the dotted line in the post-treatment window is the DiD estimate playing out in your actual data.

Here's what's happening: you group by signup week and wave, compute the mean task completion rate per cell, pivot so each wave is a column, and plot the two time series together.

A vertical dashed line marks week 20 when wave 1 got treatment. In the pre-treatment window (weeks 0 to 19) the two series should track each other closely. After week 20, wave 1 should pull ahead of wave 2 by roughly the treatment effect.

To put a number on it, run a placebo regression on the pre-treatment period only. Regress the outcome on a linear time trend interacted with the treated indicator. If the interaction coefficient is near zero and insignificant, the two groups were moving in parallel before treatment:

pre_only = analysis[analysis.post == 0].copy()
pre_only["weeks_since_start"] = pre_only.signup_week - 10  # center

placebo_model = smf.ols(
    "task_completed ~ treated * weeks_since_start + C(engagement_tier)",
    data=pre_only
).fit(
    cov_type="cluster",
    cov_kwds={"groups": pre_only.workspace_id}
)

print("Pre-trend slope difference:",
      placebo_model.params["treated:weeks_since_start"])
print("p-value:",
      placebo_model.pvalues["treated:weeks_since_start"])

Expected output:

Pre-trend slope difference: -0.00095...
p-value: 0.4435...

Here's what's happening: you restrict to pre-treatment observations, fit a regression that lets wave 1 and wave 2 follow different linear trends in the pre-period, and read off the interaction coefficient.

A coefficient close to zero with p > 0.05 means the two waves were moving in parallel before treatment. If that coefficient is large and statistically significant, the parallel-trends assumption is broken: your DiD estimate is absorbing whatever differential trend separated the groups before week 20.

If the placebo test fails, stop and rethink. Your options: restrict to a narrower pre-window where trends were parallel, find a better control group, or switch to synthetic control, which builds a weighted counterfactual from multiple untreated units.

On this synthetic dataset the placebo test passes: the pre-trend slope difference is -0.00095 with p = 0.44, so the parallel-trends assumption holds and the +0.054 estimate from step 2 is trustworthy.

When Difference-in-Differences Fails

DiD is a precise accounting method, and every precise method has specific failure modes worth knowing before you trust its output. Here are four common ones:

1. Non-parallel Pre-trends

When the treated and control groups were already diverging before treatment started, DiD mistakes that pre-existing drift for a treatment effect.

The placebo test in step 3 is your guard. Run it every time. If it fails, you have three options:

1. Restrict the analysis to a shorter pre-window where trends were parallel and re-run the placebo

2. Find a better control group whose pre-trend matches the treated group

3. Switch to synthetic control, which builds a weighted counterfactual from multiple untreated units and picks the weights to match the treated group's pre-treatment trajectory

2. Staggered Adoption

A staged rollout with three or more waves demands a different approach than a clean 2x2. Wave 1 gets treated at week 20, wave 2 at week 30, wave 3 at week 40. Once wave 2 is treated, it's no longer a valid control for wave 1 comparisons that span weeks 30 and beyond. Earlier treated units start acting as controls for later ones, which contaminates the estimate.

That's the Goodman-Bacon decomposition problem, and the standard two-way fixed effects estimator from step 2 will silently absorb it. The Callaway-Sant'Anna estimator (see their 2021 paper) fixes this by averaging only the clean 2x2 comparisons and discarding the contaminated ones. The differences package in Python implements it.

3. Time-varying Confounders that Hit Only the Treated Group

If your marketing team runs a targeted campaign in wave 1 workspaces during week 22, you've got a treatment-specific shock DiD can't net out.

Parallel trends certifies the pre-treatment period, but the post-treatment window remains your responsibility to audit.

Check every product or marketing event inside the analysis window. If you find one, the only options are to redesign the study, restrict the analysis to the window before the shock, or model the shock explicitly as a second treatment variable.

4. Anticipation Effects

If wave 1 customers knew in week 18 that the feature was coming in week 20, some will have started behaving differently before treatment technically started: signing up more, pre-configuring settings, contacting support. That contaminates the "pre" period. The tell is a bump or dip in wave 1 in the weeks immediately before week 20 on the event-study plot.

The fix is to push the pre-period cutoff back. Treat week 18 as the "treatment" start for purposes of the analysis, which removes the anticipation window from your pre-period baseline.

Each of these failure modes has a diagnostic and a specific remedy. Naming them in your analysis builds credibility with skeptical reviewers. DiD is a careful accounting identity – it produces reliable estimates exactly as long as its inputs are clean.

What to Do Next

The regression DiD above is the right tool for a two-wave rollout. If your rollout has three or more waves, switch to the Callaway-Sant'Anna estimator. If your rollout crosses a treatment threshold you set deliberately (confidence scores, query complexity), look into regression discontinuity. If you want to compare a single treated unit against a constructed counterfactual, synthetic control is the right choice.

The companion notebook for this tutorial is here. Clone the repo, generate the synthetic dataset with generate_data.py, and open did_demo.ipynb to reproduce every code block with pre-saved outputs.

If you ship AI features in waves, your rollout calendar is already a DiD study. The only question is whether you run the analysis.