The techniques they name (difference-in-differences, regression discontinuity, and doubly robust estimation, among others) are standard tools.

What's interesting is how those teams operationalized them at scale: where the methods failed in production, what they built around each one to make the estimates trustworthy, and how they connected the numbers to actual product decisions.

If you're building LLM features and making product decisions based on thumbs-up rates and session length, these posts will change how you think about measurement.

Most teams still measure feature impact with 30-day A/B tests and thumbs-up rates. That approach works until you need to know whether the metric moved because of your feature or because of a dozen other things that happened the same week.

The four teams below ran into that problem before most teams were even building with LLMs, and the patterns they settled on are worth understanding before you make the same mistakes. I've watched teams spend weeks shipping a feature, then spend additional weeks arguing about whether the numbers are real. That's avoidable.

For these organizations, causal measurement isn't an afterthought but a foundational element of product experimentation, integrated directly into their deployment architectures. The synthesis presented in this article details a comprehensive toolkit for AI product experiments in which traditional A/B testing is incompatible with the deployment model.

Whether you're managing global model transitions, threshold-based routing, staged rollouts, or observational opt-in data, each scenario necessitates a specific methodological approach. Failing to utilize this toolkit leads to more than just ambiguity. It results in product decisions driven by confounded data, a situation far more damaging than having no measurements at all.

Table of Contents

• Prerequisites

• Why Production AI Measurement is Harder Than it Looks

• Case Study 1: Airbnb's Future Value Framework

  • Short-Term A/B Tests Miss the Behavioral Change That Matters

  • The Framework

  • Reference Implementation

  • Instrumenting for Long-Term Value Cuts Experiments That Look Good in Week 2 and Fail in Month 4

• Case Study 2: Netflix's Quasi-Experiment Taxonomy

  • Deployment Structure Determines the Method

  • Reference Implementation

  • Pick the Wrong Method and Cleaner Data Won't Save You

• Case Study 3: Lyft's Doubly Robust Validation

  • Why Single-Model Approaches Fail in Production

  • Lyft's Production Diagnostics Catch Model Failure Before it Reaches a Decision

  • Reference Implementation

  • Two Hours of Diagnostics Prevent a Quarter of Misdirected Engineering Work

• Case Study 4: Uber's Causal Forecasting Pipeline

  • Merging Causal Estimates with Forecasts

  • Reference Implementation

  • Causal Forecasting in Capacity Planning

• What These Four Teams Have in Common

  • Match the Method to the Deployment Structure

  • Build Diagnostics Before Building Estimators

  • Design Every Causal Estimate Around a Specific Product Decision

  • Document Failure Modes Alongside Every Estimate

• How to Start Applying This in Your Own LLM Stack

  • 1. Instrument Before You Need the Data

  • 2. Classify Your Deployment Mechanisms

  • 3. Run One Diagnostic-Rich Causal Analysis

  • 4. Separate Short-term and Long-term Metrics

  • 5. Make Causal Estimates Forward-Looking

• When Production Causal Pipelines Break

  • Organizational Failures

  • Technical Failures

  • Interpretive Failures

• Bootstrap Confidence Intervals

• Run the Notebook, Then Instrument Your Next Feature

Every code block in this article runs end-to-end in the companion notebook at product-experimentation-causal-inference-genai-llm/tree/main/13_case_studies/. Notebook: case_studies_demo.ipynb.

Prerequisites

You need:

• Python 3.11 or newer

• Comfort with pandas, scikit-learn, and basic regression

• No prior reading on causal inference methods required: each case study explains the technique inline

Install the packages for this article:

pip install numpy pandas scikit-learn scipy matplotlib

Clone the companion repo and generate the shared 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

All four case-study code blocks in this article load that file with pd.read_csv("data/synthetic_llm_logs.csv"). The dataset has 50,000 rows and 16 columns covering user identity, session behavior, and model metadata, including user_id, session_minutes, task_completed, model_used, latency_ms, and query_complexity, among others.

Why Production AI Measurement is Harder Than it Looks

The standard story about measuring the impact of an AI feature goes like this: run an A/B test and report the lift. If your p-value is below 0.05, you ship. But this story breaks in three places.

First, randomization isn't always available. Enterprise SaaS products roll out AI features to workspaces in waves, bypassing the individual user coin flip that A/B testing assumes. Consumer products roll out features gradually by region, by cohort, or by platform. Safety-sensitive features ship to a subset of users whose risk profiles clear a threshold.

When randomization doesn't happen, A/B test logic fails. You can't just run the same analysis on non-randomized data and expect the estimate to mean anything. Confounders that correlate with both who receives the feature and how they behave will bias every coefficient you compute, often in the direction that flatters the feature.

Second, short-term metrics don't always predict long-term value. A prompt change that raises thumbs-up ratings by 8 points today might increase user dependence on the AI assistant in ways that cause churn three months out. A model routing change that improves task completion this week might degrade under a new query distribution emerging next quarter.

I initially presumed that short-term proxies would reliably mirror long-term trends, yet they fail to do so consistently. The limitation of short-term A/B testing lies in its focus on immediate metric shifts while remaining oblivious to downstream user behavioral changes, which are ultimately the most critical factors.

Finally, observational data is unavoidable. A/B testing covers a narrow slice of product decisions. The routing threshold change that shipped six months ago, the model vintage swap in Q3, or the users who opted into agent mode before the gate closed: none of these can be run as experiments after the fact.

For any question that requires looking backward, or any system with routing decisions that can't ethically be randomized, you're working from observational logs, with no experiment design to fall back on.

Observational causal inference isn't a fallback. It's a core competency, and teams that treat it as optional find out the hard way when a stakeholder asks why the numbers from last quarter's rollout don't hold up to scrutiny.

Each of the four teams below built systems that grapple with one or more of these three problems.

Case Study 1: Airbnb's Future Value Framework

Short-Term A/B Tests Miss the Behavioral Change That Matters

Airbnb's engineering team, as described by Jenny Chen in the Airbnb Tech Blog post "How Airbnb Measures Future Value to Standardize Tradeoffs", ran into a fundamental problem with their experiment infrastructure. Standard A/B tests measure outcomes at the end of the experiment window, typically 14 to 30 days.

For marketplace features that affect user behavior over months and years, that window is too short. A feature that moves 30-day bookings upward might be accelerating behavior the user was going to exhibit anyway, pulling forward demand, or genuinely adding new long-term engagement. The 30-day metric can't tell these apart.

The LLM version of this is the assistant dependence problem. A prompt redesign that makes your AI assistant more concise and confident will typically immediately raise thumbs-up ratings and task completion rates. Users prefer confident, direct answers. But if the redesign also makes users less likely to verify answers independently, you may have improved the short-term experience at the cost of calibration and long-term trust.

By the time users start churning because the assistant gave them confident wrong answers twice, the prompt change is long-shipped, and its connection to the churn signal is invisible. I've seen this gap cost teams months of diagnostic work trying to untangle prompt changes from model updates from seasonal behavior.

The Framework

You don't need to wait for long-term outcomes to arrive. You need to have estimated, from prior cohorts, which short-term signals reliably predict long-term retention and revenue. Airbnb's solution converts short-term signals into projected long-term value using a predictive model trained on that historical relationship.

In their context, the metric is a "future value" score that estimates a user's long-term booking contribution based on their current engagement pattern. Once you have that model, you can evaluate any experiment by its expected impact on future value, with the 30-day metric as one of several inputs. The experiment window stays short, and the evaluation horizon extends as far as your predictive model can reach.

The DiD step in the reference implementation requires one identifying assumption: parallel pre-treatment trends. Before the feature shipped, both cohorts must have been on equivalent behavioral trajectories. If wave 1 users were already trending toward higher retention independently of the feature, the DiD estimate mixes the feature effect with a pre-existing difference between the waves. The assumption is that most teams skip validating because it requires plotting pre-period trends, which takes 20 minutes and feels unnecessary until the results don't make sense.

For LLM teams, the equivalent requires two things. First, you need leading indicators of long-term user value: week-7 retention and return query rate. Second, you need historical data linking those leading indicators to long-term outcomes you actually care about (revenue and user lifetime). The linking model is trained once on historical cohorts and then applied to new experiments.

Reference Implementation

The code below shows the structural pattern: compute a future-value proxy for each user from short-term signals, then use it as the outcome in a DiD or IPW analysis, replacing the immediate task-completion signal.

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

# Synthetic LLM telemetry with retention signal
df = pd.read_csv("data/synthetic_llm_logs.csv")

# Step 1: Train the future-value proxy model on a historical cohort.
# In production this model is trained on users old enough that
# their long-term outcome (e.g., 90-day retained revenue) is known.
historical = df[df.signup_week < 10].copy()

feature_cols = ["task_completed", "thumbs_up", "session_minutes"]
X_hist = historical[feature_cols].fillna(0)
y_hist = historical["retained_7d"].values  # 7-day retention as long-term proxy

fv_model = LinearRegression().fit(X_hist, y_hist)
# R² computed on training data; use a holdout cohort in production
print("Future-value model R²:", round(fv_model.score(X_hist, y_hist), 3))

# Step 2: Score all users with the future-value proxy.
X_all = df[feature_cols].fillna(0)
df["future_value_score"] = fv_model.predict(X_all)

# Step 3: Compare future_value_score by wave (this is the real experiment outcome).
print("\nMean future-value score by wave:")
print(df.groupby("wave").future_value_score.mean().round(4))

# Step 4: The DiD effect on future value (rather than on task_completed).
# This is where you would plug future_value_score into your DiD regression.
analysis = df[df.signup_week < 30].copy()
analysis["post"] = (analysis.signup_week >= 20).astype(int)
analysis["treated"] = (analysis.wave == 1).astype(int)

cells = analysis.groupby(["treated", "post"]).future_value_score.mean()
did_fv = (
    (cells.loc[(1, 1)] - cells.loc[(1, 0)])
    - (cells.loc[(0, 1)] - cells.loc[(0, 0)])
)
print(f"\nDiD effect on future-value score: {did_fv:+.4f}")

Expected output:

Future-value model R²: 0.024

Mean future-value score by wave:
wave
1    0.6325
2    0.6271
Name: future_value_score, dtype: float64

DiD effect on future-value score: +0.0059

Here's what's happening: you train a lightweight linear model on a historical cohort where long-term outcomes are already known, mapping observable short-term signals to 7-day retention as a proxy for future value.

You score all users with that model, then use the future value score as the outcome in a standard DiD. Seven-day retention is an imperfect proxy, but it forces the analysis to weight short-term engagement by its historical correlation with durable value, which is more than thumbs-up rate does.

The low R² value of 0.024 is intentional, as it highlights the inherent noise when linking immediate session data to 7-day retention. While production systems should ideally utilize signals with higher predictive power such as return-visit rates or query depth, even a less precise linking model can still provide value.

The primary objective is to establish the correct direction of the correction rather than achieve absolute precision.

Instrumenting for Long-Term Value Cuts Experiments That Look Good in Week 2 and Fail in Month 4

The Airbnb framework is a direct response to the measurement horizon problem. When you evaluate AI features on 30-day or 14-day windows, you reward features that move users fast, regardless of where they're moving.

Instrumenting for leading indicators of long-term value doesn't require a longer experiment. It requires a richer measurement model. Teams that have built this capability run fewer experiments that look great in week 2 and disappoint in month 4.

If a linking model isn't yet part of your infrastructure, developing one should be your immediate priority over expanding your evaluation dashboards.

Case Study 2: Netflix's Quasi-Experiment Taxonomy

Deployment Structure Determines the Method

The Netflix Technology Blog post "Key Challenges with Quasi Experiments at Netflix" is one of the more practically useful pieces on causal inference for product teams. Its core contribution is a taxonomy: for each deployment scenario, there's a corresponding causal method, and the post names the identifying assumption and failure mode that go with it.

That framing matters because most teams don't pick methods based on deployment structure. They pick what they already know, which is often the wrong fit.

Figure 1: Deployment structure determines which identification strategy is credible. Threshold routing systems call for RDD, while opt-in analyses call for propensity methods. The assignment mechanism drives the choice, with the team's preferred estimator coming second.

Netflix's taxonomy covers four scenarios that map almost exactly to the situations LLM teams encounter:

Staged rollouts (their scenario: gradual market entry) map to difference-in-differences. When you ship an AI feature to workspace cohort A before cohort B, you've got a natural treated and control group across time. The identification strategy subtracts the shared time trend from the difference in outcomes.

The critical assumption is that the two cohorts have parallel pre-treatment trends. If one cohort was already trending up before treatment started, the method can't distinguish that from a real effect.

Threshold-based routing (their scenario: geographic score cutoffs) maps to regression discontinuity. When a continuous score determines which model or feature a user receives, users just below and just above the threshold are nearly identical in everything except the treatment.

The jump at the cutoff identifies the local average treatment effect (LATE): the causal effect for users near the threshold only, with the average treatment effect across all users outside its scope. The critical assumption is that users can't precisely manipulate the score.

Full-population upgrades (their scenario: platform-wide policy changes) map to the synthetic control design. When every user gets the new model at once, and there's no holdout group, you construct a weighted combination of historical or synthetic counterfactuals to estimate what would have happened without the upgrade.

The critical assumption is that the synthetic control fits the pre-treatment period well. Poor pre-period fit isn't a minor inconvenience. It invalidates the entire counterfactual.

Matched comparisons (their scenario: opt-in feature adoption) map to propensity score methods. When users self-select into AI features, you reweight or re-match the comparison group to approximate random assignment on observables.

The critical assumption is that all relevant confounders are observed. If users who opt in also tend to be power users in ways you haven't measured, your confounder adjustment is incomplete, and your estimate is biased in ways that are hard to detect after the fact.

The taxonomy makes method selection a structured lookup: describe your deployment structure, and find the method whose assumptions your setup most plausibly satisfies.

I've seen teams skip this step and spend two weeks running a DiD on data that was clearly a threshold routing problem. The estimates differed by 40%. Neither was wrong. They were answering different questions.

Reference Implementation

The code below implements the taxonomy as a decision function: given a deployment scenario description, print the appropriate method and its key assumption.

TAXONOMY = {
    "staged_rollout": {
        "method": "Difference-in-Differences (DiD)",
        "assumption": "Parallel pre-treatment trends between treated and control cohorts",
        "check": "Plot weekly means by cohort before treatment starts; "
                 "run pre-trend placebo regression",
        "failure_mode": "Non-parallel pre-trends, time-varying confounders, "
                        "staggered adoption without Callaway-Sant'Anna correction",
    },
    "threshold_routing": {
        "method": "Regression Discontinuity Design (RDD)",
        "assumption": "Users cannot precisely manipulate their score across the cutoff",
        "check": "McCrary density test; bandwidth sensitivity; "
                 "quadratic spec robustness",
        "failure_mode": "Score manipulation, other policies firing at same cutoff, "
                        "extrapolation bias away from the cutoff",
    },
    "full_population_upgrade": {
        "method": "Synthetic Control",
        "assumption": "Pre-treatment fit between actual and synthetic counterfactual is good",
        "check": "In-time placebo tests; in-space placebo tests; "
                 "plot pre-period fit",
        "failure_mode": "Poor pre-period fit, interference between donor units, "
                        "post-treatment structural breaks",
    },
    "opt_in_feature": {
        "method": "Propensity Score Methods (IPW / Matching)",
        "assumption": "All confounders that drive opt-in and affect outcome are observed",
        "check": "Standardized mean difference before and after weighting; "
                 "propensity overlap histogram",
        "failure_mode": "Unmeasured confounders, positivity violations, "
                        "propensity model misspecification",
    },
}

def select_method(scenario: str) -> None:
    if scenario not in TAXONOMY:
        valid = ", ".join(TAXONOMY.keys())
        print(f"Unknown scenario. Valid options: {valid}")
        return
    entry = TAXONOMY[scenario]
    print(f"Scenario:      {scenario}")
    print(f"Method:        {entry['method']}")
    print(f"Assumption:    {entry['assumption']}")
    print(f"Key checks:    {entry['check']}")
    print(f"Failure modes: {entry['failure_mode']}")

# Example: staged AI feature rollout across enterprise workspaces
select_method("staged_rollout")
print()
# Example: confidence-threshold routing between model tiers
select_method("threshold_routing")

Expected output:

Scenario:      staged_rollout
Method:        Difference-in-Differences (DiD)
Assumption:    Parallel pre-treatment trends between treated and control cohorts
Key checks:    Plot weekly means by cohort before treatment starts; run pre-trend placebo regression
Failure modes: Non-parallel pre-trends, time-varying confounders, staggered adoption without Callaway-Sant'Anna correction

Scenario:      threshold_routing
Method:        Regression Discontinuity Design (RDD)
Assumption:    Users cannot precisely manipulate their score across the cutoff
Key checks:    McCrary density test; bandwidth sensitivity; quadratic spec robustness
Failure modes: Score manipulation, other policies firing at same cutoff, extrapolation bias away from the cutoff

Each deployment scenario has a corresponding method, a main identifying assumption, the diagnostics that check whether the assumption holds, and the failure modes that invalidate the analysis.

The function is a decision aid that makes the method-selection step explicit, so the team agrees on the identification strategy before writing a single line of regression code. Without that agreement, you'll often discover mid-analysis that two people on the team were implicitly running different causal models on the same data.

Pick the Wrong Method and Cleaner Data Won't Save You

Most teams pick the causal method they know best. That's the wrong heuristic, and the Netflix taxonomy exists precisely to short-circuit it.

An LLM team with DiD experience will reach for DiD even when they're running a threshold routing system where RDD would give a cleaner answer, and a defensible local treatment effect estimate rather than an averaged-out guess.

The taxonomy highlights a vital principle: the method of selection is determined by the assignment mechanism itself, rather than by the team's familiarity. If your assignment mechanism is a cutoff score, RDD is the first tool to try, regardless of what the team already knows how to run.

Getting this wrong doesn't just produce a noisier estimate. It produces a structurally invalid one that cleaner data won't fix.

Case Study 3: Lyft's Doubly Robust Validation

Why Single-Model Approaches Fail in Production

Shima Nassiri's post on the Lyft Engineering blog, "Trusting the Untestable: Validation and Diagnostics for Doubly Robust Models", starts from a practical observation: in most real production causal analyses, at least one of your nuisance models carries specification error.

When you run an observational causal analysis, you're almost always fitting two models: a propensity model (predicting treatment from covariates) and an outcome model (predicting the outcome from treatment and covariates).

Both models are approximations of unknown true functions. If either one is wrong in ways you haven't accounted for, your causal estimate is biased, and you won't know it from the standard output alone.

Doubly robust estimation, specifically the augmented inverse probability weighting estimator (AIPW), is the response to this. AIPW combines propensity weighting with regression adjustment: if either the propensity model or the outcome model is correctly specified, the AIPW estimate is consistent. One well-specified model is enough.

That said, AIPW offers no protection against unmeasured confounders, and it still requires unconfoundedness: all factors that affect both treatment assignment and the outcome must be observed and included in the model. If a key confounder isn't in your data, AIPW can't save you.

Nassiri's post goes further than the estimator itself. What makes it practically important is the diagnostic toolkit it describes for validating observational analyses before you act on them.

In a clean randomized experiment, you check balance and run power calculations. In an observational study, you have to work harder, because the design carries no randomization guarantee. I've seen teams skip this diagnostic step and then spend weeks explaining why their causal estimate was off by a factor of two.

Lyft's Production Diagnostics Catch Model Failure Before it Reaches a Decision

The pipeline runs four checks:

1. Weight distribution check

After fitting the propensity model, plot the distribution of IPW weights. Extreme weights, say, above 20 or 30, signal that some users have near-zero propensity, which violates the positivity assumption: every unit must have nonzero probability of both treatment and control assignment.

Those users lack a comparable counterfactual, and letting a single unusual observation dominate your causal conclusion undermines the analysis. Skipping this check is how a single power user with unusual behavior skews an ATE by 15 percentage points.

2. Trim threshold

Set a maximum weight. Any observation whose weight exceeds the trim threshold is downweighted to the threshold value. Common choices are the 95th or 99th percentile of the weight distribution.

Trimming trades a small amount of bias for a large reduction in variance, making the estimate more stable under minor model misspecification. If you don't trim, you're letting the weirdest edge cases in your data drive the headline number.

3. Covariate balance plots

Plot standardized mean differences before and after weighting for every covariate in the propensity model. The target is |SMD| < 0.1 after weighting.

Covariates still above that threshold after weighting indicate that the propensity model is missing that covariate's influence on treatment assignment. This is the check that catches the "but we adjusted for everything" blind spot.

4. Placebo outcome test

Take an outcome that your treatment provably doesn't cause, for example, a pre-treatment metric from before the treatment existed, and run the full AIPW pipeline on it.

If the pipeline returns a significant effect on the placebo outcome, you have a problem: unmeasured confounders, a misspecified propensity model, or data leakage. A placebo failure is one of the clearest signals that your analysis isn't credible, and it's a signal you can get before you ship anything.

Reference Implementation

The code below shows the weight distribution check and trimming step that Lyft's pipeline applies before trusting any causal estimate.

import pandas as pd
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression

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

# Estimate propensity for opt-in to agent mode
X = pd.get_dummies(
    df[["engagement_tier", "query_confidence"]], drop_first=True
).astype(float)
y = df["opt_in_agent_mode"]

ps_model = LogisticRegression(max_iter=1000).fit(X, y)
df["propensity"] = ps_model.predict_proba(X)[:, 1]

# ATE weights: 1/P(treat) for treated, 1/(1-P) for control
df["ipw"] = np.where(
    df.opt_in_agent_mode == 1,
    1 / df.propensity,
    1 / (1 - df.propensity),
)

# Diagnostic 1: weight distribution
print("IPW weight percentiles:")
for p in [50, 75, 90, 95, 99]:
    print(f"  {p}th pct: {np.percentile(df.ipw, p):.2f}")

fig, ax = plt.subplots(figsize=(8, 4))
ax.hist(df.ipw, bins=60, edgecolor="none", alpha=0.7)
ax.axvline(np.percentile(df.ipw, 99), color="red", linestyle="--",
           label="99th pct (trim threshold)")
ax.set_xlabel("IPW weight")
ax.set_ylabel("Count")
ax.set_title("Weight distribution: check for extreme values")
ax.legend()
plt.tight_layout()
plt.savefig("weight_distribution.png", dpi=140)
print("Saved weight_distribution.png")

# Diagnostic 2: trim extreme weights at 99th percentile
trim_threshold = np.percentile(df.ipw, 99)
df["ipw_trimmed"] = df.ipw.clip(upper=trim_threshold)

# Compare ATE before and after trimming
def weighted_ate(data):
    t = data[data.opt_in_agent_mode == 1]
    c = data[data.opt_in_agent_mode == 0]
    return (
        (t.task_completed * t.ipw_trimmed).sum() / t.ipw_trimmed.sum()
        - (c.task_completed * c.ipw_trimmed).sum() / c.ipw_trimmed.sum()
    )

# Untrimmed ATE using ipw column
df["ipw_trimmed_orig"] = df["ipw"].copy()   # backup before overwrite
ate_untrimmed = (
    (df[df.opt_in_agent_mode==1].task_completed * df[df.opt_in_agent_mode==1].ipw).sum()
    / df[df.opt_in_agent_mode==1].ipw.sum()
    - (df[df.opt_in_agent_mode==0].task_completed * df[df.opt_in_agent_mode==0].ipw).sum()
    / df[df.opt_in_agent_mode==0].ipw.sum()
)
ate_trimmed = weighted_ate(df)
print(f"\nATE (untrimmed): {ate_untrimmed:+.4f}")
print(f"ATE (trimmed):   {ate_trimmed:+.4f}")
print(f"Trim threshold:  {trim_threshold:.2f}")

Expected output:

IPW weight percentiles:
  50th pct: 1.52
  75th pct: 1.57
  90th pct: 2.88
  95th pct: 8.14
  99th pct: 8.58
Saved weight_distribution.png

ATE (untrimmed): +0.0851
ATE (trimmed):   +0.0852
Trim threshold:  8.58

Figure 2: IPW weight distribution on the 50,000-user synthetic dataset. The bulk of the weights cluster between 1.0 and 3.0. 500 observations exceed the 99th-percentile trim threshold of 8.58. Trimming shifts the ATE by 0.0001, confirming extreme weights carry negligible influence on this estimate. Unlike Figure 1's conceptual map, this diagnostic runs directly on real data from the shared dataset.

Here's what's happening: you fit a propensity model, compute ATE weights, then plot the weight histogram to see whether any users have extreme weights that dominate the estimate.

The 99th percentile line is the visual trim threshold. You apply the trim and compare the untrimmed vs. trimmed ATE. If they're close, the extreme weights had minimal influence on the result. If they're far apart, you have a small cluster of influential observations, and the trimmed estimate is more trustworthy.

Two Hours of Diagnostics Prevent a Quarter of Misdirected Engineering Work

When you're measuring the causal effect of an AI feature from observational logs, you're almost always in the regime where both your propensity model and your outcome model carry error. The AIPW structure gives you protection against one of them being wrong. The Lyft diagnostic toolkit tells you how much each model is carrying before you act on the estimate.

Running the weight diagnostic and the placebo test may add about 2 hours to a causal analysis. That two hours can prevent the kind of confident-but-wrong conclusion that sends an engineering team chasing the wrong feature for a quarter, and I've watched that happen. The cost of skipping diagnostics isn't abstract: it's six engineers working on something that wasn't the cause of the outcome you were measuring.

Case Study 4: Uber's Causal Forecasting Pipeline

Merging Causal Estimates with Forecasts

The standard output of a causal analysis is a point estimate and a confidence interval: the AI feature raised task completion by 6 percentage points, 95% CI [3.8, 8.2]. That number answers a backward-looking question: what happened?

Product decisions are forward-looking. If you're considering raising the model routing threshold from 0.85 to 0.90, you want to know what the cost and quality tradeoffs will look like next quarter, a projection forward grounded in what you learned from last month's experiment.

Totte Harinen and Bonnie Li's post "Using Causal Inference to Improve the Uber User Experience" on the Uber Engineering blog describes how Uber applies causal inference to production decisions, providing the foundation for embedding causal effect estimates into forward-looking scenario models.

The structural move is to treat the causal estimate as a parameter in the forecast. Forecasting cost and quality separately and assuming a stable relationship between them leaves the causal parameter unspecified. The structural move is to model the causal effect of the routing threshold on the cost-quality tradeoff directly, then project that parameter forward under different assumptions about query volume, query distribution, and model capability.

This matters specifically for LLM systems because the relationship between routing decisions and costs is nonlinear and distribution-dependent. A routing threshold that's cost-efficient at your current query volume may break down at 3x volume. A model you optimized routing for in Q1 may be replaced by a cheaper model in Q3, shifting the cost-quality Pareto frontier entirely. Embedding causal estimates into the forecast makes those structural changes visible before they arrive.

Reference Implementation

The local comparison near the routing threshold rests on two identifying assumptions. First, engineers and users can't precisely manipulate query_confidence to cluster on one side of the 0.85 cutoff. Assignment must be as-good-as-random within a narrow band around the threshold.

Second, the potential outcome functions must be continuous across the cutoff, so the jump observed at 0.85 is attributable to routing assignment and not to any other policy firing at the same score level.

The code below illustrates the pattern: estimate the causal effect of a change in routing threshold on cost and quality, then project that effect across a range of future volume scenarios.

import pandas as pd
import numpy as np

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

# Step 1: Estimate causal effect of premium routing on quality and cost
# (Using RDD logic: compare users near the routing threshold)
cutoff = 0.85
bw = 0.10
near = df[
    (df.query_confidence > cutoff - bw)
    & (df.query_confidence < cutoff + bw)
].copy()
# Low-confidence queries route to premium model (below-threshold queries need stronger handling)
near["routed_premium"] = (near.query_confidence < cutoff).astype(int)

# Causal effects from the local comparison near the threshold
quality_effect = (
    near[near.routed_premium == 1].task_completed.mean()
    - near[near.routed_premium == 0].task_completed.mean()
)
cost_effect = (
    near[near.routed_premium == 1].cost_usd.mean()
    - near[near.routed_premium == 0].cost_usd.mean()
)

print(f"Estimated quality effect of premium routing: {quality_effect:+.4f}")
print(f"Estimated cost effect of premium routing:    {cost_effect:+.4f}")

# Step 2: Embed into forward-looking scenarios
# Suppose we're evaluating: what if we raise threshold from 0.85 to 0.90?
# Queries with confidence 0.85 to 0.90 would shift from premium to cheap routing.
threshold_change_users = df[
    (df.query_confidence >= 0.85) & (df.query_confidence < 0.90)
]
n_shifted = len(threshold_change_users)
print(f"\nQueries that would shift at threshold 0.85 to 0.90: {n_shifted}")

# Volume scenarios (monthly queries)
monthly_query_volume = [500_000, 1_000_000, 2_000_000]
shifted_fraction = n_shifted / len(df)  # fraction of total traffic shifted

print("\nForward-looking scenario: raise threshold from 0.85 to 0.90")
print(f"{'Monthly volume':>20} {'Quality change':>16} {'Cost change ($/mo)':>20}")
for vol in monthly_query_volume:
    n_affected = vol * shifted_fraction
    delta_quality = quality_effect * n_affected / vol    # rate change in overall quality
    delta_cost = -cost_effect * n_affected               # negative: saving cost by de-premiuming
    print(f"{vol:>20,.0f} {delta_quality:>+16.4f} {delta_cost:>+20,.0f}")

Expected output:

Estimated quality effect of premium routing: +0.0613
Estimated cost effect of premium routing:    +0.0080

Queries that would shift at threshold 0.85 to 0.90: 5415

Forward-looking scenario: raise threshold from 0.85 to 0.90
      Monthly volume   Quality change   Cost change ($/mo)
             500,000          +0.0066                 -436
           1,000,000          +0.0066                 -871
           2,000,000          +0.0066               -1,742

Here's what's happening: you estimate the causal effect of premium routing on quality (task completion) and cost using a local comparison near the routing threshold. You then identify the fraction of queries that would shift routing assignment if you moved the threshold from 0.85 to 0.90.

Finally, you project the quality and cost implications of that shift across different monthly query volume scenarios. The output is a scenario table that a product or finance team can read directly: raising the threshold saves roughly $X per month at current volume and costs approximately Y percentage points of task completion rate.

Causal Forecasting in Capacity Planning

The causal forecasting pattern is most useful for routing and infrastructure decisions where cost and quality effects are both significant, and you need to make choices ahead of traffic scale you haven't reached yet. Running the causal estimate forward into volume scenarios turns a retrospective finding into an actionable projection.

Skip this step and causal estimates stay buried in analysis documents, disconnected from capacity planning and pricing decisions. I've watched useful analyses go unread for this exact reason. With it, the measurement team is producing inputs that actually matter to how the product is run.

What These Four Teams Have in Common

These four teams built different methods but converged on the same operational discipline.

Match the Method to the Deployment Structure

Start from the assignment mechanism (how was treatment assigned?) and work backward to the identification strategy. Airbnb moved past short-term A/B tests because their features affect long-term value beyond a 30-day window. Netflix uses RDD for threshold routing systems because the cutoff is the natural identification strategy.

Pick the technique because your system's design makes a particular identification strategy credible. Defaulting to the method the team knows best is how identification errors happen, and those errors don't announce themselves.

Build Diagnostics Before Building Estimators

Run the assumption checks before reporting the estimate. Airbnb validates the leading-indicator model on historical cohorts. Lyft runs weight distributions and placebo tests before acting on an observational estimate.

An estimate reported without its diagnostic layer is an estimate you can't defend. That distinction matters when the product team challenges your number at the quarterly review.

Design Every Causal Estimate Around a Specific Product Decision

Airbnb estimates long-term value to inform feature-shipping decisions. Netflix runs quasi-experiments to make rollout decisions.

Analyses that don't improve any specific product choice aren't worth running: they consume analyst time, create misleading signals in the reporting backlog, and erode stakeholder trust in the measurement function over time.

Document Failure Modes Alongside Every Estimate

Each technique has a named list of ways it can break: non-parallel trends for DiD, manipulation at the cutoff for RDD, unmeasured confounders for propensity methods, and poor synthetic control fit for full-population upgrades.

Ship the estimate alongside its failure conditions labeled. The credibility of an analysis for a skeptical audience stems not from the confidence interval itself, but from a transparent disclosure of the specific assumptions that would need to be invalidated for the estimate to fail.

How to Start Applying This in Your Own LLM Stack

Most LLM teams aren't starting from a mature causal pipeline. The steps below are ordered by impact.

1. Instrument Before You Need the Data

The biggest constraint in every observational causal analysis is that the data you needed wasn't collected. Before you can run a DiD on a staged rollout, you need pre-treatment data for both cohorts.

Before you can run an AIPW on an opt-in feature, you need a rich set of covariates that predict opt-in.

Instrument your system now for the analyses you'll want to run in six months: session length, query complexity, 7-day return rate, and model routing decisions. The instrument is cheap, but retroactive data collection is impossible.

2. Classify Your Deployment Mechanisms

Apply the Netflix taxonomy to every AI feature currently running in your product. For each feature, ask: how was treatment assigned? Which causal method does that assignment mechanism support?

What's the core assumption, and do you have the data to check it? The exercise usually reveals that most features are being measured with tools that don't match their assignment mechanism. That mismatch isn't academic. It means you don't know whether those features are working.

3. Run One Diagnostic-Rich Causal Analysis

Pick one feature, run balance checks and placebo tests, stress-test sensitivity to specification choices, and write up the results. The discipline of running every check once establishes the pattern for future analyses.

It also usually surfaces one uncomfortable finding about the feature you were most confident in. I've seen this happen on three separate teams: the "obviously working" feature turns out to have a confounded comparison group.

4. Separate Short-term and Long-term Metrics

Follow Airbnb's lead and identify at least one leading indicator of long-term value that you can measure in a 30-day experiment window. Seven-day retention, return query rate in week 3, or escalation rate trajectory are all candidates.

Report this alongside immediate engagement metrics in every experiment summary. Without it, you're optimizing a proxy and discovering the gap in the next quarter's retention numbers.

5. Make Causal Estimates Forward-Looking

When you produce a causal estimate, add one row: "Under 3x current volume, this effect implies X." That translation step forces the analysis to make contact with infrastructure and product planning, and it changes who reads it.

When Production Causal Pipelines Break

Production causal pipelines break in a few predictable places.

Organizational Failures

First, no one owns the measurement design. In most teams, the data scientist writes the analysis after the feature ships. Because that's the standard workflow, you're always running retrospective analyses on data that wasn't designed for causal identification.

The fix is a measurement design review before features ship: who's the control group, how long is the pre-period, what's the core assumption, and what diagnostic will falsify it? A 30-minute review prevents a common class of unrecoverable analyses.

Second, causal results don't reach decision-makers. A correct causal estimate that doesn't inform a product decision is a failed analysis, even if the statistics are right. You can't fix that with a better methodology. Causal pipelines need fast-path reporting alongside rigorous reporting.

Technical Failures

First, instrumentation gaps are discovered after the fact. The most common technical failure is the need for a covariate that wasn't logged. You discover the gap when you try to check balance or run a propensity model, three weeks after the experiment ended.

The instrument-early principle above addresses this, but it requires buy-in from the infrastructure team to prioritize event logging that serves causal analysis as directly as it serves product dashboards. That buy-in is harder to get than the logging itself.

Second, there's treatment leakage in the synthetic dataset. For teams testing causal methods on synthetic or internal data, the data generation process can inadvertently bake in the causal effect you're trying to estimate, making any method appear to work.

Validate your analysis on external holdout data or on cohorts outside the generation window. This one is easy to miss because the synthetic data looks clean. Structural contamination within data rows can be subtle and difficult to detect.

Interpretive Failures

First, conflating LATE with ATE. RDD estimates the local average treatment effect (LATE): the effect at the cutoff, for the specific users near the threshold. Propensity matching estimates ATT: the effect for users who were treated. The ATE for the full population requires a different approach.

When a PM asks "what's the effect of this feature," they usually mean ATE. When your causal analysis gives them LATE without explaining the difference, they'll apply the estimate to decisions it wasn't designed to support, and the resulting product choice will be wrong in ways you can't trace back to the analysis.

Second, external validity assumptions that don't hold. A causal estimate from last quarter's user population may not generalize to next quarter's, particularly when you're scaling into new segments or entering an international market.

The estimated effect on power users who opted in early, as the feature rolls out to light-engagement users. Document the population your estimate applies to. Flag explicitly when it's about to be applied outside that population.

Third, reporting precision that overstates certainty. A causal estimate with two-decimal precision reported from an observational study with residual confounding risk conveys more certainty than the analysis warrants.

Report confidence intervals alongside point estimates, the assumptions the estimates depend on, and the balance after weighting, all in the summary where decision-makers will actually see them. The analysis isn't done until the uncertainty is visible to the people acting on it.

Bootstrap Confidence Intervals

Point estimates from observational analyses carry sampling uncertainty. The bootstrap below (500 replicates, seed=7) provides 95% confidence intervals for the three numerical estimates in this article: the Airbnb DiD effect on future-value score, the Lyft IPW ATE, and the Uber RDD quality effect.

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

rng = np.random.default_rng(7)
df = pd.read_csv("data/synthetic_llm_logs.csv")
n_boot = 500

# Bootstrap 1: DiD on future-value score (Airbnb)
historical = df[df.signup_week < 10].copy()
feature_cols = ["task_completed", "thumbs_up", "session_minutes"]
fv_model = LinearRegression().fit(historical[feature_cols].fillna(0), historical["retained_7d"].values)
df["future_value_score"] = fv_model.predict(df[feature_cols].fillna(0))
analysis = df[df.signup_week < 30].copy()
analysis["post"] = (analysis.signup_week >= 20).astype(int)
analysis["treated"] = (analysis.wave == 1).astype(int)

did_boots = []
for _ in range(n_boot):
    s = analysis.sample(frac=1, replace=True, random_state=rng.integers(1e9))
    c = s.groupby(["treated", "post"]).future_value_score.mean()
    try:
        did_boots.append((c.loc[(1, 1)] - c.loc[(1, 0)]) - (c.loc[(0, 1)] - c.loc[(0, 0)]))
    except KeyError:
        pass
ci_did = np.percentile(did_boots, [2.5, 97.5])
print(f"DiD future-value 95% CI: [{ci_did[0]:+.4f}, {ci_did[1]:+.4f}]")

# Bootstrap 2: IPW ATE trimmed (Lyft)
X = pd.get_dummies(df[["engagement_tier", "query_confidence"]], drop_first=True).astype(float)
ps_model = LogisticRegression(max_iter=1000).fit(X, df["opt_in_agent_mode"])
df["propensity"] = ps_model.predict_proba(X)[:, 1]
df["ipw"] = np.where(df.opt_in_agent_mode == 1, 1 / df.propensity, 1 / (1 - df.propensity))
trim_thr = np.percentile(df.ipw, 99)
df["ipw_trimmed"] = df.ipw.clip(upper=trim_thr)

ate_boots = []
for _ in range(n_boot):
    s = df.sample(frac=1, replace=True, random_state=rng.integers(1e9))
    t = s[s.opt_in_agent_mode == 1]
    c = s[s.opt_in_agent_mode == 0]
    ate_boots.append(
        (t.task_completed * t.ipw_trimmed).sum() / t.ipw_trimmed.sum()
        - (c.task_completed * c.ipw_trimmed).sum() / c.ipw_trimmed.sum()
    )
ci_ate = np.percentile(ate_boots, [2.5, 97.5])
print(f"IPW ATE trimmed 95% CI:  [{ci_ate[0]:+.4f}, {ci_ate[1]:+.4f}]")

# Bootstrap 3: RDD quality effect near routing cutoff (Uber)
cutoff = 0.85
bw = 0.10
near = df[(df.query_confidence > cutoff - bw) & (df.query_confidence < cutoff + bw)].copy()
near["routed_premium"] = (near.query_confidence < cutoff).astype(int)

qe_boots = []
for _ in range(n_boot):
    s = near.sample(frac=1, replace=True, random_state=rng.integers(1e9))
    qe_boots.append(
        s[s.routed_premium == 1].task_completed.mean()
        - s[s.routed_premium == 0].task_completed.mean()
    )
ci_qe = np.percentile(qe_boots, [2.5, 97.5])
print(f"RDD quality effect 95% CI: [{ci_qe[0]:+.4f}, {ci_qe[1]:+.4f}]")

Expected output:

DiD future-value 95% CI: [+0.0023, +0.0093]
IPW ATE trimmed 95% CI:  [+0.0727, +0.0966]
RDD quality effect 95% CI: [+0.0490, +0.0748]

Here's what's happening: three separate bootstrap loops resample the analysis dataset 500 times each with a shared seed.

The DiD bootstrap resamples the full analysis cohort and recomputes the 2x2 cell means. The interval [+0.0023, +0.0093] confirms the future-value effect is statistically distinguishable from zero.

The IPW ATE bootstrap resamples all 50,000 users and reweights each draw. The interval [+0.0727, +0.0966] covers the ground-truth +0.08 opt-in effect and excludes zero.

The RDD bootstrap resamples only users within the bandwidth window near the 0.85 cutoff. The interval [+0.0490, +0.0748] confirms the local quality effect is nonzero.

All three intervals are tight enough to be actionable and wide enough to reflect the uncertainty of observational estimates. If you're reporting a point estimate without one of these intervals, you're understating the risk your stakeholders are absorbing.

Run the Notebook, Then Instrument Your Next Feature

The companion notebook for this article lives at github.com/RudrenduPaul/product-experimentation-causal-inference-genai-llm/tree/main/13_case_studies/. Clone the repo, generate the synthetic dataset using the Prerequisites commands above, and run case_studies_demo.ipynb to reproduce every code block from this article, including all four case-study implementations and the bootstrap validation. It also contains a decision function that extends the Netflix taxonomy into a more complete method-selection guide.

The source material for the four case studies is available directly from each team's engineering blog.

1. Jenny Chen's future value post is at (Airbnb Tech Blog).

2. The quasi-experiment taxonomy is at (Netflix Technology Blog).

3. Nassiri's doubly robust validation piece is at (Lyft Engineering).

4. Harinen and Li's causal inference overview is at (Uber Engineering).

Reading the originals is worthwhile: they describe production systems in detail that a summary can't fully capture.

The teams that reliably measure AI impact share one practice: matching the method to the assignment mechanism, running diagnostics before trusting estimates, and connecting causal results to decisions before the decision window closes.

The bottleneck is almost always instrumentation. The data those analyses depend on has to exist before the feature ships. That's the gap the frameworks above can't close for you, and the reason the instrument-early step comes first.

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.

In this tutorial we’ll profile the dataset first, then let those observations drive every design decision in a MONAI segmentation pipeline.

What We'll Cover:

• Who is This For?

• About the Dataset

• What is MONAI, and Why Use it?

• What is Dice?

• Part 1 — Data Profile Before Modeling

  • Class Balance Drives the Loss

  • Patient Counts Drive the Split

• Part 2 — Building the Pipeline

  • A Single Config Object

  • The Patient-grouped Split

  • Transforms, Chosen by the Snapshot

  • Model, Loss, and Metric

• Reading the Results

• Prediction Visualization

• The Failure Modes Matter More Than the Average

• Where to Go Next

• Takeaway

• Reference

Who is This For?

This walkthrough assumes you have some comfort with Python and the basics of training a neural network. It explains the MONAI-specific pieces (dictionary transforms, DiceCELoss, DiceMetric) and the medical-imaging terms (BI-RADS, hypoechoic, patient-grouped folds) as they come up. No prior ultrasound experience is needed.

About the Dataset

The dataset is BUS-BRA, a public collection of breast ultrasound images with biopsy-proven labels and tumor segmentation masks.

Each image carries a benign/malignant label, a BI-RADS (Breast Imaging Reporting and Data System) category (a radiologist's suspicion score from 2 to 5), a histology string, and a binary tumor mask. The CSV that ships with it also includes predefined cross-validation folds.

The task is binary: separate tumor from background. BUS-BRA contains 1,875 B-mode breast ultrasound images from 1,064 patients, acquired on four scanners at a cancer institute in Brazil.

What is MONAI, and Why Use it?

MONAI (Medical Open Network for AI) is an open-source PyTorch framework built specifically for medical imaging. It's a domain-specific layer that sits on top of PyTorch: you still write standard PyTorch training loops, but MONAI provides the medical imaging-specific components so you don't have to build them yourself.

It gives you:

• Transforms for medical data, loading formats like DICOM and NIfTI, normalizing intensities, resizing, and augmenting, all in a dictionary-based pipeline that keeps an image and its mask in sync.

• Network architectures common in medical segmentation (U-Net, UNETR, SegResNet, and others) ready to instantiate.

• Loss functions and metrics designed for segmentation, including Dice-based losses and the Dice metric.

The result is less boilerplate and fewer chances for an image and its mask to drift out of alignment.

What is Dice?

Dice (the Dice similarity coefficient) measures how much two regions overlap. In segmentation, it compares the model's predicted mask against the ground-truth mask and returns a score from 0 to 1: 0 means no overlap at all, 1 means a perfect match.

The formula is:

Dice = 2 × (overlap) / (predicted area + true area)

The "2 ×" in the numerator is what keeps the score in the 0-to-1 range even though the denominator counts the overlapping pixels on both sides.

Two roles it plays in this tutorial:

• As a metric, Dice is how the run is scored. A validation Dice of 0.876 means the predicted tumor masks overlap the true masks by about 88% on average.

• As a loss (DiceCELoss), a Dice-based term is what the model trains against. This is the part that matters for the class-imbalance problem: because Dice measures overlap rather than per-pixel correctness, a model can't score well by labeling everything as background. A small tumor counts as much as a large one, so the model is pushed to actually find the tumor region.

Part 1 — Data Profile Before Modeling

This first pass is data profiling. It reads every image and mask once and answers a short list of questions whose answers determine how the pipeline must be built. Running these checks takes a few seconds and saves a lot of guesswork later.

The snapshot below summarizes the properties that directly influenced the pipeline design. We’ll let these observations determine each step of the workflow.

Two of these deserve a closer look because they shape the two most important decisions.

Class Balance Drives the Loss

Tumors are small. Across the dataset, background pixels outnumber tumor pixels by more than ten to one.

A model trained with ordinary binary cross-entropy can score around 91% pixel accuracy by labeling everything as background. This high number reflects the imbalance rather than any ability to find the tumor.

The fix is a loss that rewards overlap with the actual tumor region, which points directly at Dice.

Patient Counts Drive the Split

There are fewer patients than images because many patients contribute both a left-side and a right-side scan. If a random split puts one patient's left scan in training and their right scan in validation, the validation score is inflated by leakage.

The dataset authors already solved this: the CSV ships a K5P column: a 5-fold split where P stands for patient-grouped, meaning every image from a given patient lands in the same fold. Reusing it is safer than rebuilding the same grouping by hand.

With those answers in hand, the pipeline has a specification to build now.

Part 2 — Building the Pipeline

Everything below uses MONAI for the segmentation-specific work:
transforms, dataset wrapping, the network, the loss, and the metric.

A Single Config Object

The pipeline reads all its knobs from one dataclass. Nothing downstream hard-codes a constant, so re-running an experiment with a different fold or image size is a single edit.

from dataclasses import dataclass
from typing import Tuple, Optional
from pathlib import Path

@dataclass
class TrainConfig:
    data_root: Optional[Path] = None
     fold_column: str = "K5P"          # patient-grouped 5-fold (dev set)
    val_fold: int = 1                 # which K5P fold is validation
    test_column: str = "HOP"          # patient-grouped hold-out partition
    test_group: int = 1               # HOP value reserved as the test set

    image_size: Tuple[int, int] = (256, 256)
    batch_size: int = 16
    lr: float = 1e-3
    epochs: int = 30
    use_amp: bool = True              # mixed precision
    ckpt_path: str = "best_model.pt"

cfg = TrainConfig()

The code above defines a TrainConfig dataclass holding every setting the pipeline needs: the fold column and which fold to validate on, the target image size, batch size, learning rate, epoch count, a mixed-precision switch, and where to save the best model. Creating cfg once gives every later step a single place to read its settings from.

The Patient-grouped Split

The split uses two predefined columns. HOP (Hold-Out Partition) reserves a patient-disjoint slice as the test set, untouched until the very end. Within the remaining development set, one K5P fold becomes validation and the other four are training. Short assertions confirm no patient appears in more than one split.

dev_df   = manifest[manifest[cfg.test_column] != cfg.test_group]
test_df  = manifest[manifest[cfg.test_column] == cfg.test_group]

train_df = dev_df[dev_df[cfg.fold_column] != cfg.val_fold]
val_df   = dev_df[dev_df[cfg.fold_column] == cfg.val_fold]

# no patient may appear in more than one split
for a, b in [(train_df, val_df), (train_df, test_df), (val_df, test_df)]:
    assert not (set(a["Case"]) & set(b["Case"])), "patient leakage"

The above code first splits off the HOP test set, then divides the remaining development rows into validation (the chosen K5P fold) and training (the rest). It then checks that every pair of splits shares no patient Case. If any does, the assertion fails immediately.

Transforms, Chosen by the Snapshot

MONAI's dictionary transforms operate on records keyed by name ("image" and "label") and apply matched operations to both. Each step here answers a Part 1 data profile finding.

from monai.transforms import (
    Compose, LoadImaged, EnsureChannelFirstd, ScaleIntensityd,
    AsDiscreted, Resized, RandFlipd, EnsureTyped,
)
import torch

base = [
    LoadImaged(keys=["image", "label"], reader="PILReader", image_only=True),
    EnsureChannelFirstd(keys=["image", "label"]),
    ScaleIntensityd(keys="image"),                       # brightness spread
    AsDiscreted(keys="label", threshold=0.5),            # clean {0, 1} mask
    Resized(keys=["image", "label"],                     # hundreds of sizes
            spatial_size=cfg.image_size,
            mode=("bilinear", "nearest")),
]

train_transforms = Compose(base + [
    RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),  # horizontal
    EnsureTyped(keys=["image", "label"], dtype=torch.float32),
])
val_transforms = Compose(base + [
    EnsureTyped(keys=["image", "label"], dtype=torch.float32),
])

The above code builds a shared list of base steps, loads the PNG, moves the channel to the front, scales the image to [0, 1], binarizes the mask, and resizes both to 256×256. It then wraps that list in two pipelines. The training pipeline adds a random horizontal flip, and the validation pipeline does not, so evaluation always sees the image as-is.

Horizontal flips are a simple augmentation that preserve anatomical plausibility in this dataset. More aggressive augmentations, such as large rotations or elastic deformations, should be validated carefully because they may distort clinically meaningful structures.

Images use bilinear interpolation to preserve intensity gradients, while masks use nearest-neighbor interpolation so class labels remain strictly 0 or 1. Bilinear interpolation on masks would create artificial label values along object boundaries.

Model, Loss, and Metric

The network is a MONAI UNet with one input channel (grayscale) and one output channel (the tumor logit). The loss is the one the class-balance finding pointed at.

U-Net consists of an encoder that captures context at progressively coarser resolutions and a decoder that reconstructs fine spatial detail. Skip connections transfer high-resolution features directly from encoder to decoder, making U-Net especially effective for medical segmentation where boundaries matter.

from monai.networks.nets import UNet
from monai.losses import DiceCELoss
from monai.metrics import DiceMetric
from monai.transforms import Activations, AsDiscrete

model = UNet(
    spatial_dims=2, in_channels=1, out_channels=1,
    channels=(16, 32, 64, 128, 256), strides=(2, 2, 2, 2),
    num_res_units=2,
).to(device)

loss_fn = DiceCELoss(sigmoid=True)       # Dice handles the imbalance; CE smooths the gradient
metric  = DiceMetric(include_background=True, reduction="mean")
post_pred = Compose([Activations(sigmoid=True), AsDiscrete(threshold=0.5)])
   

The above code creates the U-Net (five resolution levels, one input and one output channel) and moves it to the GPU. It then defines the three pieces that surround it: the loss, the validation metric, and a post_pred step that turns raw model outputs into a clean 0/1 mask by applying a sigmoid and thresholding at 0.5.

DiceCELoss combines two terms. The Dice part is scale-invariant in the foreground area, so a small tumor counts as much as a large one and the model can't win by ignoring tumors. The cross-entropy part adds a smoother gradient where Dice is flat. The sigmoid=True flag tells the loss to apply the activation itself, so the model outputs raw logits and the post_pred step handles the sigmoid-and-threshold at evaluation time. This U-Net comes out to about 1.6 million parameters.

The training loop itself is mostly standard PyTorch. MONAI stays out of the optimization logic, the only segmentation-specific pieces are the loss, transforms, and evaluation metric.

for epoch in range(1, cfg.epochs + 1):
    model.train()
    for batch in train_loader:
        img, lab = batch["image"].to(device), batch["label"].to(device)
        optimizer.zero_grad(set_to_none=True)
        with torch.amp.autocast("cuda", enabled=cfg.use_amp):
            loss = loss_fn(model(img), lab)
        scaler.scale(loss).backward()
        scaler.step(optimizer); scaler.update()

    model.eval(); metric.reset()
    with torch.no_grad():
        for batch in val_loader:
            img, lab = batch["image"].to(device), batch["label"].to(device)
            pred = post_pred(model(img))
            metric(y_pred=pred, y=lab)
    val_dice = metric.aggregate().item()
    if val_dice > best_dice:
        best_dice = val_dice
        torch.save(model.state_dict(), cfg.ckpt_path)

In the above code, each epoch runs two passes. The training pass moves every batch to the GPU, computes the loss under mixed precision, and updates the weights through the gradient scaler. The validation pass then runs with gradients turned off, converts predictions with post_pred, and accumulates Dice across the fold. Whenever the epoch's Dice beats the best seen so far, the model weights are saved to disk.

Reading the Results

Two curves summarize the run. Training loss falls steadily and flattens near 0.12. Validation Dice climbs from about 0.57 to a plateau, with a best of 0.866 reached at epoch 28.

A few things are worth reading off these curves:

• The loss decreasing monotonically means the model is learning. The gradient signal is real.

• The loss flattening above zero rather than reaching it is expected. DiceCELoss has a floor, because the cross-entropy term never fully vanishes on ambiguous boundary pixels. A loss that reached zero would be a warning sign, not a triumph.

• Validation Dice plateauing above ~0.85 while training loss keeps falling is the mild-overfitting signature. Extra epochs mostly lower train loss without moving val Dice. It's not severe here, so the 30-epoch budget is fine, but a patience-based early-stopping rule would be a reasonable add.

A validation Dice of 0.866 sits in a reasonable range for a plain 2D U-Net on this dataset. But validation Dice measures a checkpoint chosen using that same set, so it runs a little optimistic.

The final, untouched check is the HOP test set, scored exactly once, after all training and model selection are done. It comes in at 0.864, essentially matching the 0.866 validation figure. The model generalizes to patients it never saw during training or selection, and the validation number wasn't hiding leakage.

Prediction Visualization

Metrics summarize overall performance, but they don’t show how the model is segmenting individual tumors.

The figure below presents a representative validation example. From left to right are the input ultrasound image, the ground-truth mask, the model’s predicted mask, and the prediction overlaid on the original image.

The close agreement between the prediction and the ground-truth annotation illustrates how the model localizes both the position and the boundary of the lesion.

The Failure Modes Matter More Than the Average

An average Dice of 0.866 can hide very different behaviors. It could mean every case is mediocre, or most cases are excellent and a few fail badly.

To distinguish between those possibilities, sort the validation set by per-image Dice and inspect the lowest-scoring predictions.

On this fold, only 4 of 299 validation cases scored below 0.5, about 1%. Looking at those four overlays surfaces a clear pattern. Three of the four worst predictions are fragmented: the model outputs several disconnected blobs where the ground truth is a single region. The fourth confuses a dark acoustic shadow, a common ultrasound artifact, for tumor tissue.

That fragmentation pattern connects straight back to a snapshot finding: the data-quality pass measured that every ground-truth mask in BUS-BRA is a single connected component. So a multi-blob prediction is wrong by a property of the dataset, which points at keeping only the largest connected component as a post-processing step:

from monai.transforms import KeepLargestConnectedComponent

post_pred = Compose([
    Activations(sigmoid=True),
    AsDiscrete(threshold=0.5),
    KeepLargestConnectedComponent(applied_labels=[1]),
])

This code rebuilds the post_pred pipeline with one extra step at the end. After the sigmoid and threshold produce a binary mask, KeepLargestConnectedComponent discards every predicted region except the largest one, so a prediction split into several blobs collapses to its single biggest piece. This matches the dataset's one-region-per-mask property.

I measured this on the validation set, and the honest result is more nuanced than "free accuracy." It recovers a few of the fragmented cases, but the net change in mean Dice is marginal and can even go slightly negative. When a real lesion is predicted as two touching pieces, discarding the smaller one throws away true-positive area. So it's a targeted lever for a specific failure mode, not a free boost: worth exploring, not adopting blindly.

The shadow-confusion case is harder still, telling a hypoechoic tumor from a dark shadow region sometimes needs context a small grayscale crop doesn't carry. This points toward higher resolution or a wider receptive field as directions for later experiments.

Where to Go Next

Once you have a reliable baseline, the next experiments become much more meaningful. Rather than randomly trying larger models, start from the failure modes you observed:

• Replace the 2D U-Net with Attention U-Net or DynUNet.

• Train at higher resolution to better capture small lesions.

• Apply connected-component analysis selectively during inference.

• Explore test-time augmentation.

• Compare DiceCE with Focal Tversky loss for highly imbalanced lesions.

Takeaway

The through-line is profile the data, then let what you find make the decisions.

The resize came from a resolution check. The loss came from a class-balance check. The split came from a patient count. The most useful post-processing idea came from a mask-component check run before training started. None of these were guesses, and none of them needed a sweep to discover.

A model is easy to build. A model whose every choice has a reason behind it is easier to trust, easier to debug, and easier to explain to the person who reads it next.

Reference

The complete, runnable code for this walkthrough is available as a MONAI notebook: busbra_segmentation_monai.ipynb. It runs top to bottom on Kaggle or Colab, and auto-downloads the dataset if it's not already attached.

For instance, current translation systems know how to differentiate the meaning of "bank" in a sentence like:

"I can't make the bank deposit today," and "We shall meet near the river bank."

Both sentences have "bank" in them, but with different meanings.

So how did we get here? This huge revolution started back in June of 2017 when a team of 8 Google researchers, notoriously known as the "8 Samurai," released a research paper titled "Attention Is All You Need". This date marked a turning point in modern AI systems and architecture.

For context, this framework is the bedrock of current LLMs like ChatGPT and all large language models.

The 8 Google researchers who created the Transformer architecture

So, what is NMT, and how were Google engineers able to develop a framework that enables machines to understand the semantic meaning of each word in a sentence?

Table of Contents

• Demystifying NMT: The Brain Behind the Screen

• How the Transformer Sees the World

• Why This Matters

• The Democratization of AI

• What is QVAC?

• The Architecture Supported by QVAC

• The Inference Pipeline

• Setting Up the Project

• Complete Implementation

• Conclusion

• Resources and Further Reading

Demystifying NMT: The Brain Behind the Screen

To understand this breakthrough, we first have to pull back the curtain on what NMT (Neural Machine Translation) actually means.

For decades, computer translation was "rule-based." The computer was essentially given a massive bilingual dictionary and a set of grammar rules. It would translate a sentence word-by-word, swap a few positions around, and hope for the best.

This is why early translations felt so incredibly stiff and robotic: the computer was trying to solve language like a math problem.

NMT changed the game by introducing Neural Networks, computer systems inspired by the human brain. Instead of memorizing strict rules, an NMT system learns by looking at millions of existing human translations. It looks at how humans translate phrases, captures patterns, and learns how words actually interact in the real world.

But even early NMT systems had a massive flaw: they read sentences sequentially, from left to right. If a sentence was too long, the system would "forget" how it started by the time it reached the end.

This is where the Google researchers made their historic leap.

How the Transformer Sees the World

The "Attention Is All You Need" paper solved the memory problem by introducing a brand-new architecture called the Transformer. Instead of reading a sentence word-by-word, the Transformer reads the entire sentence all at once.

To do this, it splits the job into two main parts: the Encoder and the Decoder.

The Encoder (The Reader)

Think of the Encoder as a highly analytical reader. When you feed a sentence into the system, the Encoder’s job is to read it and build a "mental map" of what the sentence actually means.

It does this using a mechanism called Self-Attention. You can think of Self-Attention as a series of spotlights. When the computer looks at a specific word, it shines spotlights on all the other words in the sentence to see how they relate.

Going back to our earlier example:

"We shall meet near the river bank."

When the Encoder processes the word "bank," its Self-Attention spotlight instantly flags the word "river." Because those two words are highly connected on the AI's mental map, the system immediately knows we're talking about land next to water, not a financial institution. It locks in this "semantic meaning" before moving to the next step.

The Decoder (The Writer)

Once the Encoder has mapped out the true meaning of the sentence, it hands this blueprint over to the Decoder.

The Decoder is the writer. Its only job is to translate that blueprint into the target language. But it doesn't just output a pre-written template. It builds the new sentence word-by-word, constantly looking back at the Encoder's blueprint (using a trick called Cross-Attention) to make sure it maintains the correct context, tone, and grammar.

If it's translating our river bank sentence into French, it knows to write "la rive" (the bank of the river) instead of "la banque" (the financial bank), because the Encoder's blueprint warned it ahead of time.

Why This Matters

By teaching machines to look at the whole picture rather than individual words, Google’s engineers didn't just build a better translator. They built a system that finally understands the nuances, idioms, and context of human language.

And as it turns out, if an AI can understand the context of a sentence well enough to translate it, it can also use that same context to write essays, answer complex questions, and code. The 2017 translation engine accidentally became the foundation of the entire AI era.

The Democratization of AI

A few years after the Transformer's invention, building with it was strictly a toy for the rich. If you wanted to implement even a simple translation feature, you had to pay Big Tech giants like Google a fortune once you went beyond their tiny free tier.

Trying to bypass their dominance was almost impossible because there were practically no resources for independent developers. Back then, just understanding the basic math of a Transformer required an academic PhD. Without a massive research department at your back, trying to build your own solution from scratch was an incredibly expensive nightmare.

Thankfully, the open-source developer community has worked tirelessly to democratize access to AI. Today, we have incredibly powerful models that anyone can download and use freely.

On top of that, the processors in our personal devices have become exceptionally capable. This hardware evolution means that sophisticated AI models can now run locally directly on your smartphone, ensuring maximum data privacy and removing the dependency on external servers.

As the saying goes, "Today it needs a full building to function, tomorrow it will fit in your pocket." Of course, I totally made that quote up 😅, but you get my point!

To put this in action, we'll build a mobile application with Expo and QVAC that translates English to French.

What is QVAC?

QVAC (QuantumVerse Automatic Computer) is a decentralized, local-first AI development platform and SDK created by Tether.

Unlike traditional AI tools that require cloud connectivity, QVAC allows users to run AI models entirely on their own devices. By keeping the computation local and offline, it ensures your data remains private, secure, and entirely under your control.

Key Concepts for On-Device Translation

To understand how QVAC runs on a mobile device, we must keep a few key concepts in mind:

1. On-Device Inference:

Running model calculations locally. Rather than relying on a single engine or cloud API, QVAC supports specialized local inference backends depending on the task.

For translation, it uses the Bergamot engine under the hood. These engines memory-map quantized model weights directly into the device's RAM and run calculations using native hardware acceleration.

2. Quantization

A mathematical optimization technique that compresses the model's weights. This makes it possible for models to fit into the memory constraints of consumer mobile hardware while keeping output quality high.

The Architecture Supported by QVAC

Before writing code, it's crucial to understand what's actually happening under the hood. To handle local execution without melting your device, the QVAC SDK manages the hardware binding and model lifecycle while hooking into optimized inference backends.

For translation, QVAC utilizes the Bergamot engine. Originally developed as part of the Bergamot project (which powers Firefox's offline translation), this engine is highly optimized for fast, accurate Neural Machine Translation (NMT) on consumer hardware.

At its core, the Bergamot engine takes a source sentence, processes it through its Encoder-Decoder transformer architecture, and predicts the target language tokens in a highly efficient manner.

Understanding Language Pairs

It's important to understand the mechanics of how these models are trained. Translation models like the ones used by Bergamot are strictly unidirectional language pairs. This means the BERGAMOT_EN_FR model is designed exclusively to translate from English to French. It can't reverse the process.

If you want to translate French back to English, you would need to download and load a completely separate model trained specifically for that direction.

If a model is trained to be bidirectional (English ↔ French) or multilingual (translating dozens of languages like large language models do), it has to store mathematical representations, vocabulary, and grammar rules for multiple linguistic directions inside a single neural network. This balloons the parameter count, making the file size massive and requiring heavy RAM and compute power to process.

By isolating the task to a single direction (for example BERGAMOT_EN_FR), the model only needs the neural network to "understand" English inputs and "generate" French outputs. It doesn't need the capacity to generate English text.

This extreme specialization is exactly how Bergamot shrinks the model weights down to those incredibly tiny 15–35MB files that can run instantly on a local CPU without freezing your browser.

The Inference Pipeline

To visualize how we interact with the translation engine in our codebase, think of local translation as running a dedicated interpreter right in your phone's memory:

1. Hiring the interpreter (loading the model): We map the compressed model file (in this case, the BERGAMOT_EN_FR English-to-French model) directly into the device's RAM.

2. Handing over the script (text input): We pass the source text to the loaded engine.

3. The performance (inference): The engine reads the text and mathematically predicts the translated tokens, providing the translated result once the process is complete.

4. Closing the show (unloading): Because neural network models are memory-intensive, the model can be cleared from RAM to free up resources once the translation is complete or when the user leaves the screen.

Setting Up the Project

To ensure this guide is completely self-contained, let's start by quickly generating our new Expo application and installing the QVAC SDK. Open your terminal and run the following commands:

npx create-expo-app translator-app --template blank-typescript
cd translator-app
npm install @qvac/sdk jiti

Next, you need to add the following peer dependencies to your package.json for QVAC to work correctly. Add these lines to their respective sections:

  "dependencies": {
    "bare-rpc": "^1.0.0",
    "react-native-bare-kit": "^0.11.5"
  },
  "devDependencies": {
    "bare-pack": "^1.5.1"
  }

Once added, install the dependencies by running:

npm install
npx expo install expo-file-system expo-build-properties expo-device

Configuring the Expo Plugin with JITI

Next, we need to add the QVAC SDK plugin to our Expo project. Because the QVAC SDK's Expo plugin is distributed as a modern ECMAScript Module (ESM), but Expo's configuration file (app.config.js) runs in a standard Node.js CommonJS environment, we can't use a standard require().

This is why we installed jiti. It acts as a bridge, allowing us to synchronously load ESM modules inside CommonJS files without breaking the build process.

Create or update your app.config.js file at the root of your project and configure it like this:

const createJiti = require("jiti");
const jiti = createJiti(__filename);

// Synchronously require the ESM module using jiti
const qvacModule = jiti("@qvac/sdk/expo-plugin");
const withQvacSDK = qvacModule.withQvacSDK || qvacModule.default;

// (Include your withEscapeBundleShellScript helper if needed)

module.exports = ({ config }) => {
  config.plugins = [
    [
      "expo-build-properties",
      {
        android: { minSdkVersion: 29 },
      },
    ],
    withQvacSDK,
    "expo-router",
    [
      "expo-splash-screen",
      {
        backgroundColor: "#208AEF",
      },
    ],
    withEscapeBundleShellScript, // Custom helper if applicable
  ];

  return config;
};

This configuration applies the QVAC native setup scripts and ensures Android requires at least SDK version 29 (which is necessary for the native libraries).

With our base configuration ready to go, let's jump straight into the translation code.

Complete Implementation

Let's bring it all together. We'll implement an interface that takes English text, manages the downloading and loading states for the Bergamot engine, translates the text to French, and renders the output to the screen.

Replace your entry app file src/app/index.tsx with the following implementation:

import { View, ScrollView, TextInput, Text, TouchableOpacity, StyleSheet } from "react-native";
import { useState, useEffect } from "react";
import {
  loadModel,
  translate,
  unloadModel,
  BERGAMOT_EN_FR,
  getModelInfo,
} from "@qvac/sdk";
import { Stack } from "expo-router";

type TranslationStatus =
  | "Idle"
  | "Checking model..."
  | "Downloading model..."
  | "Model downloaded successfully."
  | "Loading model..."
  | "Translating..."
  | "Streaming translation..."
  | "Translation finished."
  | `Error: ${string}`;

export default function HomeScreen() {
  const [status, setStatus] = useState<TranslationStatus>("Checking model...");
  const [translatedText, setTranslatedText] = useState<string>("");
  const [inputText, setInputText] = useState<string>("");
  const [isTranslating, setIsTranslating] = useState<boolean>(false);
  const [isDownloaded, setIsDownloaded] = useState<boolean | null>(null);
  const [downloadProgressStr, setDownloadProgressStr] = useState<string>("");

  useEffect(() => {
    const checkModelStatus = async () => {
      try {
        const model = await getModelInfo({ name: BERGAMOT_EN_FR.name });
        setIsDownloaded(model.isCached);
        console.log("Model", model);
        setStatus("Idle");
      } catch (error) {
        console.error("Error checking model:", error);
        setStatus("Error: Failed to check model status");
      }
    };
    checkModelStatus();
  }, []);

  const handleDownload = async () => {
    try {
      setIsTranslating(true);
      setStatus("Downloading model...");
      setDownloadProgressStr("");

      const modelId = await loadModel({
        modelSrc: BERGAMOT_EN_FR,
        modelType: "nmt",
        onProgress: (progress: any) => {
          let pct = progress.percentage;
          let dl = progress.downloaded;
          let tot = progress.total;
          if (progress.shardInfo) {
            pct = progress.shardInfo.overallPercentage;
            dl = progress.shardInfo.overallDownloaded;
            tot = progress.shardInfo.overallTotal;
          }
          const formatBytes = (bytes: number) => {
            if (bytes === 0) return "0 B";
            const k = 1024;
            const sizes = ["B", "KB", "MB", "GB"];
            const i = Math.floor(Math.log(bytes) / Math.log(k));
            return (
              parseFloat((bytes / Math.pow(k, i)).toFixed(2)) + " " + sizes[i]
            );
          };
          setDownloadProgressStr(
            `${pct.toFixed(2)}% (${formatBytes(dl)} / ${formatBytes(tot)})`,
          );
        },
        modelConfig: {
          engine: "Bergamot",
          from: "en",
          to: "fr",
          beamsize: 1,
          normalize: 1,
          temperature: 0.2,
          norepeatngramsize: 3,
          lengthpenalty: 1.2,
        },
      });

      await unloadModel({ modelId, clearStorage: false });

      setIsDownloaded(true);
      setStatus("Model downloaded successfully.");
    } catch (error: any) {
      console.error(error);
      setStatus(`Error: ${error.message}`);
    } finally {
      setIsTranslating(false);
      setDownloadProgressStr("");
    }
  };

  const handleTranslate = async () => {
    if (!inputText.trim()) {
      setStatus("Error: Please enter text to translate");
      return;
    }

    try {
      setIsTranslating(true);
      setTranslatedText("");
      setStatus("Loading model...");

      const modelId = await loadModel({
        modelSrc: BERGAMOT_EN_FR,
        modelType: "nmt",

        modelConfig: {
          engine: "Bergamot",
          from: "en",
          to: "fr",
          beamsize: 1,
          normalize: 1,
          temperature: 0.2,
          norepeatngramsize: 3,
          lengthpenalty: 1.2,
        },
      });

      setStatus(`Translating...`);

      const result = translate({
        modelId,
        text: inputText,
        modelType: "nmt",
        stream: false,
      });

      const text = await result.text;
      setTranslatedText(text);

      const stats = await result.stats;
      if (stats) {
        console.log(`▸ Processing stats:`, stats);
      }

      setStatus("Translation finished.");

      await unloadModel({ modelId, clearStorage: false });
    } catch (error: any) {
      console.error(error);
      setStatus(`Error: ${error.message}`);
    } finally {
      setIsTranslating(false);
    }
  };

  return (
    <>
      <Stack.Screen
        options={{
          headerTitle: "Translator",
          headerStyle: { backgroundColor: "#000" },
          headerTintColor: "#fff",
        }}
      />
      <ScrollView contentContainerStyle={styles.scrollContainer}>
        <View style={styles.card}>
          <View style={styles.header}>
            <Text style={styles.title}>
              English to French Translator
            </Text>
            <Text style={styles.subtitle}>
              Enter text to translate:
            </Text>
          </View>

          <View style={styles.content}>
            <TextInput
              style={[styles.input, isTranslating && styles.disabledText]}
              multiline
              placeholder="Type English text here..."
              placeholderTextColor="#888"
              value={inputText}
              onChangeText={setInputText}
              editable={!isTranslating}
            />

            <Text style={styles.statusText}>
              Status: {status}
              {downloadProgressStr ? `\n${downloadProgressStr}` : ""}
            </Text>

            {isDownloaded === null ? (
              <TouchableOpacity disabled style={[styles.button, styles.buttonDisabled]}>
                <Text style={styles.buttonText}>
                  Loading...
                </Text>
              </TouchableOpacity>
            ) : isDownloaded ? (
              <TouchableOpacity
                onPress={handleTranslate}
                style={[
                  styles.button,
                  (isTranslating || !inputText.trim()) && styles.buttonDisabled,
                ]}
                disabled={isTranslating || !inputText.trim()}
              >
                <Text style={styles.buttonText}>
                  {isTranslating ? "Translating..." : "Translate"}
                </Text>
              </TouchableOpacity>
            ) : (
              <TouchableOpacity
                onPress={handleDownload}
                style={[styles.button, isTranslating && styles.buttonDisabled]}
                disabled={isTranslating}
              >
                <Text style={styles.buttonText}>
                  {isTranslating ? "Downloading..." : "Download Model"}
                </Text>
              </TouchableOpacity>
            )}

            <View style={styles.outputContainer}>
              <Text style={styles.outputText}>
                {translatedText || "Translation will appear here..."}
              </Text>
            </View>
          </View>
        </View>
      </ScrollView>
    </>
  );
}

const styles = StyleSheet.create({
  scrollContainer: {
    flexGrow: 1,
    paddingHorizontal: 16,
    paddingTop: 16,
    paddingBottom: 24,
    backgroundColor: "#f9fafb",
  },
  card: {
    backgroundColor: "#ffffff",
    maxWidth: 450,
    width: "100%",
    alignSelf: "center",
    borderRadius: 12,
    padding: 16,
  },
  header: {
    marginBottom: 16,
  },
  title: {
    textAlign: "center",
    fontSize: 24,
    fontWeight: "bold",
    color: "#111827",
  },
  subtitle: {
    textAlign: "center",
    marginTop: 4,
    fontSize: 16,
    color: "#6b7280",
  },
  content: {
    gap: 24,
  },
  input: {
    borderWidth: 1,
    borderColor: "#e5e7eb",
    backgroundColor: "#ffffff",
    color: "#111827",
    padding: 12,
    borderRadius: 8,
    minHeight: 100,
    textAlignVertical: "top",
  },
  disabledText: {
    opacity: 0.5,
  },
  statusText: {
    fontSize: 14,
    color: "#3b82f6",
    fontWeight: "bold",
    textAlign: "center",
    marginTop: 12,
    marginBottom: 12,
  },
  button: {
    width: "100%",
    height: 48,
    borderRadius: 12,
    backgroundColor: "#3b82f6",
    alignItems: "center",
    justifyContent: "center",
  },
  buttonDisabled: {
    opacity: 0.5,
  },
  buttonText: {
    fontWeight: "600",
    fontSize: 18,
    color: "#ffffff",
  },
  outputContainer: {
    marginTop: 16,
    padding: 16,
    backgroundColor: "#f3f4f6",
    borderRadius: 8,
    minHeight: 100,
  },
  outputText: {
    fontSize: 16,
    color: "#111827",
  },
});

Here is a translation example from the application.

Input (English)

The location I told you was near the river bank

Output (French)

L'endroit où je vous ai dit était près de la rive de la rivière

Codebase Breakdown

Let’s lift the hood on how this local translation implementation manages native model lifecycles and processes the streamed tokens.

1. Managing the Native Lifecycle

Loading neural network weights for translation is computationally expensive. When the QVAC runtime initializes a model, it must read parameters from the local disk and copy the active weights into device RAM.

To handle this efficiently, we check if the model is cached before attempting to load it. This is used to check if the model is downloaded. That's the meaning of cached: it means the model has been downloaded to the user's disk:

const model = await getModelInfo({ name: BERGAMOT_EN_FR.name });
setIsDownloaded(model.isCached);

The loadModel function will automatically handle downloading the model from the Hugging Face hub if it hasn't been cached locally yet. Once the file is available locally, it directly memory-maps the weights.

2. Translating the Text

Once the model is loaded, we can pass our text to the translation engine:

const result = translate({
  modelId,
  text: inputText,
  modelType: "nmt",
  stream: false,
});

const text = await result.text;
setTranslatedText(text);

This waits for the full translation to complete before displaying the final result to the user.

3. Unloading the Model

After the translation is complete, we explicitly destroy the model via unloadModel:

await unloadModel({ modelId, clearStorage: false });

By unloading the model, we ensure that the device's RAM is freed up for other processes. Because the model is already downloaded and cached on the disk (and we explicitly set clearStorage: false), reloading the model the next time the user wants to translate something will be nearly instantaneous.

Conclusion

Transitioning translation from the cloud to on-device hardware offers a practical approach for mobile application developers. Running model inference locally eliminates reliance on remote internet connectivity, removes recurring API usage costs, and ensures that user text inputs never leave the physical device.

Integrating local translation can be highly beneficial for travel apps, secure communication tools, or educational platforms. As edge processors gain dedicated hardware acceleration cores and open-source models become even more efficient through quantization research, local-first architectures present a compelling alternative for developers prioritizing privacy, offline resilience, and predictable cost structures.

Resources and Further Reading

To dive deeper into local Neural Machine Translation, inspect the source code, or explore advanced configurations for your mobile applications, check out the following resources:

• QVAC Translation Docs: Official documentation for integrating local translation capabilities with QVAC.

• QVAC Expo Integration Docs: Learn more about configuring custom local models in Expo.

• Bergamot Project: Learn more about the underlying Neural Machine Translation engine.

• Attention Is All You Need: The original 2017 Google research paper that introduced the Transformer architecture.

• Full Code Example: Full code example's repository.

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.

I know that feeling – yes, that thing you’re feeling now. It’s either confidence that you know what I’m asking, or the lack of it. Worry not, buddy: I’ve got you.

In this tutorial, I’ll explain all you need to know about deep learning, neural networks, and why I think you need to know about a fancy tool called Keras.

Prerequisites:

This is a conceptual article, so you don't need any deep learning background to follow along. That's exactly what you'll be gaining here.

Basic Python familiarity is helpful but not required for this article. But if you'd like to get hands-on with Keras afterward, having Python 3.9+ and pip installed will make it easy to start experimenting.

Table of Contents

1. What is Deep Learning?

2. So What Are Neural Networks?

3. An Analogy for How Neural Networks Work

4. What Are CNNs, RNNs, and Transformers?

5. How CNNs Work

6. How RNNs Work

7. How Transformers Work

8. Keras For Building ML Models

9. Wrapping Up

What is Deep Learning?

To explain Deep Learning to you, I assume that you're already familiar with what AI (Artificial Intelligence) & ML (Machine Learning) are all about. These terms are likely not new to your ears, especially these days.

But maybe just to summarise: AI is a technology that enables computers to simulate human cognitive abilities such as learning and comprehension, problem-solving, creativity, and autonomy.

ML is a subset of artificial intelligence. It deals with the development of algorithms that can recognize patterns and learn from training data and subsequently make accurate inferences on new data without explicitly being programmed to do so.

Deep Learning is a subset of machine learning that's driven by multilayered neural networks whose design is inspired by the structure of the human brain.

Take a look at the diagram below for a visual understanding of how these concepts are layered:

Image Source — Research Gate

According to a book I’m currently reading titled Deep Learning in Python by François Chollet, the word “Deep” in Deep Learning isn’t a reference to any kind of deeper understanding achieved by this concept. Rather, it stands for this idea of successive layers of representations of data. These layers of representations are learned via models called neural networks, structured in literal layers stacked on top of each other.

Check out the image below for a vivid example of this layering:

Image Source — Research Gate

An interesting assumption Chollet also clears up is that Deep Learning models (Neural Networks) aren’t models of the brain. Rather, it's just that some central parts of deep learning were inspired by our understanding of the brain, in particular the visual cortex.

But do you know what the visual cortex is? 😅 See the image below (the circled green part):

The visual cortex is the part of the brain that processes what your eyes see. In the 1960s, neuroscientists Hubel and Wiesel discovered something surprising while studying it: neurons deeper in the visual cortex didn't respond to whole objects directly. Instead, the earliest neurons fired only for simple things, like a line at a specific angle. Only in later stages did neurons combine those simple signals into a response for more complex shapes.

In other words, the brain builds up "seeing an object" in stages — simple patterns first, complexity later. That stage-by-stage structure is the loose inspiration for how CNNs stack layers of filters, which we'll get to shortly.

Now that you have the hang of this, let’s explore the interesting and mind-boggling concept of neural networks.

📌 Note: neural networks are only mind-boggling at the start because they're a new concept. Once you take some time to understand them, they'll be easier to comprehend.

So What Are Neural Networks?

For starters, a neural network is a concept in deep learning. The “neural” in the name is derived from the neurons of the human brain.

A neural network consists of connected units or nodes called artificial neurons, which loosely model the neurons in the brain.

Here is the serious and technical definition:

A neural network is a machine learning model that stacks simple “neurons” in layers and learns pattern-recognizing weights and biases from data to map inputs to outputs. (Excerpt from IBM Blog )

Now here's the easier and more fun definition:

A Neural network is just a machine for making guessing/pattern recognition/analysis mistakes smaller, one small correction at a time.

Neural Networks come in various architectures/types such as;

1. CNNs (Convolutional Neural Networks)

2. RNNs (Recurrent Neural Networks)

3. Transformers

We’ll explore these later in this article, so for now, just understand that what makes them different is simply what they choose to focus on before they guess or recognize a pattern or provide an analysis.

An Analogy for How Neural Networks Work

Image Source

Imagine you’re learning to throw darts blindfolded and someone can tell you one thing after each throw. For example, “You're 6 inches too far left and 2 inches too low”. You're not told why or given a lecture on your positioning for the aim or your grip of the dart. You just get a distance & direction correction.

So you nudge (a light shift or twist) your arm angle a little based on that correction feedback and your throw again and get a new correction. You then nudge again and again and againnnn(!) until you hit the target.

Over time, as you do this, your arm “learns” (underline this) – not because anyone explained dart throwing physics to you, but because every throw gave you a tiny specific correction and you kept applying corrections in the direction that shrank your possibility of missing the target.

That’s exactly how neural networks work.

Let’s now learn the technical jargon that we’d use to talk about neural networks from the analogy above.

The instruction or signal “Nudge your arm this much and this way” is what we call the Gradient. The gradient indicates the direction and size/distance/rate to make the correction, along with the rate at which the weights and biases should be adjusted to decrease the loss function.

The correction detail, for example 6 Inches too far to the left, is the loss. A loss function in a neural network is a mathematical tool that measures how well the model’s predictions align with the actual target values.

When you keep applying corrections in the direction that shrinks the miss or error, we call that Gradient Descent. This is the optimization algorithm (the step-by-step process) that a neural network uses to figure out which direction to move and how big a distance (step size) to take to reach that accurate value. In this context, descent means exactly what it means in plain English: the act of moving downward.

Your arm's muscle memory adjusting is the weight update. Weights are numerical values that help each node within a network make decisions by determining which factors are more important than others.

📍 Now… Pause and let that sink in. You can re-read this analogy once again if you want to before you proceed.

What Are CNNs, RNNs, and Transformers?

I hope you’re still with me here… because you need to understand these terms, too. Remember from earlier that CNNs, RNNs, and Transformers are simply architectures or different types of neural networks. They have the same learning process, similar to the analogy of throwing darts while blindfolded: they all have the guess then measure-error/loss then nudge loop underneath.

The difference is what information they choose to focus on before they make a guess/prediction/give an output.

📌 Let me break it down for you:

• CNNs (Convolutional Neural Networks) only look at the small nearby patch and analyze it, then move to the next patch. Think of it as only seeing the dart board through a small tube pointed at one spot.

• RNNs (Recurrent Neural Networks) only look at things in order and remember a running summary as it goes. Think of it like reading a story left to right and updating your mental summary of the plot after each sentence.

• With Transformers, the neural network looks at everything all at once and figures out on the fly what matters most. Think of it like reading the whole page in one glance and deciding which words connect with each other.

Let’s take a deeper look at each type of neural network.

How CNNs Work

This is a type of neural network built for data that has spatial structure, such as images. It works with a small grid of numbers called a filter (for example, 3x3). Those numbers are weights, which are initially just random and don’t mean anything

Through the guess, measure error, nudge loop, those random numbers gradually become good at reacting strongly to a specific pattern in the image, for example a vertical edge, a certain color, and so on.

📌 Note: No one tells the filter what to look for. It discovers that on its own through training, the same way every other weight in every network we discuss here does.

That filter then slides across the image a few pixels at a time, and at each position it looks at a small path and produces a single number. This is a measure of how strongly the patterns the filter has learned to detect show up there.

When the filter has slid across the whole image, you get a grid of numbers which really just show a map of where the detected pattern shows up across the image.

📌 Note: The same filter with the same numbers is reused at every single position via a technique called parameter sharing. Hence the efficiency of CNNs

See the below example of filters sliding over an image matrix:

Image Source

In real CNNs, layers of filters are stacked on top of each other, and each layer builds on the previous one's output. Early layers, working directly on the raw pixels, tend to pick up very simple things, like edges or a patch of color. Because the next layer looks at the output of the first layer rather than raw pixels, it can combine those simple edges into slightly more complex shapes, like a curve or a corner.

Layer by layer, this keeps compounding: shapes combine into parts (like an ear or a whisker shape), and parts combine into something the network can recognize as a whole object, like a cat.

That's the real payoff of stacking filter layers: none of it happens in one step, and each layer only ever has to solve a slightly harder version of the same small problem.

Here's an image that illustrates the whole CNN process:

Image Source

Use cases of CNNs include:

• Medical Imaging: CNNs analyze medical scans, such as chest X-rays, and assist clinicians by flagging potential abnormalities for review.

• Image Generation: CNNs can create new images or manipulate existing ones.

• Autonomous Systems: CNNs can be used in autonomous systems such as self-driving cars for lane detection, obstacle detection, and traffic sign recognition.

You can learn more here.

📍 Now, pause and take note of the key things that matter: filters & parameter sharing.

How RNNs Work

RNNs handle data that's sequentially ordered, where the order itself carries meaning. Think audio data and sentences that come together to form a story.

Unlike CNNs, which slide over patches of an image in no particular order, RNNs need to read things one step at a time.

For example, a sentence is read one word at a time in sequence because what’s reviewed earlier affects how the neural network understands what comes next. This sequential flow makes it slow to review large datasets.

RNNs keep a running summary called a hidden state.

📌 Note: Think of it as a small notebook where it jots down everything important that it's understood so far. At the very start, before reading anything, that notebook is essentially blank (an initial hidden state, usually all zeros).

Here's a simple architecture:

Image Source

So the hidden state is never a lookup table of everything the RNN has seen. It’s a single running summary that gets overwritten at every step, carrying forward only what the network has learned is worth keeping.

📌 Note: The downside to RNNs is that if the sequence is long, early information reviewed can fade out almost entirely, which causes the RNN to lose the context of earlier review content. This problem in RNNs is called the vanishing gradient problem.

Use cases of RNNs include:

• Speech Recognition: RNNs are used in speech recognition systems to process audio over time. They help models understand how sounds form words and sentences.

• Voice AI Systems: In voice workflows, RNNs help process sequential audio data. Combined with technologies like text-to-speech (TTS), they contribute to natural voice generation pipelines.

• Time-Series Prediction: In finance or weather forecasting, RNNs analyze past data to predict future outcomes using probabilistic methods.

• Text Generation: RNNs can generate text by predicting the next word based on previous words. This is useful in chatbots and tools powered by generative AI.

Here's a video you can watch to learn more about RNNs.

📍 Now, pause and take note of the key things that matter: hidden state and the vanishing gradient which is a downside to RNNs.

How Transformers Work

In 2017, a group of researchers at Google Brain published a short but world-shaking paper: “Attention Is All You Need.”

It introduced the Transformer, a new architecture for processing language that quietly changed how AI engineering is done. Since then, Transformers have become the backbone of nearly every major Large Language model, including Gemini.

Here is the simplified architecture:

Image source

In a nutshell, transformers take in a bunch of data at the same time (unlike RNNs that take in data in a sequential order).

Transformers are premised on a couple key concepts. First, there's self-attention, where every element of data has a positional encoding that helps the transformer know the ordering of the elements. Second, there's embeddings that help capture the meaning of each word and the contextual relationship between all the data elements. Embeddings make it easy for the transformer to process data faster compared to other kinds of neural networks.

📍 Note: the transformer architecture reduces the vanishing gradient problem that's present with the RNNs.

Use cases of transformers include:

• NLP (Natural Language Processing) Tasks: The self-attention mechanism enhances the linguistic capabilities of machine learning models by allowing the efficient and complete analysis of an entire text.

• Computer Vision: Developments in image-recognition models suggest that self-attention is a crucial component to increase their robustness and generalization.

You can learn more about Transformers from this video.

📍 Now, pause and take note of the key things that matter: self-attention, embeddings, positional encoding, and the fact that data is ingested and processed at the same time.

Keras For Building ML Models

Image Source

Now that you understand the various types of neural networks and the use cases for each, you’re probably wondering how you can start building models.

What is Keras?

There are many options, but the one tool that I’ve come to appreciate the most is Keras. It's been used in projects such as the Google YouTube Recommendation Engine and the Waymo self-driving fleet.

Keras is an open-source, high-level neural network API that's designed to be user-friendly, modular, and extensible. It was initially developed independently and could run on top of backends like TensorFlow, Theano, or CNTK.

Since 2019, it has been the official high-level API of TensorFlow (TensorFlow 2.0+), offering high-level APIs (Sequential and Functional) and built-in support for common layers, optimizers, and loss functions.

📌 In short, Keras allows you to quickly and easily build AI/ML models.

Keras provides a complete toolkit for building deep learning models. It’s never been easier to build, train, evaluate, and deploy deep learning models.

The New Version — Keras 3.0

A significant recent development is Keras 3.0. It’s a full rewrite that lets Keras workflows run on top of multiple backends, like JAX, TensorFlow, PyTorch, and OpenVINO (inference-only), instead of being tied to TensorFlow alone.

📍 Here’s what makes Keras genuinely different from just being “another way to write neural network code”:

• It doesn’t just let you build a CNN, an RNN, or a Transformer. It lets you build all three using the same pattern.

• The training loop wrapping them- the same guess, measure error, nudge loop we’ve talked about throughout this entire article never changes. Keras is really just one consistent way of expressing that loop, no matter which architecture you’re pointing it at.

You write your model once, and you can pick the framework that suits you best. You can also switch from one to another based on your current goals without rewriting the model itself.

And the flexibility isn’t just theoretical. It matters for performance, too.

In Keras’s own benchmarks, JAX typically delivers the best training and inference performance on GPU, TPU, and CPU, though results vary from model to model.

📌 Being able to swap backends without touching your model code means you’re not locked into whichever framework happened to be fastest when you started the project.

Wrapping Up

I’ll pack it in at this. I hope you now have a good understanding of CNNs, RNNs, Transformers, and where the Deep Learning framework Keras falls into all this.

That's the mental model: one learning process, three architectures shaped by the data they're built for, and Keras as the one API that lets you build any of them. If you take one thing from this, let it be the guess → measure error → nudge loop. it's the basis for everything else you'll ever learn about deep learning.

Found this helpful? You can reach out to me via email or LinkedIn and let me know what stood out for you and what you expect to learn next.

Cheers.

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.

In this tutorial, you’ll build a healthcare AI assistant that runs entirely on your own machine using three open-source tools:

• MedGemma, Google’s open medical AI model for understanding medical text and images

• Ollama, the easiest way to download and run AI models locally

• Open WebUI, a ChatGPT-style web interface for interacting with local models

By the end, you’ll be able to chat with a medically tuned AI model, upload medical images such as chest X-rays for analysis, and do it all locally, without sending your data to the cloud.

Important disclaimer before we start: MedGemma is a developer model, not a medical device. Its outputs are not intended to directly inform clinical diagnosis, patient management, or treatment decisions.

Everything you build in this tutorial is for learning, prototyping, and research. Always consult qualified healthcare professionals for real medical questions.

What We'll Cover:

• Who is This Tutorial For?

• What is MedGemma?

• Why Run Models Locally?

• Prerequisites

• Architecture Diagram

• Step 1: Install Ollama

• Step 2: Pull MedGemma

• Step 3: Test MedGemma from the Terminal

• Step 4: Install Open WebUI

  • Option A: Docker (recommended)

  • Option B: pip (no Docker)

• Step 5: Connect Open WebUI to Ollama

• Step 6: Start Chatting with MedGemma

• Step 7: Upload Medical Images

• Example Prompts to Try

• Running Larger Models

• Troubleshooting Guide

  • Error: registry.ollama.ai/library/medgemma:latest does not support tools

  • Open WebUI shows no models in the dropdown

  • ollama pull medgemma says model not found

  • Responses are extremely slow

  • Image upload doesn't work or the model ignores the image

  • Port 3000 is already in use

  • "Out of memory" errors when loading the 27B model

• Conclusion

Who is This Tutorial For?

This tutorial is ideal if you’re:

• learning healthcare AI

• building medical RAG systems

• experimenting with radiology assistants

• developing medical education tools

• researching multimodal models

What is MedGemma?

MedGemma is a collection of open models from Google, built on the Gemma 3 architecture and specifically trained for medical text and image comprehension. Think of it as Gemma after four years of medical school and a radiology residency.

Why MedGemma?

Unlike general-purpose models such as Llama or Mistral, MedGemma is designed specifically for healthcare applications.

• Medical image understanding: Its multimodal models are trained on de-identified medical images, including chest X-rays, dermatology, ophthalmology, and pathology images.

• Medical language expertise: It has been trained on medical literature and clinical question-answer datasets, enabling it to better understand medical terminology and radiology reports.

• Multiple model sizes: MedGemma is available in 4B and 27B variants, both supporting text and image inputs with a 128K context window.

• Open weights: You can download, run, fine-tune, and build applications with the model locally under the Health AI Developer Foundation's terms of use.

MedGemma is intended as a foundation model for developers building healthcare applications, medical education tools, research assistants, report summarizers, and other AI-powered medical workflows.

Why Run Models Locally?

You could call a hosted medical model through an API. So why go local? In healthcare, the case is stronger than almost anywhere else.

First, there's the principle of privacy by architecture. When the model runs on your machine, medical text and images never leave your device. There's no API log, no third-party data processor, no data processing agreement to negotiate.

For anyone working near PHI (Protected Health Information), "the data never left the laptop" is the simplest compliance story that exists.

Next, you have zero per-token cost. Experimentation is free once the model is downloaded. You can iterate on prompts hundreds of times without watching a billing dashboard.

You also get offline access. Hospitals, labs, and field clinics often have restricted or air-gapped networks. A local model works without internet after the initial download.

And you have full control over the setup: you choose the model version, you pin it, and it never changes underneath you. No deprecation notices, no silent behavior changes.

Finally, it's a great way to learn. Running models locally demystifies them. You'll develop intuition for context windows, quantization, and memory constraints that you simply don't get from calling an API.

Prerequisites

Here's what you need before starting:

Hardware:

• 8 GB RAM minimum (16 GB recommended) for the MedGemma 4B model. The download is about 3.3 GB.

• 32 GB RAM or a 24 GB+ GPU if you want to run the 27B model (a roughly 17 GB download).

• Around 15 GB of free disk space to be comfortable (model + Docker images + working room).

• Apple Silicon Macs (M1 through M4) are excellent for this. Ollama uses Metal acceleration automatically. On Windows and Linux, an NVIDIA GPU helps a lot but isn't required. A CPU-only inference works, just slower.

Software:

• macOS, Linux, or Windows 10/11

• Docker Desktop (for the recommended Open WebUI installation), or Python 3.11 if you prefer installing Open WebUI with pip

• Basic comfort with the terminal

That's it. No API keys, no accounts, and no GPU cloud credits.

Architecture Diagram

Step 1: Install Ollama

Ollama is a lightweight runtime that handles downloading, quantizing, and serving open models through a simple CLI and a local REST API.

On macOS:

Download the app from ollama.com/download and drag it to Applications, or install via Homebrew:

brew install ollama

On Linux:

curl -fsSL https://ollama.com/install.sh | sh

On Windows:

Download the native Windows installer from ollama.com/download and run it. (Ollama now supports Windows natively, no WSL required.)

Once installed, verify it works:

ollama --version

You should see a version number printed. Ollama also starts a background service that listens on http://localhost:11434. This is the API that Open WebUI will talk to later. You can confirm the server is up with:

curl http://localhost:11434

which should return Ollama is running.

Step 2: Pull MedGemma

MedGemma is available directly in the official Ollama model library, so downloading it is one command:

ollama pull medgemma

This pulls the default 4B multimodal variant, about a 3.3 GB download.

If you want to be explicit about the size (useful when you later experiment with the 27B model):

ollama pull medgemma:4b     # 3.3 GB — multimodal, runs on most laptops
ollama pull medgemma:27b    # 17 GB — multimodal, needs serious hardware

When the download finishes, confirm the model is installed:

ollama list

You should see medgemma in the output along with its size.

Step 3: Test MedGemma from the Terminal

Before adding a UI, let's make sure the model actually works. Start an interactive session:

ollama run medgemma

You'll get a >>> prompt. Try a medical question:

>>> What are the classic radiographic signs of pneumonia on a chest X-ray?

MedGemma should respond with a structured answer covering findings like consolidation, air bronchograms, and silhouette signs — the kind of answer that shows its radiology training.

Try one more to see the clinical reasoning:

>>> Explain the difference between Type 1 and Type 2 diabetes to a first-year medical student.

A few useful commands inside the session:

• /bye — exit the session

• /clear — clear the conversation context

• /show info — display model details (parameters, quantization, context length)

You can also test image input directly from the terminal by passing a file path directly in the prompt:

>>> Describe the key findings in this image. ./chest_xray_sample.png

While this works, uploading images through Open WebUI is much more convenient.

Step 4: Install Open WebUI

Open WebUI gives you a clean, ChatGPT-style interface on top of Ollama: conversation history, model switching, image uploads, and multi-user support, all self-hosted.

Option A: Docker (recommended)

Start by installing Docker.

Make sure Docker Desktop is running, then launch Open WebUI with:

docker run -d -p 3000:8080 \
  --add-host=host.docker.internal:host-gateway \
  -v open-webui:/app/backend/data \
  --name open-webui \
  --restart always \
  ghcr.io/open-webui/open-webui:main

Let's break down what this command does:

• -d runs the container in the background

• -p 3000:8080 maps port 3000 on your machine to the WebUI's internal port 8080

• --add-host=host.docker.internal:host-gateway lets the container reach the Ollama server running on your host machine

• -v open-webui:/app/backend/data creates a Docker volume so your chats and settings survive container restarts

• --restart always brings the UI back up automatically after reboots

Option B: pip (no Docker)

If you'd rather skip Docker, you can instead install Open WebUI as a Python package (Python 3.11 is the supported version):

pip install open-webui
open-webui serve

This starts the interface at http://localhost:8080 instead of port 3000.

Step 5: Connect Open WebUI to Ollama

Open your browser and go to http://localhost:3000 (or :8080 if you used pip).

On first launch, you'll be asked to create an admin account. This account is stored locally on your machine (it's not a cloud signup).

In most setups, Open WebUI auto-detects Ollama at http://localhost:11434 and you're done.

If your models don't appear, wire up the connection manually:

1. Click your profile icon and go to Admin Panel then Settings then Connections.

2. Under Ollama API, set the URL:

   • Docker install: http://host.docker.internal:11434

   • pip install: http://localhost:11434

3. Click the refresh icon to verify the connection, then save.

Head back to the main chat screen, and medgemma should now appear in the model dropdown at the top.

You can check the troubleshooting section below if you face any errors.

Step 6: Start Chatting with MedGemma

Select medgemma from the model selector and start a conversation. A good first test might look like this:

Summarize this radiology report in plain language a patient could understand:

"Impression: Mild cardiomegaly. Small right pleural effusion.
No focal consolidation. Degenerative changes of the thoracic spine."

You should get a clear, patient-friendly explanation of each finding. This "clinical language to plain language" translation is one of MedGemma's genuine strengths.

There are a few Open WebUI features worth knowing about:

• System prompts: Click the model name and set a system prompt like "You are a medical education assistant. Always explain your reasoning and cite the relevant physiology." This shapes every response in the conversation.

• Conversation history: Every chat is saved locally and searchable from the sidebar.

• Multiple models: You can add llama3.2, gemma3, or any other Ollama model and compare their answers to the same medical question side by side. This is a great way to see the difference domain training makes.

Step 7: Upload Medical Images

This is where MedGemma really separates itself from general-purpose models. Because its vision encoder was pre-trained on medical imaging, it can meaningfully describe radiographs, skin lesions, fundus photos, and histopathology patches.

To try it:

1. Start a new chat with medgemma selected.

2. Click the + (or image) icon in the message box, or simply drag and drop an image file.

3. Add a prompt alongside the image and hit send.

For sample images you can test with (without touching any real patient data), try public teaching datasets like the NIH ChestX-ray14 dataset, MedPix, or Radiopaedia's teaching cases.

Example workflow with a chest X-ray:

[Upload: chest_xray.png]

You are an expert radiology assistant. Describe this chest X-ray
systematically: technical quality, lungs, heart, mediastinum, bones,
and soft tissues. Then summarize the key findings.

MedGemma will typically walk through the image in the systematic order you asked for, which mirrors how radiologists are trained to read films.

Two important caveats:

• Ollama and Open WebUI work with standard image formats (PNG, JPEG). Clinical DICOM files need to be converted to PNG/JPEG first — a one-liner with Python libraries like pydicom + Pillow.

• Never upload images containing patient-identifying information (names, MRNs, dates burned into the image) unless the data has been properly de-identified. Even on a local machine, good data hygiene is a habit worth building.

Example Prompts to Try

Here are prompts that showcase different capabilities. Use them as starting points:

Medical education:

Create a comparison table of ACE inhibitors vs ARBs: mechanism, common examples, key side effects, and contraindications.

Clinical documentation:

Convert these shorthand clinic notes into a structured SOAP note:"45F, 3d cough + fever 101F, no SOB, lungs clear, likely viral URI, supportive care, return if worse"

Report translation for patients:

Explain this MRI impression to a worried patient in a reassuring but honest tone: "Small disc protrusion at L4-L5 without significant canal stenosis or nerve root compression."

Image analysis (with an uploaded dermatology photo):

Describe this skin lesion using the ABCDE criteria
(Asymmetry, Border, Color, Diameter, Evolution cannot be assessed from a single image — note that explicitly).

Differential reasoning:

A 60-year-old presents with sudden painless vision loss in one eye. List the top 5 differential diagnoses and the key distinguishing feature of each.

Notice a pattern: the best results come from prompts that give MedGemma a role, a structure to follow, and explicit constraints. That's true of all LLMs, but it matters even more in a domain where precision counts.

Running Larger Models

The 4B model is impressive for its size, but the 27B variant is noticeably stronger at complex clinical reasoning, longer differential diagnoses, and nuanced report interpretation.

The trade-off is hardware:

To try the 27B model:

ollama pull medgemma:27b
ollama run medgemma:27b

Practical tips for larger models:

• Watch your memory: Run ollama ps to see how much RAM/VRAM a loaded model is using and whether it's running on GPU, CPU, or split across both. A model that spills from GPU to CPU gets dramatically slower.

• On Apple Silicon, a 32 GB M-series Mac runs the 27B model comfortably.

• Free memory between models: Ollama keeps models loaded for a few minutes after use. Unload immediately with ollama stop medgemma:27b if you need the RAM back.

• Sanity-check the speed trade-off: If the 27B model generates at 2–3 tokens per second on your machine, the 4B model at 30+ tokens/second may be the better.

You can keep both installed and switch between them in the Open WebUI dropdown — 4B for fast iteration, 27B when you need the deeper reasoning.

Troubleshooting Guide

Error: registry.ollama.ai/library/medgemma:latest does not support tools

This is the most common MedGemma-specific error, and it means Open WebUI is sending native tool/function definitions with your request. MedGemma (like base Gemma 3) doesn't support Ollama's tools API, so the request is rejected before the model even sees your message.

Hunt down whatever is attaching tools, in this order:

1. Model capabilities (most likely culprit): Go to the Admin Panel, then Settings, then Models, then medgemma, then uncheck Builtin Tools, Web Search, Code Interpreter, and Terminal under Capabilities, and make sure every item in the Builtin Tools checklist is unticked. Keep Vision, File Upload, and File Context checked. Newer Open WebUI versions enable builtin tools by default, so a fresh install will hit this immediately.

2. Task model: Go to Admin Panel, then Settings, then Interface, and make sure neither the local nor external Task Model is set to medgemma. Background jobs like title and follow-up generation use tool calls — route them to llama3.2 or similar.

3. Function Calling mode: Set to Default (not Native) in the model's Advanced Params and in your user Settings, General, Advanced Parameters.

4. Global functions/filters: Go to Admin Panel, then Functions, and disable the Global toggle on any active function, since global functions attach to every model.

5. Per-chat toggles: In the message box, make sure web search and code interpreter toggles are off, and no Tools are attached via the + menu.

Then start a new chat (old chats can carry stale settings) and test. To confirm the model itself is fine, run ollama run medgemma "hello" in your terminal. If that works, the issue is purely Open WebUI configuration.

Open WebUI shows no models in the dropdown

The container can't reach Ollama. Check that:

• Ollama is actually running: curl http://localhost:11434 should return Ollama is running.

• The connection URL in Admin Panel, Settings, Connections is http://host.docker.internal:11434 (Docker) — localhost won't work from inside a container because it refers to the container itself.

• On Linux, if host.docker.internal doesn't resolve, add --network=host to your docker run command instead and use http://localhost:11434.

ollama pull medgemma says model not found

Update Ollama, as MedGemma requires a recent version. Re-run the installer or, on macOS, click the menu bar icon and then Update. Then retry the pull.

Responses are extremely slow

• Check ollama ps — if the model shows a large CPU percentage, it doesn't fit in your GPU/unified memory. Switch to the 4B model.

• Close memory-hungry apps (browsers with 40 tabs are the usual suspect).

• On first message, models take several seconds to load into memory, subsequent messages are much faster.

Image upload doesn't work or the model ignores the image

• Make sure you selected medgemma (multimodal) and not a text-only model in the dropdown.

• Use PNG or JPEG. DICOM files must be converted first.

• Very high-resolution images can cause issues — resize to something reasonable (e.g., 1024px on the long edge) before uploading.

Port 3000 is already in use

Map a different host port: change -p 3000:8080 to -p 3001:8080 and access the UI at http://localhost:3001.

"Out of memory" errors when loading the 27B model

Your machine doesn't have enough free RAM/VRAM. Stick with medgemma:4b, or free memory and try again. There is no shame in the 4B model — it punches well above its weight.

Conclusion

In this tutorial, you built a complete, private healthcare AI assistant from scratch — and it took three tools and a handful of terminal commands.

Let's recap what you accomplished:

• Installed Ollama and pulled MedGemma, a medically-tuned multimodal model, onto your own machine

• Verified the model from the terminal, then put a full chat interface on top of it with Open WebUI

• Configured the model's capabilities correctly so tool-calling features don't break a model that doesn't support them

• Chatted with a model that understands radiology reports, clinical terminology, and medical images — and uploaded images for analysis

• Learned how to scale up to the 27B model and how to diagnose the most common errors along the way.

You now have a fully private AI assistant running entirely on your own machine. From here, you can extend it with retrieval-augmented generation (RAG), integrate it with medical imaging pipelines, or connect it to de-identified clinical datasets to build more advanced healthcare AI applications.

Happy building!

Further reading:

• MedGemma on the Ollama library

• MedGemma model documentation (Google Health AI Developer Foundations)

• Ollama documentation

• Open WebUI documentation

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.

Four months later, a four-agent LangChain loop ran for eleven days and cost 47,000 USD. Nobody noticed until the invoice arrived. The pipeline worked correctly in testing, and the agents were doing exactly what they were told. Same pattern.

This tutorial is about that missing instruction.

You'll build five small Python primitives that catch most agent loop failures before they ship:

• A spec writer that forces you to define done before the loop starts

• A circuit breaker that kills the loop when it exceeds hard limits

• A ledger that records every turn in an append-only SQLite audit trail

• An agent loop that ties all three together

• A review surface that forces human attestation before downstream systems receive anything

By the end you'll have a working repo you can drop into any agent project. The full code is at github.com/dannwaneri/production-safe-agent-loop.

Table of Contents

1. Why This Keeps Happening

2. Prerequisites

3. Phase 1: Define Done Before You Build

4. Phase 2: Enforce Done at Runtime

5. Phase 3: Record Everything

6. Phase 4: The Loop That Respects Its Boundaries

7. Phase 5: The Review Surface

8. Phase 6: A Real Example, SEO Audit Agent

9. Pluggable LLM Client

10. Running the Tests

11. What You've Built

12. Next Steps

Why This Keeps Happening

The math that got companies into trouble was simple. A chatbot costs roughly 0.04 USD per interaction. An orchestrated multi-agent workflow costs 1.20 USD. That's a 30x multiplier — and production benchmarks show it can reach 70x on complex tasks.

The problem isn't that agents are expensive. The problem is that most teams budgeted for chatbot costs and deployed agent architectures. Gartner found the token consumption gap between pilot chatbots and production agent workflows sits at 5-30x. The FinOps Foundation's 2026 State of FinOps report found 73% of enterprises say AI costs exceeded original projections.

The mechanism is straightforward once you see it. When an agent fails a task and retries, it doesn't start fresh. It re-reads the entire context window — every prior failed attempt — before trying again. Iteration one costs 100 tokens. Iteration two costs 200. Iteration ten costs thousands. You're paying for every failure, over and over, in milliseconds.

# This is the entire problem in three lines
while True:
    result = agent.run(task)
    # done when...?

That question mark is where the money goes.

The other thing making it worse: agents don't fail loudly. Traditional code hits an undefined state and crashes. An LLM hits ambiguity and tries to be helpful. It retries. It reformats the tool call. It spins up a verification agent. The verification agent finds something. A correction agent fires. Nobody defined what "correct" means. The loop looks beautiful on every dashboard you have — activity, tool calls, completion rate — while quietly burning through your budget.

Gartner predicts that 40% of agentic projects will be scrapped by 2027 due to economic failure. Most of that failure is preventable. Not with better models, but with exit conditions.

Prerequisites

• Python 3.10+

• An Anthropic API key (or any provider — more on that later)

• Basic familiarity with Python classes and SQLite

git clone https://github.com/dannwaneri/production-safe-agent-loop
cd production-safe-agent-loop
pip install -r requirements.txt
export ANTHROPIC_API_KEY=sk-...

Phase 1: Define Done Before You Build

The most expensive mistake in agent development isn't a bad model choice or a missing retry limit. It's starting the build before you can answer one question in one sentence:

What does done look like?

Most teams can't answer it. Not because they're careless, but because nothing forces them to before they open the terminal. The spec writer is that forcing function.

# spec_writer.py
from spec_writer import SpecWriter

spec = SpecWriter(db_path="spec.db").run()

When you call .run(), it won't return until you've answered three questions:

1. What does this do?

2. What does this NOT do?

3. What does done look like in one sentence?

The third question is the one that matters. It's also the hardest. "The agent audits the site" is not an answer. "The agent crawls the target URL, extracts all <title> and <meta description> tags, flags any missing or over-length, and stops" is an answer. One of those gives the circuit breaker something to enforce.

The spec stores to SQLite and returns a SpecResult dataclass with a session_id. That ID becomes the thread connecting your spec, your ledger rows, and your loop result. One session, traceable end to end.

@dataclass(frozen=True)
class SpecResult:
    what_it_does: str
    what_it_does_not: str
    done_looks_like: str
    session_id: str

frozen=True matters. The spec is a commitment, not a draft. Once it's written, the loop runs against it. No mid-run revisions.

For testing, SpecWriter accepts injectable input_fn and output_fn callables. No stdin monkey-patching required. See tests/test_spec_writer.py for working examples — the suite uses a small scripted_input helper that returns answers from a generator, and writes to a per-test SQLite file via pytest's tmp_path fixture. SQLite's :memory: isn't safe here, because SpecWriter opens a fresh connection per method and each :memory: connection is its own isolated database.

Phase 2: Enforce Done at Runtime

Defining the exit condition upstream is discipline. The circuit breaker is enforcement.

# circuit_breaker.py
from circuit_breaker import CircuitBreaker, CircuitBreakerError

breaker = CircuitBreaker(turn_limit=5, token_limit=15000)
breaker.check(turn_count, accumulated_tokens)  # raises on breach

Two ceilings. Both hard.

turn_limit caps how many times the loop can call the LLM. token_limit caps total token consumption across all turns. Either one tripping raises CircuitBreakerError immediately.

The boundary is strict: turn_count == turn_limit is allowed. turn_count == turn_limit + 1 trips. No grace periods or warnings. A hard stop forces a human checkpoint.

from dataclasses import dataclass


@dataclass
class CircuitBreakerError(Exception):
    reason: str          # "turn_ceiling" or "token_ceiling"
    turn_count: int
    accumulated_tokens: int

    def __post_init__(self) -> None:
        super().__init__(
            f"circuit breaker tripped: {self.reason} "
            f"(turn={self.turn_count}, tokens={self.accumulated_tokens})"
        )


class CircuitBreaker:
    def __init__(self, turn_limit: int = 5, token_limit: int = 15000) -> None:
        self.turn_limit = turn_limit
        self.token_limit = token_limit

    def check(self, turn_count: int, accumulated_tokens: int) -> None:
        if turn_count > self.turn_limit:
            self._trip("turn_ceiling", turn_count, accumulated_tokens)
        if accumulated_tokens > self.token_limit:
            self._trip("token_ceiling", turn_count, accumulated_tokens)

    def _trip(self, reason: str, turn_count: int, accumulated_tokens: int) -> None:
        print(
            "\n=== CIRCUIT BREAKER CHECKPOINT ===\n"
            f"reason         : {reason}\n"
            f"turn_count     : {turn_count} / limit {self.turn_limit}\n"
            f"tokens_used    : {accumulated_tokens} / limit {self.token_limit}\n"
            "action         : halt loop, surface to human reviewer\n"
            "=================================="
        )
        raise CircuitBreakerError(
            reason=reason,
            turn_count=turn_count,
            accumulated_tokens=accumulated_tokens,
        )

CircuitBreakerError is an exception, not a return code. That's intentional. A return code can be ignored. An uncaught exception can't. Silent breach is impossible. The human-readable checkpoint banner is printed to stdout by _trip() before the exception is raised, so even if a caller swallows the exception the operator still sees state.

The critical rule: call .check() before every LLM call, not after. Post-flight checking means you've already burned the tokens before you knew the limit was exceeded.

# Wrong — post-flight
result = client.messages.create(...)
breaker.check(turn_count, accumulated_tokens)  # too late

# Right — pre-flight
breaker.check(turn_count, accumulated_tokens)  # raises before any spend
result = client.messages.create(...)

The defaults (5 turns, 15,000 tokens) match a tight tutorial demo. Your production budget is different. Tune at instantiation:

# Production example — tighter token budget, more turns
breaker = CircuitBreaker(turn_limit=10, token_limit=50000)

Phase 3: Record Everything

The circuit breaker protects your bank account. The ledger protects your understanding of what happened.

Most teams log for debugging — they want to know what went wrong after it went wrong. The ledger has a different purpose. It's governance. Every row is proof that the loop stayed within its boundaries, or didn't, and exactly when.

# ledger.py
from ledger import Ledger

ledger = Ledger(db_path="ledger.db")
ledger.write(
    session_id=spec.session_id,
    turn_count=1,
    state_origin="llm",
    input_str=task,
    token_delta=523,
    execution_time_ms=1240,
    pass_fail=True,
)

One row per turn. Append-only, no updates, and no deletes. The immutability is the point: a ledger you can edit isn't a ledger, it's a notebook.

The schema:

CREATE TABLE IF NOT EXISTS ledger (
    id                 INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id         TEXT    NOT NULL,
    turn_count         INTEGER NOT NULL,
    state_origin       TEXT    NOT NULL,
    input_hash         TEXT    NOT NULL,
    token_delta        INTEGER NOT NULL,
    execution_time_ms  INTEGER NOT NULL,
    pass_fail          INTEGER NOT NULL,  -- 1=pass, 0=fail
    breach_reason      TEXT,              -- NULL unless circuit breaker fired
    created_at         TEXT    NOT NULL   -- ISO 8601, UTC
);
CREATE INDEX IF NOT EXISTS idx_ledger_session ON ledger(session_id);

The index makes get_session(session_id) — the primary read path — a constant-time lookup as the ledger grows.

Three decisions worth explaining:

1. input_hash not input_text. The raw input string never persists. Only its SHA-256 hash does. There are two benefits to this: identical inputs across runs are detectable, and PII never enters the audit trail.

2. pass_fail as INTEGER not BOOLEAN. SQLite has no boolean type. 1 and 0 are canonical. Clean Python ergonomics at the API edge, correct SQL types on disk.

3. created_at as datetime.now(timezone.utc).isoformat(). datetime.utcnow() was deprecated in Python 3.12. Timezone-aware timestamps avoid the footgun in any system that crosses timezones.

Retrieve by session:

rows = ledger.get_session(spec.session_id)
for row in rows:
    print(f"Turn {row.turn_count}: {'PASS' if row.pass_fail else 'FAIL'} "
          f"| {row.token_delta} tokens | {row.execution_time_ms}ms")

Phase 4: The Loop That Respects Its Boundaries

The agent loop wires the three primitives together. It's the only component that calls the LLM. Everything else is local.

# agent_loop.py
from agent_loop import AgentLoop

loop = AgentLoop(spec, breaker, ledger, client)
result = loop.run(task)
# LoopResult(success, turns, total_tokens, session_id, breach_reason)

The anatomy of a turn, in order:

1. circuit_breaker.check(turn_count, accumulated_tokens) — raises if either ceiling is exceeded

2. client.messages.create(...) — the actual LLM call

3. ledger.write(...) — one row, append-only

4. If stop_reason == "end_turn", return. Otherwise loop.

Pre-flight checking before every LLM call, with no exceptions.

def run(self, task: str) -> LoopResult:
    session_id = self.spec.session_id
    messages: list[dict] = [{"role": "user", "content": task}]
    turn = 0
    total_tokens = 0

    try:
        while True:
            turn += 1
            self.circuit_breaker.check(turn, total_tokens)

            started = time.perf_counter()
            response = self.client.messages.create(
                model=self.model,
                max_tokens=self.max_tokens,
                system=self._system_prompt(),
                messages=messages,
            )
            elapsed_ms = int((time.perf_counter() - started) * 1000)

            turn_tokens = (
                getattr(response.usage, "input_tokens", 0)
                + getattr(response.usage, "output_tokens", 0)
            )
            total_tokens += turn_tokens

            text = self._text_from(response)
            messages.append({"role": "assistant", "content": text})

            self.ledger.write(
                session_id=session_id,
                turn_count=turn,
                state_origin="llm",
                input_str=task,
                token_delta=turn_tokens,
                execution_time_ms=elapsed_ms,
                pass_fail=True,
            )

            if getattr(response, "stop_reason", "end_turn") == "end_turn":
                return LoopResult(
                    success=True,
                    turns=turn,
                    total_tokens=total_tokens,
                    session_id=session_id,
                )

            messages.append({"role": "user", "content": "continue"})

    except CircuitBreakerError as err:
        self.ledger.write(
            session_id=session_id,
            turn_count=turn,
            state_origin="circuit_breaker",
            input_str=task,
            token_delta=0,
            execution_time_ms=0,
            pass_fail=False,
            breach_reason=err.reason,
        )
        return LoopResult(
            success=False,
            turns=turn,
            total_tokens=total_tokens,
            session_id=session_id,
            breach_reason=err.reason,
        )

def _system_prompt(self) -> str:
    return (
        "You are an agent working on a tightly-scoped task.\n\n"
        f"What this does: {self.spec.what_it_does}\n"
        f"What this does NOT do: {self.spec.what_it_does_not}\n"
        f"Done looks like: {self.spec.done_looks_like}\n"
    )

@staticmethod
def _text_from(response) -> str:
    content = getattr(response, "content", None)
    if not content:
        return ""
    block = content[0]
    return getattr(block, "text", "") or ""

A few choices worth calling out in this body:

• The whole while True: is wrapped in one try/except CircuitBreakerError. The check happens at the top of every turn, so a breach is caught the same way whether it fires on turn 1 or turn 6.

• input_str=task on every ledger row — the original task, not the last assistant message. The input_hash column then groups rows that share the same starting input across the run.

• pass_fail=True for every LLM turn that returns, False only on breach. The pass/fail flag tracks whether the loop reached the row legitimately, not whether the model's output was good. Quality scoring is a separate concern.

• _system_prompt() uses all three spec fields, not just done_looks_like. The model needs the negative scope (what_it_does_not) at least as much as the positive scope.

• time.perf_counter() not time.time() — monotonic, immune to wall-clock adjustments mid-run.

LoopResult.session_id is inherited from spec.session_id. The ledger rows tie back to the spec without a join. One session ID, one traceable run, start to finish.

Phase 5: The Review Surface

The circuit breaker protects your bank account. The ledger records what happened. But neither tells you whether what happened matched what you promised.

That gap is where bad loops get approved. Polished output, green dashboard, missed commitment. A reviewer sees the artifact, decides it looks acceptable, and signs off. Nobody asked whether the original promise was kept.

The review surface closes that gap. It reads the session from SQLite, assembles the five-element frame, and forces a comparison before anything downstream receives the output.

from review_surface import ReviewSurface

rs = ReviewSurface(spec_db_path="spec.db", ledger_db_path="ledger.db")
print(rs.render(session_id))

Here's the five-element frame, in order:

1. Original promise — pulled from the spec table: what it does, what it doesn't do, what done looks like

2. Acceptance criteria — the done_looks_like field rendered as the explicit benchmark

3. Diff — first turn input vs final turn output, turns completed, total tokens, whether the loop breached

4. Evidence — all ledger rows for the session: turn-by-turn pass/fail, token delta, execution time

5. Unresolved assumptions — derived from breach rows and failed turns. Empty when clean.

When the reviewer is satisfied, they attest:

attestation = rs.attest(
    session_id=result.session_id,
    reviewer="daniel",
    notes="Output matches spec. Approved."
)
print(attestation.frame_hash)

.attest() writes to the attestations table in ledger.db. The frame_hash is a SHA-256 of the canonical frame data — deterministic across reviewers attesting the same session. It's the audit receipt. It proves the reviewer saw the exact frame as rendered, not a summary or a paraphrase.

Approval confirms the process ran. Attestation confirms the reviewer compared output to commitment. When the loop touches something regulated, those are different legal documents.

@dataclass(frozen=True)
class ReviewFrame:
    session_id: str
    original_promise: SpecResult
    acceptance_criteria: str
    diff: DiffResult
    evidence: tuple  # tuple[LedgerRow, ...]
    unresolved_assumptions: tuple  # tuple[str, ...]
    created_at: str

ReviewFrame is frozen for the same reason SpecResult is — the frame is evidence, not a draft. evidence and unresolved_assumptions are tuples because lists aren't hashable and frozen dataclasses need hashable fields.

The full end-to-end flow with the review surface lives in examples/review_example.py in the repo. Run it after any completed session: it renders the five-element frame, prompts for attestation, and writes the receipt if you approve.

The loop runs to you. Downstream systems get nothing until someone signs.

Phase 6: A Real Example — SEO Audit Agent

The pattern only makes sense against a real problem. This is the same agent architecture behind my seo-agent project.

SEO audits have a natural cadence: crawl, surface what's broken, fix, wait for reindex. Running the agent continuously doesn't change that cadence. It just burns tokens in the empty space between the moments that matter. A cron job wired to the loop is the honest architecture.

# examples/seo_audit_example.py
import requests
from bs4 import BeautifulSoup
import anthropic
from spec_writer import SpecWriter
from circuit_breaker import CircuitBreaker
from ledger import Ledger
from agent_loop import AgentLoop

def crawl_url(url: str) -> str:
    response = requests.get(url, timeout=10)
    soup = BeautifulSoup(response.text, "html.parser")
    title = soup.find("title")
    meta_desc = soup.find("meta", attrs={"name": "description"})
    h1_tags = soup.find_all("h1")
    return (
        f"URL: {url}\n"
        f"Title: {title.text if title else 'MISSING'}\n"
        f"Meta description: "
        f"{meta_desc['content'] if meta_desc else 'MISSING'}\n"
        f"H1 count: {len(h1_tags)}\n"
        f"H1 tags: {[h.text[:50] for h in h1_tags]}"
    )

def run_seo_audit(url: str) -> None:
    # Step 1: Define done before the loop starts
    spec = SpecWriter(db_path="spec.db").run()

    # Step 2: Initialise circuit breaker and ledger
    breaker = CircuitBreaker(turn_limit=5, token_limit=15000)
    ledger = Ledger(db_path="ledger.db")
    client = anthropic.Anthropic()

    # Step 3: Crawl the URL
    site_data = crawl_url(url)

    # Step 4: Run the loop
    # AgentLoop catches CircuitBreakerError internally and returns
    # LoopResult(success=False, breach_reason=...). Branch on the
    # result — do NOT wrap loop.run() in try/except CircuitBreakerError.
    loop = AgentLoop(spec, breaker, ledger, client)
    result = loop.run(
        f"Audit this page for SEO issues:\n\n{site_data}"
    )

    # Step 5: Print the ledger
    print(f"\nResult: {'SUCCESS' if result.success else 'BREACH'}")
    if not result.success:
        print(f"Breach reason: {result.breach_reason}")
    print(f"Turns: {result.turns} | Tokens: {result.total_tokens}")
    print("\nAudit trail:")
    for row in ledger.get_session(result.session_id):
        status = "PASS" if row.pass_fail else "FAIL"
        print(f"  Turn {row.turn_count}: {status} | "
              f"{row.token_delta} tokens | {row.execution_time_ms}ms")

if __name__ == "__main__":
    import sys
    run_seo_audit(sys.argv[1] if len(sys.argv) > 1 else "https://example.com")

Run it:

python examples/seo_audit_example.py https://yourdomain.com

The spec writer prompts you. The loop runs, the circuit breaker fires if the limits are exceeded, and the ledger records every turn. The output lands in front of you and you decide what to fix.

The loop runs to you, not into a void.

Pluggable LLM Client

The loop works with any client that satisfies the LLMClient protocol (Anthropic by default). Bring your own via a ~20-line adapter.

# agent_loop.py
from typing import Protocol, runtime_checkable


@runtime_checkable
class MessagesEndpoint(Protocol):
    def create(self, *, model: str, max_tokens: int,
               system: str, messages: list) -> object: ...


@runtime_checkable
class LLMClient(Protocol):
    messages: MessagesEndpoint

messages is an instance attribute (not a nested class) because that's how the real Anthropic SDK exposes it — anthropic.Anthropic().messages.create(...). Modeling it as a nested class would mean the real client wouldn't satisfy the Protocol. The @runtime_checkable decorator lets you sanity-check conformance with isinstance(client, LLMClient), and the repo's test suite uses exactly that assertion against the FakeClient test double.

Here's an OpenAI adapter example (This is illustrative. A production adapter would also map streaming, tool-use, and error shapes.):

# openai_adapter.py — illustrative pseudocode, not production-ready.
from openai import OpenAI as _OpenAI


class _MessagesAdapter:
    def __init__(self, client):
        self._client = client

    def create(self, *, model, max_tokens, system, messages):
        completion = self._client.chat.completions.create(
            model=model,
            max_tokens=max_tokens,
            messages=[{"role": "system", "content": system}] + messages,
        )
        # Reshape OpenAI's response into the Anthropic-shaped surface
        # AgentLoop reads: response.usage.{input,output}_tokens,
        # response.content[0].text, response.stop_reason.
        return _adapt_response(completion)


class OpenAIAdapter:
    def __init__(self, api_key: str):
        self._client = _OpenAI(api_key=api_key)
        self.messages = _MessagesAdapter(self._client)  # instance attr, not a nested class

The adapter pattern is worth teaching explicitly. Provider APIs don't share a shape. Anthropic puts system at the top level. OpenAI puts it inside the messages array. An adapter shim is ~20 lines and makes the loop provider-agnostic without rewriting anything. Note that self.messages is assigned in __init__ so it's a real attribute on each adapter instance, the same shape as the actual SDK.

Running the Tests

python -m pytest tests/

With coverage:

python -m coverage run --source=circuit_breaker,ledger,spec_writer,agent_loop,review_surface -m pytest tests/
python -m coverage report -m

80 tests, 100% coverage on all five core modules. The loop is exercised against a FakeClient test double defined inline in tests/test_agent_loop.py. It satisfies the LLMClient protocol via duck typing: messages is set to self, so client.messages.create(...) routes back to the same object and ships with scripted responses for each test scenario. Clone the repo and run pytest to see all 80 tests pass without touching the network or needing an API key.

circuit_breaker.py has 100% coverage — no untested paths. It's the financial safety component. Every path through it is exercised.

What You've Built

In this tutorial, you've build five small primitives, each independently usable.

The pattern: upstream discipline defines the boundaries. Downstream enforcement breaks the circuit. Neither trusts the model to police itself.

A loop that runs without an exit condition isn't autonomous. It's a billing event waiting to happen.

Define what done looks like before you start. That's the job, and always has been.

Next Steps

The repo is at github.com/dannwaneri/production-safe-agent-loop.

There are three natural extensions if you want to go further:

1. Graduation to Distributed Systems

The SQLite ledger works for isolated sequential loops. The moment you run multiple agents against shared state, you need serializable isolation — concurrent writes to flat JSON corrupt silently. The README documents the three tipping points where a flat ledger needs to graduate.

2. Cryptographic Signing

For compliance-scale systems where the auditor wasn't present when the loop ran, SQLite rows aren't enough. A database admin can run an UPDATE query. Ed25519 signing wraps each ledger row in a receipt that proves the log wasn't altered after execution. But that's a different tutorial.

Wiring a Cron Job

The honest architecture for the SEO audit agent isn't 24/7 autonomous operation. It's a cron job that runs on schedule, surfaces what's broken, and stops. 0 3 * * 2 python examples/seo_audit_example.py https://yourdomain.com is the whole thing. The loop runs to you, not into a void.

If you need this architecture built for your own stack (circuit breakers, audit trails, production-safe agent loops), I do freelance work. dannwaneri.com/ai-agents/

Each new generation seemed to confirm a simple belief: bigger models lead to better capabilities. Yet there was one area where progress appeared frustratingly limited. When problems required multiple steps of reasoning, language models often struggled in ways that were difficult to ignore.

A math word problem, a common sense question, or a symbolic puzzle could expose a surprising gap between fluent language generation and genuine problem solving. Models could frequently produce confident answers, but confidence alone wasn't enough. The challenge was whether they could reason through a problem before arriving at an answer.

Against this backdrop, the paper Chain-of-Thought Prompting Elicits Reasoning in Large Language Models introduced an idea that was both simple and unexpected. Rather than asking a model to produce an answer immediately, the authors encouraged it to work through intermediate reasoning steps first.

What followed was one of the most influential discoveries in modern AI research: many reasoning abilities that appeared absent in large language models weren't necessarily missing. In many cases, they simply hadn't been elicited in the right way.

This paper went on to reshape how researchers think about prompting, reasoning, and the capabilities of large language models. More importantly, it laid the intellectual foundation for many of the reasoning-oriented techniques and systems that emerged in the years that followed.

Paper Overview

In this article, we'll explore the paper Chain-of-Thought Prompting Elicits Reasoning in Large Language Models, published by researchers at Google Research in 2022.

This paper introduced one of the most influential ideas in modern AI: Chain-of-Thought (CoT) Prompting. At a time when researchers were focused on scaling language models to ever-larger sizes, this study revealed that performance improvements were not always about building bigger models. Sometimes, the key was changing how we communicate with them.

The paper investigates a simple but powerful question: what happens if a language model is encouraged to show its reasoning process before giving an answer? Instead of responding directly, the model is guided to generate intermediate reasoning steps that lead to the final solution.

What makes this paper historically important is that it changed how researchers think about reasoning in large language models. The authors demonstrated that many reasoning capabilities can be unlocked through prompting alone, without additional training, fine-tuning, or architectural modifications.

The impact of this idea quickly extended beyond arithmetic reasoning. It influenced a new generation of research on reasoning, including Self-Consistency, Process Supervision, Verification-based methods, and the reasoning-oriented models that followed in subsequent years.

In many ways, this paper marked a shift from asking language models what the answer is to asking them how they arrived at the answer.

Here's the original paper if you'd like to explore it directly:

Chain-of-Thought Prompting Elicits Reasoning in Large Language Models

And here's a quick infographic of what we'll cover throughout this review.

Table of Contents:

• Abstract

• Introduction

• Chain-of-Thought Prompting

• Arithmetic Reasoning

• Results

• Ablation Study

• Robustness of Chain-of-Thought Prompting

• Common Sense Reasoning

• Symbolic Reasoning

• Discussion

• Related Work

• Conclusion

• Resources

Prerequisites

To get the most out of this breakdown, it helps to already be familiar with a few foundational ideas and the evolution of large language models that led to Chain-of-Thought prompting.

Reading the previous reviews in this series will be especially helpful:

• AI Paper Review: Improving Language Understanding by Generative Pre-Training (GPT-1)

• AI Paper Review: Language Models are Unsupervised Multitask Learners (GPT-2)

• AI Paper Review: Language Models are Few-Shot Learners (GPT-3)

• AI Paper Review: Training Language Models to Follow Instructions with Human Feedback (InstructGPT)

The GPT-3 review is particularly important because the Chain-of-Thought paper builds directly on one of GPT-3's most surprising capabilities: in-context learning. Rather than changing the model architecture or retraining the model, the authors discovered that reasoning performance could be dramatically improved simply by changing how examples were presented in the prompt.

It also helps to have:

• A general understanding of natural language processing (NLP) and large language models

• A basic understanding of Transformer-based autoregressive models

• Familiarity with prompting, few-shot learning, and in-context learning

• A high-level understanding of how language models generate text token by token

• General machine learning concepts such as training, inference, scaling laws, and model evaluation

• Some exposure to reasoning tasks, logic problems, and mathematical word problems

• A basic understanding of benchmark datasets and model performance evaluation

You don't need a deep background in mathematics or machine learning research to follow this article.

I'll keep the explanations intuitive and practical, focusing on why Chain-of-Thought prompting became one of the most influential reasoning techniques in modern AI and how a simple prompting strategy changed the way researchers think about language model reasoning.

Abstract

One of the long-standing challenges for large language models has been reasoning. While these models can generate fluent text and answer a wide variety of questions, they often struggle when a task requires multiple logical steps.

This paper introduces a remarkably simple idea to address that limitation: instead of prompting a model with only questions and answers, you should provide examples that also include the intermediate reasoning steps leading to the solution.

The authors call this approach Chain-of-Thought (CoT) Prompting. By showing a model a few demonstrations of step-by-step reasoning, they find that sufficiently large language models can generate their own reasoning chains and solve complex problems more effectively. Importantly, this improvement doesn't require additional training or fine-tuning, only a different style of prompting.

Through experiments on arithmetic, common sense, and symbolic reasoning tasks, the paper demonstrates that chain-of-thought prompting consistently improves performance. The gains become especially pronounced at larger model scales, suggesting that reasoning abilities emerge naturally as models grow and are given the right prompting strategy.

The paper's most striking result comes from the GSM8K math benchmark, where PaLM 540B, using only eight chain-of-thought examples, achieved state-of-the-art performance and even surpassed a fine-tuned GPT-3 system equipped with a verifier. This finding revealed that prompting alone could unlock reasoning capabilities that standard prompting often fails to expose.

The figure below compares standard prompting with Chain-of-Thought (CoT) prompting using a simple arithmetic example.

Source: Chain-of-Thought Prompting Elicits Reasoning in Large Language Models

In standard prompting, the model is shown question–answer pairs and is expected to produce an answer directly, which can lead to mistakes on multi-step problems.

In Chain-of-Thought prompting, the examples include intermediate reasoning steps before the final answer. When faced with a new problem, the model follows a similar step-by-step process, arriving at the correct solution.

This paper shows that providing reasoning demonstrations can substantially improve performance on arithmetic, common sense, and symbolic reasoning tasks, particularly in large language models.

Introduction

By 2022, large language models had already transformed natural language processing. Models such as GPT-3 demonstrated that scaling model size could unlock impressive capabilities, from text generation to few-shot learning.

But there was an important limitation: larger models weren't necessarily better at reasoning. Tasks that required multi-step arithmetic, common sense inference, or symbolic manipulation remained surprisingly difficult, even for some of the largest models available.

The authors begin by observing two promising research directions. The first comes from prior work showing that reasoning tasks can benefit from natural language explanations or intermediate solution steps. Instead of jumping directly to an answer, a model can generate a rationale that mirrors how a human might solve the problem.

The second direction is few-shot prompting, where a model learns a task from a handful of examples provided in the prompt, eliminating the need for task-specific fine-tuning.

Still, both approaches have drawbacks. Training models on large collections of human-written rationales is expensive and time-consuming, while standard few-shot prompting often struggles on tasks that require genuine reasoning.

The key insight of this paper was to combine the strengths of both ideas. Rather than providing only input-output examples, the prompt includes an additional component: the reasoning process itself. Each example follows the structure of input → chain of thought → output.

This simple modification led to Chain-of-Thought Prompting. By exposing intermediate reasoning steps, the model is encouraged to break complex problems into smaller, more manageable stages before arriving at a final answer.

To evaluate the idea, the authors tested chain-of-thought prompting across arithmetic, common sense, and symbolic reasoning benchmarks. The results showed substantial improvements over standard prompting, with some gains being remarkably large.

Chain-of-Thought Prompting

At the heart of this paper is a simple observation about how humans solve difficult problems. When faced with a multi-step reasoning task, we rarely jump directly to the answer. Instead, we break the problem into smaller pieces, solve each intermediate step, and gradually work toward a conclusion. The authors argued that large language models could benefit from a similar process.

This idea led to Chain-of-Thought (CoT) Prompting, where examples in the prompt included not only the question and answer, but also the reasoning steps connecting them. By seeing a few demonstrations of this reasoning process, sufficiently large language models learned to generate their own chains of thought before producing a final answer.

The significance of this approach extends beyond improving accuracy. First, it allows complex problems to be decomposed into manageable intermediate steps, making multi-step reasoning easier to perform.

Second, the generated reasoning process offers a degree of interpretability, giving researchers and users a glimpse into how the model arrived at its answer. While these reasoning traces don't fully reveal the model's internal computations, they can help identify where mistakes occur.

Another important aspect of chain-of-thought prompting is its generality. The authors proposed it not as a solution for a single benchmark, but as a broad reasoning framework that can be applied to arithmetic problems, common sense reasoning tasks, symbolic manipulation, and potentially many other challenges that require sequential reasoning.

Perhaps most importantly, this capability can be elicited from existing language models through prompting alone, without additional training or architectural modifications.

This section establishes the paper's central claim: reasoning abilities don't necessarily require new model architectures or specialized fine-tuning. In sufficiently large language models, these capabilities can emerge when the model is guided to generate intermediate reasoning steps rather than being asked to produce an answer immediately.

Arithmetic Reasoning

The authors begin their empirical evaluation with arithmetic reasoning, a domain that had long exposed a weakness of large language models.

Although solving math word problems is relatively straightforward for humans, it often requires a sequence of intermediate calculations and logical deductions.

Previous research had shown that even large language models struggled with these tasks, making arithmetic reasoning an ideal setting for testing whether chain-of-thought prompting could genuinely improve reasoning ability.

To evaluate their approach, the authors selected five established benchmarks covering a variety of math word problems. These datasets differ in style and difficulty, ranging from straightforward arithmetic questions to more complex problems that require multiple reasoning steps before arriving at a solution. Together, they provide a broad picture of how well language models handle mathematical reasoning.

The experiments compare two prompting strategies. The first is standard few-shot prompting, where the model is shown examples consisting only of questions and their corresponding answers. This was the dominant prompting approach at the time and serves as the baseline throughout the paper.

The second is chain-of-thought prompting, where each example is expanded to include the intermediate reasoning steps that connect the question to the final answer.

To ensure a fair comparison, the authors manually created a small set of eight reasoning demonstrations and reused them across the arithmetic benchmarks. Importantly, these examples weren't heavily optimized or engineered for specific datasets. Instead, they were intended to test whether a modest number of natural reasoning demonstrations could reliably encourage models to reason through new problems on their own.

The study also evaluates a diverse collection of language models, including GPT-3, LaMDA, PaLM, UL2, and Codex, spanning model sizes from hundreds of millions to hundreds of billions of parameters. This broad range allowed the authors to examine not only whether chain-of-thought prompting works, but also how its effectiveness changes as models become larger.

With this experimental framework in place, the paper investigated a central question: can providing a few examples of step-by-step reasoning enable large language models to solve mathematical problems that standard prompting struggles to handle?

Results

The arithmetic reasoning experiments revealed that the success of chain-of-thought prompting depends heavily on model scale.

One of the clearest patterns across the benchmarks was that smaller models gained little benefit from generating reasoning steps. In some cases, their performance even deteriorated because the models produced explanations that sounded plausible but were logically flawed.

The advantages of chain-of-thought prompting only became apparent once the models reached very large scales, suggesting that the ability to effectively use intermediate reasoning steps is itself an emergent capability.

Another important observation was that the benefits of chain-of-thought prompting grew as problems became more challenging. On simpler tasks that required only a single reasoning step, standard prompting was already sufficient and the additional reasoning process provided little value.

But as the complexity of the problems increased, the gap between standard prompting and chain-of-thought prompting widened substantially. The GSM8K benchmark provides the strongest example of this trend, where the largest GPT and PaLM models more than doubled their performance when allowed to reason step by step.

Perhaps the most significant result is that chain-of-thought prompting enabled large language models to compete with, and in some cases surpass, specialized systems trained directly for these tasks.

Using only a handful of reasoning demonstrations, PaLM 540B established new state-of-the-art results on several arithmetic benchmarks, despite relying solely on prompting rather than task-specific fine-tuning. This outcome challenged the prevailing assumption that strong performance on reasoning tasks necessarily required dedicated training datasets and specialized models.

To better understand these improvements, the authors manually inspected the reasoning traces generated by the models. When the model arrived at the correct answer, the reasoning process was usually correct as well, indicating that the model was often following a coherent sequence of logical steps rather than guessing the final answer.

Even among incorrect predictions, many reasoning chains were largely accurate and failed only because of small mistakes such as arithmetic slips, incorrect symbol mappings, or a missing intermediate step. More serious failures tended to arise from misunderstanding the problem itself or producing incoherent reasoning.

The error analysis also offered an explanation for why larger models benefited more from chain-of-thought prompting. Comparing PaLM 62B with PaLM 540B showed that increasing scale reduced many of the semantic misunderstandings and incomplete reasoning patterns that appeared in smaller models.

In other words, larger models were not merely generating longer explanations. They were producing reasoning chains that were more logically complete and more faithful to the underlying problem.

Ablation Study

Before diving into this section, it's worth briefly explaining what an ablation study is. In machine learning research, an ablation study systematically removes or modifies parts of a method to determine which components are actually responsible for its performance. Rather than asking whether a method works, an ablation study asks why it works.

In this paper, the authors use ablation experiments to identify which aspects of Chain-of-Thought prompting contribute most to its reasoning improvements.

After demonstrating that chain-of-thought prompting improved reasoning performance, the authors turned to a more fundamental question: why does it work? Simply observing higher accuracy isn't enough. To understand the source of these gains, they designed a series of ablation experiments that isolated different aspects of the prompting strategy.

One possible explanation is that chain-of-thought prompting helps because it encourages the model to generate mathematical equations before producing an answer. If this were true, then the natural language reasoning itself might not be necessary.

To test this idea, the authors replaced the reasoning steps with equations alone. The results showed that this approach provides only limited benefits on complex benchmarks such as GSM8K. While equations can help with simpler problems, they are often insufficient for tasks that require understanding the meaning of the question before translating it into mathematical operations. This suggests that the value of chain-of-thought prompting comes from more than symbolic calculation.

The authors then examined another hypothesis: perhaps chain-of-thought prompting succeeds simply because it allows the model to generate more tokens and therefore spend more computation on difficult problems.

To isolate this factor, they created a prompt that produces additional tokens without any meaningful reasoning content. Performance remained close to the standard prompting baseline, indicating that extra computation alone doesn't explain the observed improvements. What mattered wasn't the number of intermediate tokens, but the reasoning expressed within them.

A third possibility was that chain-of-thought prompts merely activated relevant knowledge already stored in the model. If that were the case, the reasoning steps wouldn't need to appear before the answer.

The authors tested this by moving the reasoning process to after the final answer. Once again, performance largely fell back to the baseline. This result suggested that the sequence of reasoning steps plays an active role in helping the model arrive at the correct solution rather than simply serving as an explanation after the fact.

Taken together, these experiments strengthen the paper's central argument. The success of chain-of-thought prompting can't be explained by equation generation, additional computation, or easier access to stored knowledge alone.

Instead, the evidence points toward the reasoning process itself as the critical ingredient. The intermediate steps aren't merely decorative explanations. They appear to guide the model through a sequence of decisions that makes complex problem solving more effective.

Robustness of Chain-of-Thought Prompting

One of the long-standing concerns with prompting methods is their sensitivity to the examples included in the prompt. Small changes in wording, example selection, or even the order of examples can sometimes produce noticeably different results.

Once they established that chain-of-thought prompting improves reasoning performance, the authors investigated whether these gains were robust or whether they depended on a particular set of carefully crafted demonstrations.

To answer this question, the researchers asked multiple authors of the paper to independently write reasoning traces for the same examples. They also experimented with a more concise writing style and tested prompts built from entirely different sets of examples.

The goal was to determine whether chain-of-thought prompting was succeeding because of a specific wording choice or because the underlying reasoning structure was genuinely useful.

The results provided reassuring evidence that the technique isn't tied to a particular author, writing style, or collection of exemplars. While some variation in performance naturally appeared across different prompts, every version of chain-of-thought prompting consistently outperformed standard prompting by a substantial margin. Whether the reasoning steps were detailed or concise, manually written or drawn from an independent dataset, the overall pattern remained remarkably stable.

The authors further broadened their analysis by varying the order and number of exemplars used in the prompt. Once again, the central finding persisted: although prompt design still influenced performance to some degree, the effectiveness of chain-of-thought prompting didn't depend on a single carefully engineered prompt.

This robustness analysis strengthens one of the paper's most important claims that the success of chain-of-thought prompting isn't an artifact of a particular phrasing or annotation style. Instead, the benefits appear to arise from exposing the model to a reasoning process itself, suggesting that the method captures a more general principle rather than a prompt-specific trick.

Common Sense Reasoning

Up to this point, the paper focused primarily on mathematical reasoning. While the results are impressive, they leave an important question unanswered: is chain-of-thought prompting useful only for arithmetic problems, or can it improve reasoning more broadly?

To investigate this, the authors turned to common sense reasoning tasks. Unlike math problems, these tasks often require background knowledge about the world, an understanding of human behavior, or the ability to connect multiple pieces of information before arriving at a conclusion. In many cases, the challenge isn't performing calculations but reasoning through situations that humans find intuitive.

The evaluation spanned a diverse collection of benchmarks, including common sense question answering, multi-hop reasoning, date understanding, sports-related reasoning, and even tasks that involved converting natural language instructions into robot actions.

Despite their differences, these tasks share a common requirement: solving them often involves a sequence of intermediate inferences rather than an immediate answer.

The results showed that the benefits of chain-of-thought prompting extend well beyond mathematics. Across most benchmarks, models consistently performed better when encouraged to generate intermediate reasoning steps before producing a final answer.

The improvements became particularly noticeable for larger models, suggesting that the same pattern observed in arithmetic reasoning also applies to common sense reasoning.

Some of the strongest gains appeared on tasks that required multi-step inference. On StrategyQA, for example, chain-of-thought prompting enabled PaLM 540B to surpass the previous state of the art. Similarly, on the Sports Understanding benchmark, the model achieved performance that exceeded that of an unaided human sports enthusiast.

These results suggest that the reasoning process encouraged by chain-of-thought prompting can help models connect facts, evaluate plausibility, and navigate more complex decision-making scenarios.

At the same time, the improvements weren't uniform across every dataset. The gains on CommonsenseQA were relatively modest, indicating that not all reasoning tasks benefit equally from explicit reasoning traces. This serves as an early reminder that chain-of-thought prompting isn't a universal solution, even though it consistently proves valuable across a wide range of settings.

More broadly, this section strengthens the paper's central argument by showing that chain-of-thought prompting isn't merely a technique for solving math word problems. Its effectiveness across diverse common sense tasks suggests that the method taps into a more general reasoning capability that emerges in sufficiently large language models.

Symbolic Reasoning

The final evaluation moves away from mathematics and real-world knowledge altogether. Instead, the authors focus on symbolic reasoning tasks, where success depends on following abstract rules rather than recalling facts or performing calculations. These tasks are simple for humans, yet they provide a useful way to test whether language models can consistently apply a sequence of reasoning steps.

To explore this question, the authors designed two controlled tasks. The first required the model to extract and concatenate the last letters of words in a name. The second asked the model to track the state of a coin after a sequence of flips and non-flips.

Although these tasks may appear simple, they required the model to perform precise symbolic manipulations without relying on memorized knowledge about the world.

What made these experiments particularly interesting was the introduction of an out-of-distribution setting. During prompting, the model only saw examples involving short reasoning chains. At evaluation time, it was asked to solve versions of the same tasks that required more steps than any example it had previously encountered.

This setup allowed the authors to test not only whether the model could follow a reasoning procedure, but also whether it could extend that procedure to longer and unfamiliar cases.

The results revealed a familiar pattern. Large models benefitted substantially from chain-of-thought prompting, while smaller models struggled even when the required reasoning process was straightforward.

On the in-domain tasks, where the evaluation closely matched the examples provided in the prompt, the largest models achieved near-perfect performance when guided by chain-of-thought reasoning. This indicated that they could successfully learn and apply the underlying procedure demonstrated in the prompt.

The more revealing results come from the out-of-distribution evaluations. Standard prompting largely fails when the reasoning chain becomes longer than those seen in the examples. In contrast, chain-of-thought prompting enabled performance to improve as model size increased, demonstrating an ability to extend learned reasoning patterns beyond the exact situations shown during prompting.

Although accuracy declines compared to the in-domain setting, the models were still able to generalize in ways that standard prompting couldn't.

This section provided some of the strongest evidence that chain-of-thought prompting is doing more than improving benchmark performance. By helping models apply reasoning procedures to longer and previously unseen inputs, it suggests that the generated reasoning steps serve as a scaffold for systematic problem solving rather than merely a mechanism for producing better answers on familiar examples.

Discussion

The most important contribution of this paper wasn't a new model architecture, a new training objective, or a larger dataset. Instead, it demonstrated that a simple change in prompting could unlock capabilities that standard prompting often failed to reveal.

Across arithmetic, common sense, and symbolic reasoning tasks, chain-of-thought prompting consistently allowed large language models to solve problems that were previously difficult or inaccessible.

A recurring theme throughout the paper was the relationship between reasoning and scale. The authors repeatedly observed that chain-of-thought prompting became effective only once models reached a sufficient size. Smaller models generated fluent reasoning traces, but those traces were often logically inconsistent.

Larger models, in contrast, were able to use intermediate reasoning steps in a way that genuinely improved problem-solving performance.

This finding reinforced a broader lesson emerging from language model research at the time: some capabilities don't appear gradually, but emerge once a model crosses a certain scale threshold.

Perhaps the most intriguing implication was that standard prompting may significantly underestimate what large language models are capable of doing.

Before this work, many reasoning tasks appeared to have reached a performance ceiling. Chain-of-thought prompting revealed that the limitation wasn't always the model itself, but sometimes the way the model was being asked to solve the problem. In that sense, the paper shifted attention from building more capable models to discovering better ways of interacting with the capabilities that already exist within them.

At the same time, the authors were careful not to overstate their conclusions. Although chain-of-thought outputs can resemble human reasoning, the paper doesn't prove that language models reason in the same way humans do. The generated reasoning traces may reflect genuine problem-solving processes, post-hoc rationalizations, or something in between. Determining the relationship between generated reasoning and internal model computation remains an open research question.

The authors also acknowledged several practical limitations. Constructing high-quality reasoning demonstrations can require additional effort, particularly if the approach is extended beyond few-shot prompting.

Also, generating a chain of thought doesn't guarantee that the reasoning itself is correct. Models can still produce convincing but flawed reasoning paths, leading to incorrect answers.

Finally, the strongest benefits appear only in very large models, raising questions about computational cost and whether similar reasoning abilities can be induced in smaller systems.

Viewed from a historical perspective, this paper marked a turning point in research on language model reasoning. Rather than treating reasoning as something that must be explicitly trained into a model, it suggested that reasoning abilities could be elicited through the right prompting strategy.

Many influential ideas that followed, including self-consistency, reasoning supervision, process supervision, and the reasoning-focused models that emerged in later years, can trace part of their intellectual foundation back to the simple insight introduced here: sometimes a model performs better when it's encouraged to show its work.

Related Work

The ideas behind Chain-of-Thought prompting didn't emerge in isolation. Instead, the paper sits at the intersection of two research directions that had been evolving independently for several years.

The first direction focused on helping models solve complex problems through intermediate reasoning steps. Earlier work had already shown that tasks such as mathematical reasoning become easier when a model generates natural language rationales rather than producing an answer directly. Researchers explored methods that trained models to generate explanations, reasoning traces, or intermediate computations before arriving at a final solution.

Other approaches relied on formal symbolic representations, translating problems into structured equations or logical forms. Despite their differences, these efforts shared a common intuition: difficult reasoning tasks are often easier to solve when they're decomposed into smaller steps.

Chain-of-thought prompting inherits this intuition but introduces an important shift. Earlier methods typically required dedicated training procedures, specialized datasets, or task-specific fine-tuning.

In contrast, this paper demonstrated that reasoning traces could be elicited through prompting alone. Rather than teaching a model to reason through additional training, the authors showed that providing a handful of reasoning examples may be enough to unlock capabilities that already exist within sufficiently large language models.

The second research direction concerns prompting itself. Following the success of GPT-3 and few-shot learning, a growing body of work explored how prompts could be used to improve model performance without retraining.

Researchers experimented with prompt engineering, prompt tuning, and natural language instructions to better communicate tasks to language models. Most of these techniques focused on improving the input side of the interaction by changing how a task was described to the model.

Chain-of-thought prompting takes a different approach. Instead of modifying the instructions that precede a task, it augments the examples that follow them by exposing the reasoning process that connects inputs and outputs. This distinction may seem subtle, but it represents one of the paper's key insights: the contribution goes beyond a better prompt template. It focuses on the realization that demonstrating how to reason can be just as important as describing what task should be solved.

Viewed in this broader context, the paper acts as a bridge between research on reasoning traces and research on prompting. It combines the strengths of both traditions and, in doing so, lays the foundation for many later advances in language model reasoning, including self-consistency, STaR, process supervision, and the reasoning-oriented systems that followed in subsequent years.

Conclusion

Chain-of-Thought Prompting introduced a simple idea that changed how researchers think about reasoning in large language models. Rather than modifying model architectures or relying on additional training, the authors showed that reasoning abilities could often be unlocked by encouraging models to generate intermediate reasoning steps before producing an answer.

Across arithmetic, common sense, and symbolic reasoning tasks, the results demonstrated that large language models become significantly more capable when allowed to work through a problem step by step. More importantly, the paper revealed that many of these improvements emerge at larger scales, suggesting that reasoning isn't simply a product of prompting but a capability that becomes increasingly accessible as models grow more powerful.

What made this work particularly influential wasn't the complexity of the method, but the insight behind it. A model may possess the knowledge required to solve a problem, yet still fail to use that knowledge effectively when asked for an immediate answer. By exposing the reasoning process, Chain-of-Thought prompting showed that how a model arrives at an answer can be just as important as the answer itself.

This idea helped shift the focus of AI research beyond what language models know toward how they reason, plan, and solve problems. Many of the techniques that followed (including Self-Consistency, process supervision, verification-based methods, and modern reasoning-focused systems) build upon the foundation established by this paper.

Viewed in retrospect, Chain-of-Thought Prompting was more than a prompting technique. It marked a turning point in the study of language model reasoning, demonstrating that some capabilities aren't absent from a model but simply require the right conditions to emerge.

The infographic below highlights some of the most influential papers and milestones that shaped modern AI, from the introduction of GPT-1 and the scaling era of GPT-2 and GPT-3, to instruction tuning, Chain-of-Thought reasoning, Self-Consistency, process supervision, and the latest generation of reasoning-focused models. Together, these works reveal how the field evolved from teaching models to predict language toward helping them reason, verify, and solve increasingly complex problems.

Resources

• Pytorch Projects for GPT series

• Attention Is All You Need

• Improving Language Understanding by Generative Pre-Training (GPT-1)

• Language Models are Unsupervised Multitask Learners (GPT-2)

• Language Models are Few-Shot Learners (GPT-3)

• Scaling Laws for Neural Language Models

• Training Language Models to Follow Instructions with Human Feedback (InstructGPT)

• Finetuned Language Models are Zero-Shot Learners (FLAN)

• LaMDA: Language Models for Dialog Applications

• PaLM: Scaling Language Modeling with Pathways

• Program Induction by Rationale Generation: Learning to Solve and Explain Algebra Word Problems

• Training Verifiers to Solve Math Word Problems

• Show Your Work: Scratchpads for Intermediate Computation with Language Models

• Self-Consistency Improves Chain of Thought Reasoning in Language Models

• STaR: Bootstrapping Reasoning with Reasoning

• Emergent Abilities of Large Language Models

• Sparks of Artificial General Intelligence: Early Experiments with GPT-4

• Let's Verify Step by Step

• Learning to Reason with LLMs

• GPT-4 Technical Report

Contact Me

• Github

• X

• Linkedin

Naturally, my first instinct was to look at cloud providers. While services like ElevenLabs offer incredible voice quality, I quickly ran the numbers. Between the API pricing, token consumption for lengthy tutoring sessions, and the sheer volume of users I anticipated, the math got ugly very quickly. Relying on a paid API for every single sentence spoken within the app simply wasn't sustainable for an independent developer.

If you’re about to ask, "How far did you get with QuizRope?", well honestly, I straight-up gave up on the project back then because I couldn't find a sane, affordable solution for the TTS feature.

Beyond the prohibitive cost, there was the latency. Waiting for a server to process a prompt, generate the audio, and stream it back down to a mobile device completely breaks the conversational illusion. And worst of all, it meant every question a student asked would be beamed to a third-party server.

That frustration became the catalyst for my search to find a reliable, offline, and completely zero-cost solution.

In this article, we’re going to build a React Native application that performs high-fidelity Text-to-Speech (TTS) completely offline using your device's own hardware.

If you haven't set up your environment or need a refresher on local inference fundamentals, I highly recommend reading my previous article, How to Run a Local LLM Offline in React Native with QVAC, where I cover project initialization, prebuilding, and native hardware dependencies.

This guide assumes you already have a project with the QVAC SDK configured and ready to run on a physical device.

Table of Contents

• Prerequisites

• What is QVAC?

• The Architecture Supported by QVAC

• The Inference Pipeline

• Environment and Dependency Config

• The Audio Utility Packaging

• Complete Implementation

• Codebase Breakdown

• Conclusion

• Resources and Further Reading

Prerequisites

To get the most out of this article, you should have a solid foundation in modern web and mobile development:

• JavaScript/TypeScript & React: Familiarity with React concepts and hooks, especially useState, useEffect, and useRef.

• React Native & Expo: Basic understanding of layout structures (such as View, ScrollView, TextInput) and styling conventions.

• Asynchronous JavaScript & Binary Buffers: Experience with async/await, Promises, and basic manipulation of arrays like Int16Array or Buffer.

• Development Build Environment: Familiarity with running local development compilation commands, specifically npx expo prebuild to build native iOS and Android modules.

• Physical Mobile Device: Because local machine learning models leverage device-specific hardware acceleration and native optimizations, the QVAC SDK doesn't support simulator environments. You must have a physical iOS or Android testing device with Developer Mode enabled.

What is QVAC?

To help you follow along more effectively, let’s establish what QVAC is and why it exists.

Developed by Tether, QVAC is a local-first AI SDK designed for building cross-platform, peer-to-peer (P2P) applications and systems.

Many mobile applications that utilize Large Language Models (LLMs) or Text-to-Speech (TTS) engines rely on network requests to cloud-hosted APIs (such as OpenAI or ElevenLabs). While convenient, this model introduces dependencies on network connectivity, recurring API usage fees, and transmission of user data to third-party servers.

QVAC provides an alternative by executing AI models directly on the client device. This local-first architecture offers several practical advantages:

• Local-first execution: Runs inference directly on the client hardware, eliminating the need for external APIs or active internet connections.

• Peer-to-peer (P2P) support: Allows distributing inference tasks across local networks, helping coordinate workloads without centralized servers.

• Cross-platform compatibility: Provides a single JavaScript/TypeScript interface that works consistently across different hardware and runtime environments.

• Unified capabilities: Exposes text generation, transcription, image generation, and speech synthesis within a single package.

Key Concepts for On-Device Inference

To understand how QVAC runs on a mobile device, we must keep a few key concepts in mind:

• On-Device Inference: Running model calculations locally. Rather than relying on a single engine, QVAC supports multiple specialized local inference backends depending on the task (such as llama.cpp for text, whisper.cpp for transcription, or custom diffusion backends for image generation). Under the hood, these engines memory-map quantized model weights directly into the device's RAM and run calculations using native GPU hardware acceleration.

• Quantization (GGUF format): A mathematical optimization technique that compresses the model's weights (for example, from a standard 16-bit floating-point precision down to 4-bit or 8-bit integers). This makes it possible for models to fit into the memory constraints of consumer mobile hardware while keeping output quality high.

• KV (Key-Value) Cache: A memory area that stores calculated states of previous tokens so the model doesn't have to re-evaluate the entire context window with every word or token it generates.

The Architecture Supported by QVAC

Before writing code, it's crucial to understand what's actually happening under the hood. To handle local execution without melting your device, the QVAC SDK manages the hardware binding and model lifecycle while hooking into optimized, community-maintained GGML inference backends.

Instead of a one-size-fits-all approach, the QVAC SDK supports two distinctly different neural architectures for speech synthesis. Depending on your application's needs — whether you want instant voice cloning or ultra-high-fidelity pre-trained voices — you'll choose between Chatterbox and Supertonic.

1. The Chatterbox Engine

Chatterbox is built on a transformer-based language model architecture. It treats audio generation similarly to how an LLM predicts the next word in a sentence, but instead, it predicts discrete acoustic tokens.

Because of this architecture, Chatterbox excels at zero-shot voice cloning. Instead of relying purely on pre-baked voices, you can pass an optional referenceAudioSrc (a short WAV file of someone speaking) alongside your text. The transformer analyzes the reference audio's acoustic properties and generates a cloned voice based on those features.

2. The Supertonic Engine

Supertonic takes a completely different approach, utilizing diffusion-based latent denoising — the same fundamental architecture used by AI image generators like Stable Diffusion, but applied to audio.

It starts with pure digital noise and iteratively refines it into a 44.1 kHz high-fidelity speech waveform based on the text prompt. Supertonic uses a single, unified GGUF file rather than a split model. Instead of dynamic voice cloning, it relies on highly optimized, pre-trained voice styles (for example, voice: "F1" or voice: "M1") baked directly into the model. This makes it incredibly efficient for generating crystal-clear, studio-quality speech when you don't need dynamic cloning capabilities.

For this tutorial, we'll use Supertonic. It yields fantastic results out of the box and avoids the complexity of loading multiple companion files.

The Inference Pipeline

To visualize how we interact with these engines in our codebase, think of local TTS (Text to Speech) as running a virtual recording studio right in your phone's memory:

1. Hiring the actor (loading the model): We map the compressed GGUF file directly into the device's RAM or GPU VRAM.

2. Handing over the script (text input): We pass plain text to the loaded engine.

3. The performance (inference): The engine reads the text and mathematically predicts the sound waves. Crucially, the AI doesn't emit a finished audio file. Instead, it outputs raw digital sound waves known as PCM samples.

4. Packaging the audio: Because a raw list of numbers can't be played by standard media players, we must manually wrap the PCM data in a standard WAV header.

5. Closing the studio (unloading): Because speech synthesis is memory-intensive and maintains a persistent state, the model is cleared from RAM to free up resources and flush its context.

Environment and Dependency Config

Before we jump into the codebase, there's a crucial dependency setup to keep in mind if your project uses the pnpm package manager.

Because QVAC plugins rely on transitive native peer dependencies, strict package managers like pnpm will lock these dependencies down inside hidden .pnpm subfolders.

To ensure the QVAC native bundler (bare-pack) can resolve your worker plugins correctly at build time, create a .npmrc file in the root of your project:

shamefully-hoist=true

IMPORTANT: After creating this file, you must run a clean dependency install (pnpm install). This ensures a flat layout in your root node_modules so that all QVAC-specific helper packages are resolved properly during your local npx expo prebuild compilation step.

The Audio Utility Packaging

Because QVAC outputs raw PCM arrays, we need to construct a valid WAV file in memory and write it to the device's storage before the native audio player can play it.

To achieve this, let's create a utility module inside src/lib/utils.ts to build the required WAV header, convert raw audio samples into a binary buffer, and write it to local storage.

import { Buffer } from "buffer";
import * as FileSystem from "expo-file-system/legacy";

/**
 * Creates a WAV header for 16-bit PCM audio
 */
export function createWavHeader(
  dataLength: number,
  sampleRate: number,
): Buffer {
  const buffer = Buffer.alloc(44);
  const channels = 1; // Mono
  const byteRate = sampleRate * channels * 2; // 16-bit audio
  const blockAlign = channels * 2;

  buffer.write("RIFF", 0);
  buffer.writeUInt32LE(36 + dataLength, 4);
  buffer.write("WAVE", 8);
  buffer.write("fmt ", 12);
  buffer.writeUInt32LE(16, 16); // Subchunk1Size
  buffer.writeUInt16LE(1, 20); // AudioFormat (PCM)
  buffer.writeUInt16LE(channels, 22);
  buffer.writeUInt32LE(sampleRate, 24);
  buffer.writeUInt32LE(byteRate, 28);
  buffer.writeUInt16LE(blockAlign, 32);
  buffer.writeUInt16LE(16, 34); // BitsPerSample
  buffer.write("data", 36);
  buffer.writeUInt32LE(dataLength, 40);

  return buffer;
}

/**
 * Converts the raw Int16Array samples from QVAC to a binary Buffer
 */
export function int16ArrayToBuffer(int16Array: Int16Array): Buffer {
  const buffer = Buffer.alloc(int16Array.length * 2);
  for (let i = 0; i < int16Array.length; i++) {
    buffer.writeInt16LE(int16Array[i] ?? 0, i * 2);
  }
  return buffer;
}

/**
 * Main function to package and save the file to local mobile storage
 */
export async function saveAudioToDevice(
  audioBuffer: Int16Array,
  sampleRate: number,
): Promise<string> {
  try {
    const audioData = int16ArrayToBuffer(audioBuffer);
    const wavHeader = createWavHeader(audioData.length, sampleRate);
    const finalWavBuffer = Buffer.concat([wavHeader, audioData]);
    const base64Data = finalWavBuffer.toString("base64");

    const filename = `tts-speech-${Date.now()}.wav`;
    const fileUri = `\({FileSystem.documentDirectory}\){filename}`;

    await FileSystem.writeAsStringAsync(fileUri, base64Data, {
      encoding: FileSystem.EncodingType.Base64,
    });

    console.log(`✅ File saved locally at: ${fileUri}`);
    return fileUri;
  } catch (error) {
    console.error("❌ Failed to save audio file locally:", error);
    throw error;
  }
}

Complete Implementation

Let's bring it all together. We'll implement an interface that takes user input, manages download and loading states for the Supertonic engine, packages generated raw waves into a playable local file, and renders an interactive visual waveform player.

Replace your entry app file src/app/index.tsx with the following implementation:

import { useState, useEffect } from "react";
import {
  TextInput,
  KeyboardAvoidingView,
  Platform,
  ScrollView,
} from "react-native";
import {
  loadModel,
  unloadModel,
  textToSpeech,
  downloadAsset,
  TTS_EN_SUPERTONIC_Q8_0,
  getModelInfo,
  type ModelProgressUpdate,
} from "@qvac/sdk";
import { saveAudioToDevice } from "@/lib/utils";
import { TtsModelLoader } from "@/components/tts-model-loader";
import { AudioPlayer } from "@/components/audio-player";
import {
  Card,
  CardContent,
  CardDescription,
  CardHeader,
  CardTitle,
} from "@/components/ui/card";
import { Button } from "@/components/ui/button";
import { Text } from "@/components/ui/text";

const SUPERTONIC_SAMPLE_RATE = 44100;

// Global reference for our model ID
let globalModelId: string | null = null;

type TtsStatus =
  | { phase: "idle" }
  | { phase: "synthesizing" }
  | { phase: "done"; audioUri: string }
  | { phase: "error"; message: string };

export default function TextToVoiceScreen() {
  const [text, setText] = useState("");
  const [status, setStatus] = useState<TtsStatus>({ phase: "idle" });

  const [isModelLoaded, setIsModelLoaded] = useState(!!globalModelId);
  const [isDownloading, setIsDownloading] = useState(false);
  const [downloadProgress, setDownloadProgress] = useState(0);

  const isBusy = status.phase === "synthesizing";

  useEffect(() => {
    async function checkAndAutoLoad() {
      if (globalModelId) return;
      try {
        const info = await getModelInfo({ name: TTS_EN_SUPERTONIC_Q8_0.name });
        if (info.isCached) {
          setIsDownloading(true);
          setDownloadProgress(1);

          globalModelId = await loadModel({
            modelSrc: TTS_EN_SUPERTONIC_Q8_0,
            modelConfig: {
              ttsEngine: "supertonic",
              language: "en",
              voice: "F1",
              ttsSpeed: 1.05,
              ttsNumInferenceSteps: 5,
            },
          });

          setIsModelLoaded(true);
          setIsDownloading(false);
        }
      } catch (err: unknown) {
        console.warn("Failed to auto-load cached model on mount:", err);
        setIsDownloading(false);
      }
    }
    checkAndAutoLoad();
  }, []);

  const handleDownloadModel = async () => {
    if (isDownloading || isModelLoaded) return;

    try {
      setIsDownloading(true);
      setDownloadProgress(0);

      await downloadAsset({
        assetSrc: TTS_EN_SUPERTONIC_Q8_0,
        onProgress: (p: ModelProgressUpdate) => {
          setDownloadProgress(p.percentage / 100);
        },
      });

      setDownloadProgress(1);

      globalModelId = await loadModel({
        modelSrc: TTS_EN_SUPERTONIC_Q8_0,
        modelConfig: {
          ttsEngine: "supertonic",
          language: "en",
          voice: "F1",
          ttsSpeed: 1.05,
          ttsNumInferenceSteps: 5,
        },
      });

      setIsModelLoaded(true);
      setIsDownloading(false);
    } catch (err: unknown) {
      console.error("Failed to download or load model:", err);
      setIsDownloading(false);
      setStatus({
        phase: "error",
        message: err instanceof Error ? err.message : String(err),
      });
      setIsModelLoaded(false);
    }
  };

  const handleSubmit = async () => {
    if (!text.trim() || isBusy || !globalModelId) return;

    try {
      setStatus({ phase: "synthesizing" });

      // 1. Unload and reload the model to reset its state and clear the KV cache.
      if (globalModelId) {
        await unloadModel({ modelId: globalModelId });
      }
      globalModelId = await loadModel({
        modelSrc: TTS_EN_SUPERTONIC_Q8_0,
        modelConfig: {
          ttsEngine: "supertonic",
          language: "en",
          voice: "F1",
          ttsSpeed: 1.05,
          ttsNumInferenceSteps: 5,
        },
      });

      // 2. Synthesize text to raw PCM samples
      const result = textToSpeech({
        modelId: globalModelId,
        text: text.trim(),
        inputType: "text",
        stream: false,
      });

      const audioBuffer = await result.buffer;

      // 3. Package and save WAV file using our local util
      const samplesInt16 = new Int16Array(audioBuffer);
      const wavUri = await saveAudioToDevice(
        samplesInt16,
        SUPERTONIC_SAMPLE_RATE,
      );

      // 4. Show player
      setStatus({ phase: "done", audioUri: wavUri });
    } catch (err: unknown) {
      console.error("TTS error:", err);
      const msg = err instanceof Error ? err.message : String(err);
      setStatus({ phase: "error", message: msg });
    }
  };

  const buttonLabel =
    status.phase === "synthesizing" ? "Synthesizing…" : "Synthesize Speech";

  if (!isModelLoaded) {
    return (
      <TtsModelLoader
        onDownload={handleDownloadModel}
        isDownloading={isDownloading}
        progress={downloadProgress}
      />
    );
  }

  return (
    <KeyboardAvoidingView
      behavior={Platform.OS === "ios" ? "padding" : "height"}
      className="flex-1 bg-black"
    >
      <ScrollView contentContainerClassName="flex-grow p-6  justify-center">
        <Card className="border border-border bg-card max-w-md w-full mx-auto">
          <CardHeader>
            <CardTitle variant="h3" className="text-white text-center">
              Text to Voice
            </CardTitle>
            <CardDescription className="text-center mt-1">
              Type or paste your content to synthesize speech
            </CardDescription>
          </CardHeader>

          <CardContent className="gap-6">
            <TextInput
              className="bg-muted text-white border border-border rounded-lg p-4 h-48 text-base leading-6"
              multiline
              numberOfLines={8}
              placeholder="Type your message here..."
              placeholderTextColor="#666"
              value={text}
              onChangeText={setText}
              style={{ textAlignVertical: "top" }}
              editable={!isBusy}
            />

            {status.phase === "error" && (
              <Text className="text-destructive text-sm text-center">
                {status.message}
              </Text>
            )}

            {status.phase === "done" && <AudioPlayer uri={status.audioUri} />}

            <Button
              onPress={handleSubmit}
              className="w-full h-12 rounded-xl"
              disabled={!text.trim() || isBusy}
            >
              <Text className="font-semibold text-lg">{buttonLabel}</Text>
            </Button>
          </CardContent>
        </Card>
      </ScrollView>
    </KeyboardAvoidingView>
  );
}

Codebase Breakdown

Let’s lift the hood on how this local Text-to-Speech implementation manages native model lifecycles and processes raw audio arrays.

1. Managing the Native Lifecycle

Loading neural network weights for speech synthesis is computationally expensive. When the QVAC runtime initializes a model, it must read parameters from the local disk and copy the active weights into device RAM.

To handle this efficiently, we declared the reference variable outside the component scope:

let globalModelId: string | null = null;

If globalModelId were tracked inside component states, navigating away from the text-to-speech screen would clean up the state, causing the app to unnecessarily drop the reference. Storing the ID globally ensures we hold onto it across layout transitions.

2. Flushing the KV Cache: Unload and Reload

One of the most important aspects of offline generation using GGML engines is state management:

// 1. Unload and reload the model to reset its state and clear the KV cache.
if (globalModelId) {
  await unloadModel({ modelId: globalModelId });
}

globalModelId = await loadModel({ ... });

WARNING about acoustic hallucinations: If you continuously synthesize sentences on a single TTS model instance without resetting it, the model's Key-Value (KV) cache fills up. It begins treating your new sentence as a continuation of the previous one, leading to heavy robotic distortion, echoing, and repeated voices.

By explicitly destroying the model via unloadModel and immediately booting a fresh instance with loadModel, we're forcing a pristine, empty context window. Since the model is already downloaded and memory-mapped, reloading the model directly from local flash storage is extremely fast, typically completing in a fraction of a second on modern mobile hardware to ensure a seamless user experience while guaranteeing artifact-free audio.

3. Demystifying the WAV Header Structure

Operating systems and built-in mobile media decoders are unable to parse raw, naked PCM (Pulse Code Modulation) sound waves directly. A raw PCM buffer is simply a stream of numerical coordinates representing audio wave amplitudes.

We resolve this by prepending-formatting our PCM buffer with a standard 44-byte RIFF/WAVE header.

This header acts as a passport, defining:

• AudioFormat (1): Signals uncompressed linear PCM.

• NumChannels (1): Mono audio.

• SampleRate (44100): The clock frequency required for Supertonic playback.

• BitsPerSample (16): 16-bit word length (2 bytes per sample).

Additionally, writing the file is handled via Base64 encoding to safely cross React Native's JavaScript-to-Native bridge without dropping binary data:

const base64Data = finalWavBuffer.toString("base64");
await FileSystem.writeAsStringAsync(fileUri, base64Data, {
  encoding: FileSystem.EncodingType.Base64,
});

4. Visual Waveform Player

Rather than using a basic headless native audio player that fires immediately in the background, we pass the local WAV file path to a custom <AudioPlayer> component powered by @simform_solutions/react-native-audio-waveform.

This module analyzes our newly written WAV file and draws a sleek, WhatsApp-inspired interactive visual waveform, giving the user full control over playback, dynamic speed adjustments (1x, 1.5x, 2x), and seeking. It's a vast UX improvement that makes the final result feel premium and polished.

Conclusion

Transitioning Text-to-Speech from the cloud to on-device hardware offers a practical approach for mobile application developers. Running model inference locally eliminates reliance on remote internet connectivity, removes recurring API usage costs, and ensures that user text inputs never leave the physical device.

Integrating local speech synthesis can be highly beneficial for interactive, educational, or conversational apps. For example, in voice-guided systems, on-device TTS allows applications to function in private or offline environments. As edge processors gain dedicated hardware acceleration cores and open-source models decrease in memory size through quantization research, local-first architectures present a compelling alternative for developers prioritizing privacy, offline resilience, and predictable cost structures.

Resources and Further Reading

To dive deeper into local Text-to-Speech inference, inspect the source code, or explore advanced configurations for your mobile applications, check out the following resources:

• QVAC Expo Integration Docs: Learn more about configuring custom local models in Expo.

• react-native-audio-waveform: Learn more about interactive React Native audio visualizations.

• GGUF Model Hub on Hugging Face: Browse compatible quantized open-source models.

• Latent Denoising Deep Dive: Technical deep dive into Diffusion-based acoustic generation.

• https://github.com/DjibrilM/QVAC-TTS-Expo-Implementation: Full implementation code.

In this article, you’ll learn how to prepare a real-world medical imaging dataset for machine learning, from initial data validation to a complete preprocessing pipeline.

We’ll use the Chest X-Ray Pneumonia dataset as our running example, but the lessons apply broadly to healthcare imaging data, including ultrasound, MRI, CT, and dermatology images.

What You'll Learn in This Article

By the end of this article, you'll know how to:

• Approach healthcare data preprocessing differently from preprocessing structured data, and recognize where standard techniques fall short

• Validate a medical imaging dataset before training to catch corrupted files, mislabels, and data leakage between train and test

• Apply six core preprocessing techniques for medical images

• Build a complete preprocessing pipeline for chest X-rays using Python with OpenCV.

What We'll Cover:

• Why Preprocessing Data Matters More in Healthcare

• The Dataset

• Before Preprocessing: Validate the Dataset

• The Six Pillars of Healthcare Imaging Preprocessing

• Pillar 1: Scaling — Making the Numbers Play Fair

• Pillar 2: Normalization — Centering the Data

• Pillar 3: Guiding the Model's Attention

• Pillar 4: Handling Missing Data

• Pillar 5: Resizing & Resampling — Fitting Everything in the Same Frame

• Pillar 6: Denoising & Artifact Handling — Cleaning the Window

• Putting it All together: A Complete Pipeline

• Try it Yourself

• Conclusion

Why Preprocessing Data Matters More in Healthcare

Imagine handing a toddler a jigsaw puzzle with missing pieces, warped edges, and pieces from three different puzzles mixed together. The toddler can't solve it, but that isn't really the toddler's fault.

The same thing happens when raw, messy data gets fed into a machine learning model. A bad prediction on a clinical image can mean a missed diagnosis.

Healthcare data tends to be messier than what most ML practitioners are used to:

• Images come from different machines, hospitals, and acquisition protocols

• Labels are inconsistent, sometimes missing, sometimes wrong

• Patient data is incomplete

• Image sizes, contrast levels, and orientations vary across sources

Poor preprocessing often leads to models that perform well on benchmark datasets but struggle to generalize to data collected from different hospitals or imaging devices.

The Dataset

This guide uses the Chest X-Ray Pneumonia dataset by Paul Mooney on Kaggle. It's a strong choice for learning preprocessing because:

• It contains around 5,800 pediatric chest X-rays

• It has two clear classes — Normal and Pneumonia

• It's already organized into train, validation, and test folders

• The images are recognizable without specialized medical training

• It exhibits almost every preprocessing challenge worth learning

The dataset is available at Kaggle: Chest X-Ray Pneumonia.

Folder Structure

After downloading, the dataset is organized like this:

chest_xray/
├── train/
│   ├── NORMAL/
│   └── PNEUMONIA/
├── val/
│   ├── NORMAL/
│   └── PNEUMONIA/
└── test/
    ├── NORMAL/
    └── PNEUMONIA/

Side-by-side comparison — Normal vs Pneumonia chest X-ray:

A quick first look at one of the images:

import os
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import cv2

DATA_DIR = "chest_xray"
TRAIN_DIR = os.path.join(DATA_DIR, "train")

# Peek at a sample image
sample_path = os.path.join(TRAIN_DIR, "NORMAL", os.listdir(os.path.join(TRAIN_DIR, "NORMAL"))[0])
sample_image = cv2.imread(sample_path, cv2.IMREAD_GRAYSCALE)

print(f"Image shape: {sample_image.shape}")
print(f"Pixel range: {sample_image.min()} to {sample_image.max()}")
print(f"Data type: {sample_image.dtype}")

The output reveals a few useful things right away: most images are large (often around 1500×2000 pixels), pixel values fall in the 0–255 range, and image sizes vary across the dataset. Each of these observations will inform a preprocessing step.

Before Preprocessing: Validate the Dataset

Before applying any transformations, it's worth checking that the data itself is intact. This step alone catches issues that would otherwise cause training to fail silently or produce misleading results.

A simple validation function:

def validate_dataset(data_dir):
    """Scan a dataset folder and flag common data quality issues."""
    corrupted = []
    too_small = []
    nearly_black = []
    total = 0
    
    for class_name in os.listdir(data_dir):
        class_path = os.path.join(data_dir, class_name)
        if not os.path.isdir(class_path):
            continue
        for fname in os.listdir(class_path):
            fpath = os.path.join(class_path, fname)
            total += 1
            try:
                img = cv2.imread(fpath, cv2.IMREAD_GRAYSCALE)
                if img is None:
                    corrupted.append(fpath)
                    continue
                if img.shape[0] < 100 or img.shape[1] < 100:
                    too_small.append(fpath)
                if img.mean() < 5:
                    nearly_black.append(fpath)
            except Exception:
                corrupted.append(fpath)
    
    print(f"Total files scanned: {total}")
    print(f"Corrupted: {len(corrupted)}")
    print(f"Too small: {len(too_small)}")
    print(f"Nearly black: {len(nearly_black)}")
    return corrupted, too_small, nearly_black

validate_dataset(TRAIN_DIR)

Common issues this catches:

• Corrupted files — files that won't open at all

• Empty or nearly-black images — failed acquisitions or saved-as-blank files

• Wrong dimensions — thumbnails or partial downloads mixed in

• Duplicate images — the same scan appearing in both train and test (this causes data leakage)

• Mislabeled images — a normal X-ray placed in the pneumonia folder

⚠️ This step is critical, One corrupted file can crash a training loop hours into a run. One duplicate between train and test can inflate accuracy scores by several percentage points without anyone noticing.

The Six Pillars of Healthcare Imaging Preprocessing

Preprocessing for medical images can be organized around six core concerns. Two of them carry over directly from preprocessing structured data. Two need to be adapted because the mechanics change when the input is an image. And two are entirely new, they only exist once the data becomes pictures of human bodies.

Pillar 1: Scaling — Making the Numbers Play Fair

Imagine two children comparing their collections. One has 3 seashells. The other has 3,000 stickers. Asking who has more makes the answer seem obvious, but the scales are completely different. Comparing them meaningfully means putting both collections on the same measuring system.

In medical images, pixels usually range from 0 to 255 in 8-bit images, or 0 to 65,535 in some 16-bit medical DICOM images. Neural networks tend to train faster and more reliably when input values are small numbers close to zero.

The fix: Divide every pixel by its maximum possible value, bringing everything into the 0-to-1 range.

image = cv2.imread(sample_path, cv2.IMREAD_GRAYSCALE)

# Scale to [0, 1]
image_scaled = image.astype(np.float32) / 255.0

print(f"Before scaling: {image.min()} to {image.max()}")
print(f"After scaling:  {image_scaled.min():.3f} to {image_scaled.max():.3f}")

Takeaway: Pixel scaling follows the same principle as scaling any numerical feature. The values simply happen to be arranged as an image rather than a column.

Pillar 2: Normalization — Centering the Data

Imagine a teacher asks a class to rate a movie from 1 to 10. One child always gives 9s and 10s. Another spreads ratings evenly from 1 to 10. Comparing their opinions fairly requires adjusting each child's score relative to their own average.

In medical imaging even after scaling to 0–1, the overall brightness of images can vary. Some X-rays are taken with stronger exposure than others. Normalization shifts and rescales each image (or each channel) so the values are centered around zero with a standard deviation of one.

The fix: Subtract the mean, divide by the standard deviation.

# Compute mean and std from the TRAINING set only — never from validation or test
def compute_train_stats(train_dir, sample_limit=1000):
    """Compute pixel mean and std across the training set."""
    pixel_values = []
    count = 0
    for class_name in os.listdir(train_dir):
        class_path = os.path.join(train_dir, class_name)
        for fname in os.listdir(class_path):
            if count >= sample_limit:
                break
            img = cv2.imread(os.path.join(class_path, fname), cv2.IMREAD_GRAYSCALE)
            if img is not None:
                pixel_values.append(img.astype(np.float32).flatten() / 255.0)
                count += 1
    pixels = np.concatenate(pixel_values)
    return pixels.mean(), pixels.std()

train_mean, train_std = compute_train_stats(TRAIN_DIR)
image_normalized = (image_scaled - train_mean) / train_std

⚠️ Avoid this common mistake: Statistics for normalization should be computed from the training set only, never from validation or test. Including those in the calculation leaks information from the evaluation data into the model. The same statistics should then be applied to validation, test, and any new data at inference time.

Takeaway: Centering and scaling each image around the dataset's statistics is the imaging equivalent of standardizing a feature column. The pixels are now comparable across images, regardless of how bright or dim each scan happened to be.

Pillar 3: Guiding the Model's Attention

Imagine a child walking into a crowded pet store. Instead of describing every animal in sight, a parent points to the features that matter: “Look at the soft fur, the fluffy tail, and the nice small size.” The child learns where to focus their attention.

Medical image preprocessing does something similar. It highlights the regions and features most relevant to the diagnostic task.

• Region-of-interest (ROI) cropping — focus on the lung field and discard the patient's arms, machine borders, and any imprinted text

• Contrast enhancement — use techniques like CLAHE (Contrast Limited Adaptive Histogram Equalization) to make subtle lung textures more visible

• Channel selection — for images stored as RGB but containing grayscale information, convert to single-channel input to reduce noise

CLAHE applied to an X-ray:

# CLAHE enhances local contrast — useful for X-rays
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
image_enhanced = clahe.apply(image)

# Visualize the difference
fig, axes = plt.subplots(1, 2, figsize=(12, 6))
axes[0].imshow(image, cmap='gray')
axes[0].set_title('Original')
axes[1].imshow(image_enhanced, cmap='gray')
axes[1].set_title('After CLAHE')
plt.show()

Takeaway: The goal of teaching the model what to look at hasn't changed. With structured data, the answer is in new columns. With images, the answer is in cropping, enhancement, and emphasizing the regions that carry diagnostic signal.

Pillar 4: Handling Missing Data

Imagine reading a storybook with a few damaged pages. You don’t throw away the entire book, you decide whether to skip the page, infer what might be missing, or mark it for review.

In medical imaging, missing data can mean corrupted files, missing labels, or incomplete studies rather than empty spreadsheet cells.

The same three strategies — drop, impute, flag — still apply, just with different mechanics:

# Strategy 1: Drop — remove unreadable or empty images
def is_valid_image(path):
    try:
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            return False
        if img.mean() < 5:           # nearly black
            return False
        if img.shape[0] < 50 or img.shape[1] < 50:  # too small
            return False
        return True
    except Exception:
        return False

# Strategy 2: Impute — rare for images, but possible (e.g., in painting to fill in missing patches). Generally avoided for diagnostic data.

# Strategy 3: Flag — track which patients are missing which modalities,
#   and let the model condition on availability. Common in multi-modal healthcare ML.

Takeaway: "Missing" in imaging data is rarely just a NaN. It can be a broken file, an unlabeled scan, an absent modality, or a black corner inside an image. The same three strategies still apply.

Pillar 5: Resizing & Resampling — Fitting Everything in the Same Frame

Imagine displaying children’s drawings on a classroom wall. If every drawing is a different size, they won’t fit neatly into the display. You resize them while preserving their proportions.

Medical images must often be resized to a common input size, but anatomical structures should retain their original shape.

The fix: Resize all images to a common shape. For medical data, how the resizing is done matters.

TARGET_SIZE = (224, 224)

# Simple resize (may distort aspect ratio)
image_resized = cv2.resize(image, TARGET_SIZE)

# Better: preserve aspect ratio with padding
def resize_with_padding(image, target_size):
    h, w = image.shape[:2]
    target_h, target_w = target_size
    scale = min(target_h / h, target_w / w)
    new_h, new_w = int(h * scale), int(w * scale)
    resized = cv2.resize(image, (new_w, new_h))
    
    pad_h = target_h - new_h
    pad_w = target_w - new_w
    top, bottom = pad_h // 2, pad_h - pad_h // 2
    left, right = pad_w // 2, pad_w - pad_w // 2
    padded = cv2.copyMakeBorder(resized, top, bottom, left, right,
                                 cv2.BORDER_CONSTANT, value=0)
    return padded

image_clean_resize = resize_with_padding(image, TARGET_SIZE)

⚠️ Why aspect ratio matters in healthcare: Squishing a chest X-ray horizontally makes the lungs look unnatural. Models trained on distorted anatomy often perform worse on real scans. Preserving aspect ratio is generally the safer choice.

Takeaway: Models need a consistent input size, but the geometry of the anatomy needs to be preserved. Resize, but resize carefully.

Pillar 6: Denoising & Artifact Handling — Cleaning the Window

Imagine looking through a window with dust and smudges on the glass. Cleaning the window makes the view clearer, but scrubbing too aggressively could scratch the glass.

Similarly, medical images often contain noise and acquisition artifacts that should be reduced carefully without removing clinically important details.

For chest X-rays, the most common issues are mild noise and burned-in text or markers. A gentle median or bilateral filter helps with the first, while cropping or masking helps with the second.

# Gentle denoising — careful not to blur away clinical detail
image_denoised = cv2.medianBlur(image, ksize=3)

# Bilateral filter preserves edges better than a median filter
image_bilateral = cv2.bilateralFilter(image, d=5, sigmaColor=50, sigmaSpace=50)

⚠️ A note of caution: Aggressive denoising can erase the features a model needs to detect a disease. For diagnostic ML, gentle filtering is generally preferred. A useful rule of thumb: if a radiologist can't distinguish the cleaned image from the original, the filtering has gone too far.

Takeaway: Imaging data carries noise that structured data doesn't have. The window can be cleaned, but never so aggressively that the view is wiped away with the smudges.

Putting it All Together: A Complete Pipeline

Here's how the six pillars combine into a single preprocessing function for chest X-ray images:

def preprocess_xray(image_path, target_size=(224, 224),
                    train_mean=0.482, train_std=0.236):
    """
    Full preprocessing pipeline for chest X-ray images.
    Applies all six pillars in order.
    """
    # Pillar 4: Validate first — skip corrupted files
    if not is_valid_image(image_path):
        return None
    
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    
    # Pillar 5: Resize with aspect ratio preserved
    image = resize_with_padding(image, target_size)
    
    # Pillar 6: Gentle denoising
    image = cv2.medianBlur(image, 3)
    
    # Pillar 3: Enhance contrast to highlight lung texture
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    image = clahe.apply(image)
    
    # Pillar 1: Scale to [0, 1]
    image = image.astype(np.float32) / 255.0
    
    # Pillar 2: Normalize using training set statistics
    image = (image - train_mean) / train_std
    
    return image

Try it Yourself

Every code snippet in this article is bundled into a runnable Kaggle notebook: Chest X-Ray Preprocessing — Kaggle Notebook. Fork it, attach the dataset, and run all the cells to see each preprocessing pillar in action on real chest X-rays.

Conclusion

Here's a summary of what we've discussed in this article:

Preprocessing for structured data is about making numbers play fair so a model can see them clearly.

Preprocessing for healthcare imaging is about respecting the messy reality of how medical data is captured, stored, and labeled. Some standard techniques carry over directly. Some need to be adapted. And a few preprocessing concerns only emerge once the data becomes pictures of human bodies.

Stepping back, whether it's a child learning to organize their toy box, or a model learning to spot pneumonia in a chest X-ray, the quality of learning depends on the quality of data preparation. Get the data right.

If this was useful, you can find a related conceptual primer on preprocessing more broadly here: Data Preprocessing for Machine Learning.