That score looked like a success, but the model had learned very little about the real world. CDC had calculated the sixth column from the other five, so the model simply worked out CDC's formula.

Data scientists call this problem target leakage, and it happens when the inputs you give a model already contain the answer in some form.

Leakage like this hides easily in public data, because a large share of public data is calculated from other public data. A government index might be built from survey columns, and a second index might be built from the first one. Agencies explain these recipes in their methodology PDFs, yet data catalogues rarely store them in a form a computer can check.

In this tutorial, you'll write that record yourself and then build a small Python tool that reads it. The tool works like the dependency checker inside a package manager: you tell it what you want to predict and which columns you plan to use, and it refuses any column that sits on a derivation path to or from your target.

By the end, you'll know how to:

• reproduce a real leak using live CDC data and scikit-learn

• describe what a dataset was built from in a small YAML file called a manifest

• walk that graph with breadth-first search and depth-first search

• make a checking tool that fails loudly on typos, broken files, and empty inputs

• run the check automatically on every push with GitHub Actions

Table of Contents

• Prerequisites

• Key Terms in Plain English

• Step 1: See the Leak for Yourself

• Step 2: Understand Why Public Data Leaks

• Step 3: Borrow an Idea from Package Managers

• Step 4: Write the Dependency Manifest

• Step 5: Build the Linter

• Step 6: Run the Linter on Real Cases

• Step 7: Make the Linter Fail Loudly on Bad Input

• Step 8: Run the Check Automatically in CI

• Step 9: Learn from My Mistakes

• Going Further with the Full Tool

• Conclusion

Prerequisites

To follow along, you'll need:

• Python 3.10 or newer

• a basic idea of what a pandas DataFrame is

• a terminal where you can run commands

• about 7 MB of free disk space for the CDC data file

Create a fresh project folder with a virtual environment inside it, so these libraries stay separate from the rest of your system. Then install the three libraries this tutorial uses:

mkdir leak-tutorial
cd leak-tutorial
python3 -m venv .venv
source .venv/bin/activate
pip install pandas scikit-learn pyyaml

On Windows, run .venv\Scripts\activate in place of the source line, and type python wherever this article says python3. Run every command in this tutorial from inside the leak-tutorial folder.

Every file you build in this tutorial is also in the companion repository, derives-from-tutorial, so you can compare your work against it if you get stuck.

The full version of the tool lives in a public GitHub repository, and I link to it at the end of the article.

Key Terms in Plain English

Here are five words that come up again and again in this tutorial:

• Target: the column you want your model to predict.

• Covariate: a column you feed into the model to help it predict the target (many people call these features).

• R² (R-squared): a score that tells you how closely a model's predictions match the real values. A score of 1.0 means a perfect match, and a score near 0 means the model explains very little.

• Census tract: a small area of the United States that usually holds about 4,000 people, roughly the size of a neighbourhood.

• Cross-validation: a fair way to test a model. You split the data into five parts, train on four, test on the fifth, and repeat until every part has had a turn as the test set.

Step 1: See the Leak for Yourself

The US Centers for Disease Control and Prevention (CDC) publishes the Social Vulnerability Index, or SVI. Emergency planners use it to find communities that may need extra help during a flood, a heatwave, or a disease outbreak.

CDC builds the SVI in layers. It starts with 16 columns from the American Community Survey (ACS), a large survey run by the US Census Bureau. Each column is a percentage, such as the share of people living in poverty or the share of households with zero vehicles.

CDC groups those 16 columns into four themes and ranks every census tract within each theme. It then combines the four theme ranks into one overall rank called RPL_THEMES.

How CDC builds the SVI: 16 ACS survey columns feed four theme ranks, and the four theme ranks feed the one overall rank. Every yellow box is calculated from the boxes below it.

Here's the detail that matters for this tutorial: CDC ships the raw ACS columns and the finished ranks together in the same CSV file. That makes it very easy to grab both and put them into one model.

Download the California file:

curl -L -o California.csv https://svi.cdc.gov/Documents/Data/2022/csv/states/California.csv

I use curl here because some Python installs on macOS fail to verify the website's security certificate when they download files directly.

Now create a file called leak_demo.py:

"""leak_demo.py: predict a published index from the columns it was built from."""
import pandas as pd
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.model_selection import KFold, cross_val_score

# CDC marks missing values as -999, so turn those into proper blanks.
df = pd.read_csv("California.csv", low_memory=False).replace(-999, float("nan"))


def score(inputs, target):
    data = df[inputs + [target]].dropna()
    model = HistGradientBoostingRegressor(random_state=0)
    folds = KFold(n_splits=5, shuffle=True, random_state=0)
    r2 = cross_val_score(model, data[inputs], data[target],
                         cv=folds, scoring="r2").mean()
    print(f"{target:<11} from {len(inputs)} column(s)  "
          f"tracts={len(data)}  R2 = {r2:.3f}")


# Theme 1 is built from exactly these five columns.
score(["EP_POV150", "EP_UNEMP", "EP_HBURD", "EP_NOHSDP", "EP_UNINSUR"],
      "RPL_THEME1")

# EP_NOINT ships in the same file, and CDC leaves it out of the index.
score(["EP_NOINT"], "RPL_THEMES")

Here's what the script does:

1. It loads the CSV and turns CDC's -999 markers into blank values, because CDC uses -999 to flag a missing value.

2. The score function trains a gradient boosting model, which is a strong and popular choice for tables of numbers, and it measures R² with five-fold cross-validation.

3. The first call predicts Theme 1 using the exact five columns CDC used to build Theme 1.

4. The second call predicts the overall rank using EP_NOINT, the share of households lacking a broadband internet subscription. CDC includes this column in the same file and leaves it out of the index.

Run it:

python3 leak_demo.py

The output of leak_demo.py. Theme 1, predicted from the five columns CDC built it from, scores R² = 0.998. The overall rank, predicted from a column CDC leaves out of the index, scores 0.384.

The first score is 0.998, which means the model rebuilt CDC's Theme 1 almost perfectly. CDC's formula is a fixed recipe, and the model had every ingredient.

The second score is 0.384. EP_NOINT sits beside the index in the file, and its score shows the size of an ordinary link between two related measures.

Now imagine a paper that reports R² = 0.998 for predicting social vulnerability. That number would look like a breakthrough, yet it would only show that the model had found CDC's recipe.

The scores in this article came from scikit-learn 1.8.0. They stay the same to three decimal places across scikit-learn 1.3.2 to 1.9.0, so your run should match.

Step 2: Understand Why Public Data Leaks

The SVI example is easy to spot because the inputs and the index sit in one file. Most real cases are harder, because the chain runs across several agencies.

Here's one real chain that crosses three organisations. FEMA's National Risk Index (NRI) includes a social vulnerability score. According to FEMA's technical documentation (version 1.20, December 2025), that score comes from the Census Bureau's Community Resilience Estimates. The Census Bureau builds those estimates from ACS survey data.

So a FEMA risk score and an ACS column can sit at two ends of one chain, even though they come from different agencies and different websites.

To see why computers miss this, you need to know about two kinds of history a number can have:

• Provenance answers the question "Where did this number arrive from?" For example, a value came from California.csv, which came from svi.cdc.gov.

• Derivation answers the question "What was this number calculated from?" For example, Theme 1 was calculated from five ACS columns.

Two kinds of history a number can have. Provenance, on the left, records the file and the website a value arrived from. Derivation, on the right, records the five ACS columns Theme 1 was calculated from.

Most data catalogues store provenance well, and Google's Data Commons is a good example: it defines provenance as "the physical unit of an import", which tells you the file a number came in. Derivation usually lives only in PDF methodology documents written for humans.

So when an automated pipeline searches for helpful covariates, it can happily collect columns that the target was built from. The pipeline sees high scores and keeps those columns.

Step 3: Borrow an Idea from Package Managers

Software developers solved a very similar problem long ago.

When you run pip install requests, pip reads a list of what requests depends on, then what those packages depend on, and so on down the tree. Because every dependency is written down, pip can spot trouble anywhere in the tree before it installs anything.

The diagram below puts that tree beside the data version of the same problem.

The same shape twice. On the left, pip's dependency tree: urllib3 and certifi feed requests, which feeds my-app. On the right, the data version: ACS.EP_UNEMP goes into a model that predicts FEMA_NRI.risk_score, and the same column also climbs through two other products into that risk score.

Public data needs the same kind of record. In the right-hand half of the diagram above, ACS.EP_UNEMP (the unemployment rate) goes into the model as a covariate. The same column also climbs up through two other products into FEMA_NRI.risk_score, which is the target. The column sits at both ends of the loop.

In computer science, this kind of diagram is a graph. Each box is a node, and each arrow is an edge. In a data graph, following the arrows always leads you upward and away from where you started, so the graph is a directed acyclic graph, or DAG for short. "Acyclic" means the arrows form zero loops.

Throughout this article, every arrow points from an ingredient to the product made from it.

Two family words help describe positions in the graph:

• An ancestor of a node is anything you reach by following arrows backwards from it, at any distance. The ACS columns are ancestors of the SVI.

• A descendant of a node is anything you reach by following arrows forwards from it. The SVI is a descendant of the ACS columns.

Your leak check then becomes one simple rule: every covariate must stay clear of the target's ancestors and descendants.

Step 4: Write the Dependency Manifest

A manifest is a file that lists every product and what each one was built from. You'll use YAML here because people can read and edit it easily.

Here's how one measured product looks:

ACS.EP_UNEMP:
  label: Unemployment rate
  measurementBasis: measured
  derivesFrom: []

The empty list in derivesFrom: [] records zero parents, because this product comes straight from a survey.

And here is a product built from another product:

FEMA_NRI.social_vulnerability:
  label: FEMA National Risk Index, social vulnerability
  measurementBasis: composite
  derivesFrom:
    - {variable: CENSUS_CRE.social_vulnerability, relation: identity, confidence: documented}

Each entry in derivesFrom is one edge in the graph, and every edge carries three facts:

• variable holds the name of the parent product, and that name must match a product defined elsewhere in the file.

• relation describes how the parent was used.

• confidence records how sure you are about the edge.

These are the five relations:

These are the three confidence levels:

The confidence field matters more than it first appears. A lineage graph full of guesses would recreate the same problem it aims to solve, so each edge should say how much evidence stands behind it.

Each product also has a measurementBasis:

This field records something public catalogues usually leave out: whether a number was counted or predicted. A census count and a random forest prediction look identical in a spreadsheet, yet they're very different kinds of evidence.

Now create mini-manifest.yaml with the content below. It's a trimmed slice of the full manifest with 11 products from real US data infrastructure, and each product keeps a few of its real edges so the file stays short.

# A small slice of derivation-manifest.yaml, used in the tutorial.
schema: derives-from/0.2

products:

  # ---- measured: counted or surveyed directly
  ACS.EP_POV150:
    label: Population below 150% of the poverty line
    measurementBasis: measured
    derivesFrom: []

  ACS.EP_UNEMP:
    label: Unemployment rate
    measurementBasis: measured
    derivesFrom: []

  ACS.EP_NOVEH:
    label: Households with zero vehicles
    measurementBasis: measured
    derivesFrom: []

  SAT.chirps_rainfall:
    label: CHIRPS satellite rainfall
    measurementBasis: measured
    derivesFrom: []

  NVSS.mortality:
    label: Death certificate records
    measurementBasis: measured
    derivesFrom: []

  # ---- built from other products
  CENSUS_CRE.social_vulnerability:
    label: Census Community Resilience Estimates, social vulnerability
    measurementBasis: modelled
    derivesFrom:
      - {variable: ACS.EP_POV150, relation: modelled_from, confidence: documented}
      - {variable: ACS.EP_UNEMP,  relation: modelled_from, confidence: documented}
      - {variable: ACS.EP_NOVEH,  relation: modelled_from, confidence: documented}

  FEMA_NRI.social_vulnerability:
    label: FEMA National Risk Index, social vulnerability
    measurementBasis: composite
    derivesFrom:
      - {variable: CENSUS_CRE.social_vulnerability, relation: identity, confidence: documented}

  HVRI.bric:
    label: Baseline Resilience Indicators for Communities
    measurementBasis: composite
    derivesFrom:
      - {variable: ACS.EP_UNEMP, relation: component, confidence: documented}
      - {variable: ACS.EP_NOVEH, relation: component, confidence: documented}

  FEMA_NRI.community_resilience:
    label: FEMA National Risk Index, community resilience
    measurementBasis: composite
    derivesFrom:
      - {variable: HVRI.bric, relation: identity, confidence: documented}

  FEMA_NRI.expected_annual_loss:
    label: FEMA National Risk Index, expected annual loss
    measurementBasis: modelled
    derivesFrom: []

  FEMA_NRI.risk_score:
    label: FEMA National Risk Index, overall risk score
    measurementBasis: composite
    derivesFrom:
      - {variable: FEMA_NRI.expected_annual_loss, relation: component, confidence: certain}
      - {variable: FEMA_NRI.social_vulnerability, relation: component, confidence: certain}
      - {variable: FEMA_NRI.community_resilience, relation: component, confidence: certain}

Step 5: Build the Linter

A linter is a tool that reads something and warns you about problems before they cause harm. Code linters such as Flake8 read source code, while this linter reads your manifest and your list of covariates.

Create a file called mini_lint.py. You'll build it in six parts, and the finished file stays under 200 lines.

Part 1: Load YAML and Refuse Duplicate Keys

"""mini_lint.py: refuse covariates that sit on a derivation path to or from the target."""
import argparse
import sys
from collections import deque
from itertools import combinations

import yaml

RANK = {"certain": 3, "documented": 2, "inferred": 1}
DETERMINISTIC = {"component", "identity", "denominator"}


def fail(message):
    """Exit code 2 means the manifest or the command itself is broken."""
    print(message, file=sys.stderr)
    sys.exit(2)


# ---------------------------------------------------------------- step 1
class StrictLoader(yaml.SafeLoader):
    """A YAML loader that stops on a repeated key."""


def refuse_duplicates(loader, node, deep=False):
    seen = {}
    for key_node, _ in node.value:
        key = loader.construct_object(key_node, deep=deep)
        line = key_node.start_mark.line + 1
        if key in seen:
            fail(f"manifest error: key {key!r} appears twice "
                 f"(line {seen[key]} and line {line})")
        seen[key] = line
    return loader.construct_mapping(node, deep=deep)


StrictLoader.add_constructor(
    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG, refuse_duplicates)

RANK turns confidence words into numbers, so the tool can find the weakest edge in a route. DETERMINISTIC lists the relations that are pure arithmetic.

fail prints a message and exits with code 2. Later in the tutorial, you'll see why code 2 must stay separate from code 1.

The loader deals with a sneaky YAML behaviour. If a key appears twice in the same block, PyYAML quietly keeps the last copy and throws the first one away. In a manifest, that can erase every edge of a product, and the tool would then see zero routes and happily clear a leaky covariate.

refuse_duplicates runs every time PyYAML builds a mapping (a Python dictionary). It walks through the keys, remembers the line number of each one, and stops the program as soon as a key repeats.

Part 2: Read Products and Edges

# ---------------------------------------------------------------- step 2
def load(path):
    try:
        with open(path) as fh:
            products = yaml.load(fh, StrictLoader)["products"]
    except (OSError, yaml.YAMLError, KeyError, TypeError) as e:
        fail(f"manifest error: unable to read {path}: {e}")

    edges = {}
    for name, product in products.items():
        edges[name] = []
        for e in product.get("derivesFrom") or []:
            if RANK.get(e.get("confidence")) is None:
                fail(f"manifest error: {name} has an edge with "
                     f"confidence {e.get('confidence')!r}")
            edges[name].append((e["variable"], e["relation"], e["confidence"]))
    return products, edges

load opens the file with the strict loader. If anything goes wrong while reading, such as a bad path or broken YAML, it calls fail.

It then builds a dictionary called edges. For each product name, it stores a list of (parent, relation, confidence) tuples. For example:

edges["FEMA_NRI.social_vulnerability"]
# [("CENSUS_CRE.social_vulnerability", "identity", "documented")]

It also checks that every confidence value is one of the three allowed words. A typo such as documneted would otherwise slip through and break the ranking later.

Part 3: Check the Manifest Before Trusting It

# ---------------------------------------------------------------- step 3
def undefined_names(products, edges):
    mentioned = {parent for rows in edges.values() for parent, _, _ in rows}
    return sorted(mentioned - set(products))


def find_cycle(edges):
    state = {}

    def visit(node, stack):
        state[node] = "open"
        stack.append(node)
        for parent, _, _ in edges.get(node, []):
            if state.get(parent) == "open":
                return stack[stack.index(parent):] + [parent]
            if parent in state:
                continue
            cycle = visit(parent, stack)
            if cycle:
                return cycle
        stack.pop()
        state[node] = "closed"
        return None

    for node in edges:
        if node in state:
            continue
        cycle = visit(node, [])
        if cycle:
            return cycle
    return None

A typo in a parent name, such as HVRI.brick in place of HVRI.bric, creates an edge that points at a product the file lacks. The traversal would stop at that dead end, and every covariate beyond it would look safe.

undefined_names collects every parent mentioned in any edge and subtracts the set of defined products. Anything left over is either a typo or a product you forgot to add.

find_cycle makes sure the manifest really is a DAG. A product built from itself is impossible in real data, and a loop would send the route finder around in circles forever.

The function uses depth-first search with two labels. When the search enters a node, it marks that node open. When it has finished exploring everything above the node, it marks it closed. If the search reaches a node that's still open, it has walked in a circle, and the function returns that circle so you can see it.

Part 4: Walk the Graph

# ---------------------------------------------------------------- step 4
def ancestors(edges, node):
    """Every product that `node` was built from, at any distance."""
    found = set()
    queue = deque([node])
    while queue:
        current = queue.popleft()
        for parent, _, _ in edges.get(current, []):
            if parent in found:
                continue
            found.add(parent)
            queue.append(parent)
    return found


def routes(edges, start, goal):
    """Every path from start up to goal. Safe because step 3 ruled out cycles."""
    found = []
    for parent, relation, confidence in edges.get(start, []):
        step = (parent, relation, confidence)
        if parent == goal:
            found.append([step])
        else:
            for rest in routes(edges, parent, goal):
                found.append([step] + rest)
    return found


def describe(start, route):
    chain = " -> ".join([start] + [parent for parent, _, _ in route])
    weakest = min(route, key=lambda step: RANK[step[2]])[2]
    arithmetic = all(rel in DETERMINISTIC for _, rel, _ in route)
    kind = "deterministic" if arithmetic else "statistical"
    return [chain, f"{kind}, weakest link: {weakest}"]

ancestors uses breadth-first search (BFS), so picture a queue at a ticket counter: you put the starting product in the queue. On each turn, you take the product at the front, look up its parents, and add each parent you have yet to see to the back of the queue. When the queue is empty, the found set holds every ancestor at every distance.

The found set also stops the search from visiting the same product twice. That matters because many products share parents.

ancestors tells you whether a covariate is upstream, and routes tells you how it gets there.

routes uses depth-first search with recursion. For each parent of start, it checks whether that parent is the goal. If it is, that single step is a complete route. Otherwise, the function calls itself to find every route from the parent to the goal, then puts the current step on the front of each one.

The function returns every route, and that choice is deliberate. An earlier version of my full tool reported only the shortest route, so the report showed whichever route had the fewest hops, even when a longer route rested on stronger evidence.

The recursion is safe here only because Part 3 already confirmed that the graph is a DAG.

describe turns a route into two readable lines. The first line is the chain of names. The second line says whether the route is deterministic (arithmetic at every step) or statistical (at least one model in the chain), and it names the weakest confidence level along the route, since a chain is only as strong as its weakest link.

Part 5: The Audit

The audit looks for three shapes in the graph:

Each panel shows one shape and the verdict it produces: an arrow running into the target (FAIL), an arrow running out of the target (FAIL), and two covariates hanging off one shared input (REVIEW).

• Ancestor: the covariate went into the target, directly or through other products. This is the classic leak, so the tool reports it as an error.

• Descendant: the target went into the covariate. Predicting a parent from its own child leaks just as badly, so this is also an error.

• Shared ancestor: two covariates came from the same input. This is a softer problem, because the pair carries overlapping information, so the tool raises a warning for a person to review.

# ---------------------------------------------------------------- step 5
def audit(products, edges, target, covariates):
    findings = []

    unknown = [n for n in [target, *covariates] if products.get(n) is None]
    if unknown:
        return [("ERROR", f"unknown name: {n}", ["check the spelling"])
                for n in unknown]

    target_ancestors = ancestors(edges, target)
    for cov in covariates:
        if cov in target_ancestors:
            found = routes(edges, target, cov)
            details = [line for r in found for line in describe(target, r)]
            findings.append(("ERROR", f"{cov} is an ancestor of the target "
                                      f"({len(found)} route(s))", details))
        if target in ancestors(edges, cov):
            found = routes(edges, cov, target)
            details = [line for r in found for line in describe(cov, r)]
            findings.append(("ERROR", f"{cov} is a descendant of the target "
                                      f"({len(found)} route(s))", details))

    for a, b in combinations(covariates, 2):
        if a in ancestors(edges, b) or b in ancestors(edges, a):
            findings.append(("ERROR", f"{a} and {b}: one is built from the other", []))
        elif ancestors(edges, a) & ancestors(edges, b):
            shared = sorted(ancestors(edges, a) & ancestors(edges, b))
            findings.append(("WARN", f"{a} and {b} share ancestors", shared))
    return findings

The audit starts with name checks: if you misspell a covariate, the tool reports an error straight away, because it holds zero information about a name outside the manifest, and calling that name safe would be a guess.

Next, it computes the target's ancestors once and tests each covariate against that set. It also computes each covariate's ancestors to see whether the target appears among them, which is how it catches descendants.

Finally, combinations from the itertools module produces every pair of covariates. If one covariate is built from the other, that's an error. If the pair shares any ancestor, that's a warning.

Part 6: Verdicts and Exit Codes

# ---------------------------------------------------------------- step 6
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", default="mini-manifest.yaml")
    parser.add_argument("--target", required=True)
    parser.add_argument("--covariates", nargs="+", required=True)
    args = parser.parse_args()

    products, edges = load(args.manifest)
    missing = undefined_names(products, edges)
    if missing:
        fail(f"manifest error: undefined names: {', '.join(missing)}")
    cycle = find_cycle(edges)
    if cycle:
        fail(f"manifest error: cycle: {' -> '.join(cycle)}")

    findings = audit(products, edges, args.target, args.covariates)
    severities = {severity for severity, _, _ in findings}
    if "ERROR" in severities:
        verdict = "FAIL"
    elif "WARN" in severities:
        verdict = "REVIEW"
    elif ancestors(edges, args.target):
        verdict = "PASS"
    else:
        verdict = "UNTRACED"

    basis = products.get(args.target, {}).get("measurementBasis", "unknown")
    print(f"target      {args.target}  [{basis}]")
    print(f"covariates  {', '.join(args.covariates)}")
    print(f"verdict     {verdict}\n")
    for severity, message, details in findings:
        print(f"  {severity:<5} {message}")
        for line in details:
            print(f"        {line}")

    sys.exit(1 if verdict == "FAIL" else 0)


if __name__ == "__main__":
    main()

main reads the command-line flags, loads the manifest, runs both self-checks, and then runs the audit. It turns the findings into one of four verdicts:

A broken manifest or a malformed command exits with code 2.

PASS and UNTRACED deserve a closer look. PASS means the tool walked a real family tree and found every covariate outside it. UNTRACED means the manifest holds an empty family tree for the target, so the walk had zero steps to take. Calling that a pass would flatter the tool, so it gets its own name.

The exit codes matter just as much. Exit code 1 means the check ran and found a leak, while exit code 2 means the check itself is broken. A CI pipeline needs to tell these two apart, because a leak asks you to change your covariates and a broken manifest asks you to fix the file.

Step 6: Run the Linter on Real Cases

Case 1: FEMA's Risk Score

FEMA's composite risk score multiplies Expected Annual Loss by a community risk factor. That factor is built from a social vulnerability score and a community resilience score, and both of those reach back to ACS survey columns through different organisations.

FEMA_NRI.risk_score and everything it was built from. Two routes, drawn in blue and red, both start at the same ACS unemployment column: one climbs through the Census Bureau's resilience estimates, the other through HVRI's BRIC index.

Suppose you want to predict the risk score using the poverty rate and the unemployment rate:

python3 mini_lint.py --target FEMA_NRI.risk_score --covariates ACS.EP_POV150 ACS.EP_UNEMP

The verdict is FAIL. The poverty rate reaches the target by one route, and the unemployment rate reaches it by two, one statistical and one deterministic.

The verdict is FAIL, with exit code 1.

ACS.EP_POV150 reaches the target by one route. ACS.EP_UNEMP reaches it by two.

The first unemployment route passes through the Census model, so the tool labels it statistical. The second route passes through HVRI's BRIC index, where unemployment is a direct ingredient, so the tool labels it deterministic.

Try tracing both routes by eye in a spreadsheet of column names and you'll quickly see why a graph helps. The traversal finds both in a fraction of a second.

Case 2: A Descendant

Now flip the direction and suppose you want to predict the Census score using FEMA's republished copy of it as a covariate:

python3 mini_lint.py --target CENSUS_CRE.social_vulnerability --covariates FEMA_NRI.social_vulnerability

The flipped case, and another FAIL. FEMA's republished copy of the Census score is a descendant of the target by one deterministic route.

FEMA's score is built directly from the Census score, so using it as an input hands the model the answer. The tool catches this as a descendant.

Case 3: REVIEW, PASS, and UNTRACED

Here are three runs that should come back clean or nearly clean:

python3 mini_lint.py --target NVSS.mortality --covariates FEMA_NRI.social_vulnerability HVRI.bric
python3 mini_lint.py --target FEMA_NRI.risk_score --covariates SAT.chirps_rainfall
python3 mini_lint.py --target NVSS.mortality --covariates SAT.chirps_rainfall

Three cleaner runs: REVIEW for the pair that shares two ACS parents, PASS for satellite rainfall against the risk score, and UNTRACED for a target with zero listed parents.

The first run returns REVIEW, because the two covariates share two ACS parents and overlap in what they tell the model.

The second run returns PASS, because the risk score has a traced family tree and satellite rainfall sits outside it.

The third run returns UNTRACED, because death certificate records are a direct count with zero listed parents.

These quiet results matter as much as the failures. A checker that raised an alarm on every input would be useless, so a good test set always includes cases that should pass.

Step 7: Make the Linter Fail Loudly on Bad Input

A safety tool earns trust by failing clearly. The worst outcome for a leak checker is a green PASS on a check that quietly skipped its work, and there are three common ways that can happen.

Trap 1: Typos in Names

To try this, copy mini-manifest.yaml to typo-manifest.yaml and change HVRI.bric to HVRI.brick inside the FEMA_NRI.community_resilience entry. Then run these two commands:

python3 mini_lint.py --target FEMA_NRI.risk_score --covariates ACS.EP_POV15
python3 mini_lint.py --manifest typo-manifest.yaml --target FEMA_NRI.risk_score --covariates ACS.EP_UNEMP

Two typos, two exit codes. The misspelled covariate becomes a FAIL with exit code 1, and the misspelled parent inside the manifest becomes a manifest error with exit code 2.

A misspelled covariate (ACS.EP_POV15) becomes a FAIL with exit code 1. A misspelled parent inside the manifest (HVRI.brick) becomes a manifest error with exit code 2. Both stop the run before any traversal happens.

Trap 2: Duplicate YAML Keys

Make another copy of the manifest called sneaky-manifest.yaml, then add one extra line at the very end of the file, inside the FEMA_NRI.risk_score block:

    derivesFrom: []

The risk score now has two derivesFrom keys. Create peek.py to see what plain PyYAML does with that:

import yaml

with open("sneaky-manifest.yaml") as fh:
    doc = yaml.safe_load(fh)

print(doc["products"]["FEMA_NRI.risk_score"]["derivesFrom"])

Now run peek.py, and then run the linter on the same file:

python3 peek.py
python3 mini_lint.py --manifest sneaky-manifest.yaml --target FEMA_NRI.risk_score --covariates ACS.EP_UNEMP

The duplicate key, seen two ways. Plain yaml.safe_load prints an empty list and says nothing, while the strict loader names the repeated key, both line numbers, and exits with code 2.

Plain yaml.safe_load returns an empty list, because PyYAML kept the second key, threw away all three real edges, and stayed silent about it. A linter built on that loader would find zero routes and let every leaky covariate through.

The strict loader stops with exit code 2 and points at both line numbers.

Trap 3: An Empty Covariate List

The third trap is easy to overlook. In CI, you might build the covariate list from a file or a shell variable. If that file is empty or the variable name has a typo, the command ends up with zero covariates.

An earlier version of my full tool accepted that and printed PASS, which is a clean bill of health for a check that skipped all its work.

python3 mini_lint.py --target FEMA_NRI.risk_score --covariates

An empty --covariates list now stops the run at argparse, before any traversal, with exit code 2.

In mini_lint.py, nargs="+" tells argparse that --covariates needs at least one value, and required=True makes the flag itself mandatory. An empty list now turns the pipeline red with exit code 2.

Step 8: Run the Check Automatically in CI

CI (continuous integration) runs checks for you every time you push code. This GitHub Actions workflow runs the linter on every push and every pull request. Save it as .github/workflows/lineage.yml in a repository that holds mini_lint.py and mini-manifest.yaml at its root:

name: lineage-check

on: [push, pull_request]

jobs:
  lint-lineage:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v5
      - uses: actions/setup-python@v6
        with:
          python-version: "3.12"
      - run: pip install pyyaml
      - name: Check covariates against the target's lineage
        run: |
          python3 mini_lint.py \
            --target FEMA_NRI.risk_score \
            --covariates $(cat features.txt)

Put your covariate names in a file called features.txt at the root of the repository, one name per line:

SAT.chirps_rainfall

The $(cat features.txt) part pastes those names into the command, and with this file the job passes and turns green. If someone later adds a leaky covariate such as ACS.EP_UNEMP, the job exits with code 1 and turns red. If someone empties features.txt by accident, argparse exits with code 2, and the job turns red as well.

REVIEW and UNTRACED also exit with code 0. If you want your pipeline to stop on those as well, change the last line of main so that every verdict other than PASS exits with code 1.

The companion repository runs this exact workflow, and you can see its results on the repository's Actions tab.

Step 9: Learn from My Mistakes

Building the full manifest taught me that a linter is only as good as the graph it reads. I got the SVI wrong twice, in opposite directions, and both mistakes came from the same habit of trusting the columns in a file over the methodology behind it.

Mistake 1: The SVI California file contains 24 columns whose names start with EP_, but CDC ranks only 16 of them into the index. My first manifest counted all 24 and recorded EP_NOINT (broadband subscriptions) as an ingredient of the index. That column sits in the file, and CDC leaves it out of the ranking, so I removed the edge.

Mistake 2: I then over-corrected and listed all eight unranked columns as safe bystanders that ship beside the index. Seven of those eight are race and ethnicity columns. When I checked the raw counts, those seven added up to E_MINRTY exactly, and the largest difference across all 9,109 tracts was zero.

EP_MINRTY is the single input to Theme 3, and Theme 3 feeds the overall index.

Why those seven columns aren't bystanders. They sum exactly to EP_MINRTY, which feeds Theme 3 one hop up, which feeds the overall index one hop above that. EP_NOINT, in the dashed box, ships in the same file and stays outside the index.

So those seven columns are ancestors of the index, two hops up. My own linter had been clearing them as safe covariates for an SVI target, which is exactly the kind of false clearance the tool exists to prevent.

Cross-validated R² for each product, predicted from the columns listed beside it. Every product predicted from its own inputs scores 0.987 or higher, while EP_NOINT, which only ships beside the index, reaches 0.384.

The chart makes the difference plain: the seven race and ethnicity columns rebuild Theme 3 at R² = 0.992, while EP_NOINT alone reaches only 0.384 against the overall index.

I now follow a stricter rule before I mark any column as safe. I read the methodology first, and then I test the relationship in the data.

The full manifest records the result with a field called coPublishedNonInputs, which lists columns that ship in the same file as an index and take zero part in computing it. For the SVI, EP_NOINT is now the only entry.

Going Further with the Full Tool

The mini linter in this tutorial covers the core ideas. The full project adds:

• a manifest with 60 products and 75 derivation edges across US and global data, including CDC PLACES, FEMA's National Risk Index, WorldPop, AlphaEarth satellite embeddings, and WFP's HungerMap LIVE

• a written evidence note for every product that carries edges, plus a correction field wherever an earlier claim turned out to be wrong

• a --graph mode that prints the whole derivation graph

• a built-in suite of eight real audit cases

• a reproduce_svi.py script that checks every R² figure in this article against the live CDC file

• a pinned Dockerfile, so the figures reproduce exactly

To try it:

git clone https://github.com/Adeniyikayodee/dependency_manifest.git
cd dependency_manifest
python3 lint_lineage.py
python3 lint_lineage.py --graph
python3 reproduce_svi.py

The full tool on the complete manifest: 60 products, 75 derivation edges, and eight audits that come back as 5 FAIL, 1 REVIEW, 1 PASS, and 1 UNTRACED.

The manifest is clear about its limits. It covers 60 products out of an estimated 400 or more official composite indices worldwide, and four of its edges are still marked inferred. The first audit of the file found six errors, and four of them sat in edges I had already labelled certain or documented.

The people best placed to write this kind of record are the agencies themselves, since they already describe their methods in PDF form. Two new fields on a public data schema, derivesFrom and measurementBasis, would give every producer a place to store what they already know.

If you work with public data, you can help by adding products you know well, or by checking the edges marked inferred against their source documents.

Conclusion

Target leakage in public data hides inside the recipes that agencies use to build their indices. A model can score close to perfect by rediscovering one of those recipes, and that score says very little about the real world.

In this tutorial, you:

• rebuilt CDC's Theme 1 at R² = 0.998 from its own five input columns

• separated provenance (where a number arrived from) from derivation (what it was calculated from)

• wrote a YAML manifest that records derivation edges with a relation and a confidence level

• built a linter that uses breadth-first search to find ancestors and depth-first search to list every route

• made the linter fail loudly on typos, duplicate YAML keys, cycles, and empty covariate lists

• wired the check into GitHub Actions with clear exit codes

Before you trust a high score on public data, ask yourself what your target was built from. Once you write the answer down, a few lines of Python can check it every time you train a model.

The code from this tutorial lives in derives-from-tutorial, and you can find the full tool, the manifest, and the reproduction script in the main dependency_manifest repository. The project is archived on Zenodo with the DOI 10.5281/zenodo.22274757, and you are free to use it under the CC0 licence.

Sources

• CDC/ATSDR Social Vulnerability Index: https://svi.cdc.gov/

• FEMA National Risk Index Technical Documentation v1.20, December 2025: https://www.fema.gov/sites/default/files/documents/fema_national-risk-index_technical-documentation.pdf

• Census Bureau Community Resilience Estimates: https://www.census.gov/programs-surveys/community-resilience-estimates.html

• CDC PLACES methodology, Preventing Chronic Disease, 2022: https://www.cdc.gov/pcd/issues/2022/21_0459.htm

• Data Commons data model: https://docs.datacommons.org/data_model.html

The truth about building with Large Language Models is that traditional software testing falls apart. You can't write a simple assert output ==expected when your system generates different text every time it runs.

Most tutorials out there will teach you how to build a chatbot or wire up a RAG pipeline and then they just...stop. "Deploy to production" they say, as if the hard part is over. But the hard part is actually knowing whether your AI is any good and catching it when it stops being good.

In this article, I'll walk you through building a complete evaluation pipeline. We'll also cover three different evaluation strategies that work at different levels of cost and depth.

What We'll Cover:

• Why Traditional Testing Breaks Down for LLM Applications

• The Three Layers of LLM Evaluation

• How to Build Layer 1: Deterministic Checks

• How to Build Layer 2: LLM-as-Judge Evaluation

• How to Build Layer 3: Human Evaluation Loops

• How to Build the Regression Testing Pipeline

• How to Know If Your AI Actually Got Better: Statistical Significance

• How to Put It All Together: The Complete Evaluation Architecture

• What I Wish I Knew Earlier

What You'll Need

To follow along, you should have Python 3.10+ and some basic experience calling an LLM API. It doesn't matter if you're using OpenAI, Anthropic, or a local model because the evaluation patterns work the same way.

You'll also need an OpenAI API key for the LLM-as-judge examples (we're using gpt-4o-mini since it's cheap and good enough for scoring).

If you already have an LLM-powered app you want to evaluate, even a tiny one, that's perfect. If not, the examples are self-contained so you can still follow everything.

Grab the dependencies here:

pip install openai numpy pandas scikit-learn python-dotenv

Why Traditional Testing Breaks Down for LLM Applications

If you've written tests for regular software, you know the drill. Function goes in, value comes out, you assert they match. Clean, simple, and done.

But LLMs break that entire model. And not in one way, but in several that compound on each other.

First, the outputs aren't deterministic. You can send the exact same prompt twice and get back different wording. Even setting temperature=0 doesn't fully save you because model providers update their models behind the scenes. The same API call in January and March might behave differently.

Second, there's no single right answer. If your app summarizes a document, what does a correct summary even look like? Two humans would write different summaries and both could be perfectly good. You can't assertEqual your way through that.

And third, nothing breaks visibly and there's no error, crash, or red line in your logs. The model just quietly returns a polished, confident wrong answer. Your uptime dashboard says 100% while your users are getting nonsense. This is the one that really gets you when an LLM fails.

So you can't just test LLM apps the way you test a REST API. You need scoring instead of pass/fail. You need to evaluate batches of outputs not individual ones. And you need something that runs continuously because the quality can drift over time without you changing a single line of code.

The Three Layers of LLM Evaluation

The approach I've landed on after a lot of trial and error uses three layers stacked from cheap-and-fast to expensive-and-thorough.

1. Layer 1 is deterministic checks. Think of these as bouncers at the door. Is the output valid JSON when it should be? Is it suspiciously short or absurdly long? Does it contain a hallucinated URL? These checks are instant, free and catch more problems than you'd expect.

2. Layer 2 is LLM-as-judge. This is where you use a separate LLM call to grade your main LLM's output. "Was this answer relevant? Was it accurate? Did it actually help?" A model like gpt-4o-mini is surprisingly good at scoring other models' work as long as you give it a clear rubric.

3. Layer 3 is human evaluation. Real people reviewing real outputs. You don't do this on every response, as that would be impossibly slow. But you do it periodically, to make sure your automated layers haven't drifted away from what good actually means.

The trick is knowing when to use which layer, and we'll build each one of them.

How to Build Layer 1: Deterministic Checks

When I first started building eval pipelines, I skipped straight to the fancy stuff: LLM judges, embedding similarity scores, the works. Meanwhile, my app was occasionally returning completely empty strings and I didn't notice for two weeks. Two weeks!

That's why I now start every project with deterministic checks. They're dead simple: no ML and no API calls, just plain Python asking basic sanity questions about the output. Does it exist? Is it the right format? Is it suspiciously short? Did the model hallucinate a URL?

You might be thinking these are too basic to matter. I thought so too. Then I ran them on a month of production logs and found that roughly a third of the bad outputs I'd missed would've been caught by checks you could write in five minutes.

Here's the DeterministicEvaluator class I now drop into every project on day one:

import json
import re
from dataclasses import dataclass


@dataclass
class EvalResult:
    """Holds the result of a single evaluation check."""
    check_name: str
    passed: bool
    score: float  # 0.0 to 1.0
    details: str


class DeterministicEvaluator:
    """Layer 1: Fast, rule-based checks for LLM outputs."""

    def check_json_validity(self, output: str) -> EvalResult:
        """Verify the output is valid JSON when JSON is expected."""
        try:
            json.loads(output)
            return EvalResult("json_validity", True, 1.0, "Valid JSON")
        except json.JSONDecodeError as e:
            return EvalResult("json_validity", False, 0.0, f"Invalid JSON: {e}")

    def check_length_bounds(
        self, output: str, min_chars: int = 10, max_chars: int = 5000
    ) -> EvalResult:
        """Check that output length falls within acceptable bounds."""
        length = len(output)
        if length < min_chars:
            return EvalResult(
                "length_bounds", False, 0.0,
                f"Too short: {length} chars (minimum: {min_chars})"
            )
        if length > max_chars:
            return EvalResult(
                "length_bounds", False, 0.0,
                f"Too long: {length} chars (maximum: {max_chars})"
            )
        return EvalResult("length_bounds", True, 1.0, f"Length OK: {length} chars")

    def check_no_hallucinated_links(self, output: str) -> EvalResult:
        """Detect URLs in output that the model may have fabricated."""
        url_pattern = r'https?://[^\s\)\]\}\"\'<>]+'
        urls = re.findall(url_pattern, output)
        if urls:
            return EvalResult(
                "no_hallucinated_links", False, 0.0,
                f"Found {len(urls)} URLs that may be hallucinated: {urls[:3]}"
            )
        return EvalResult("no_hallucinated_links", True, 1.0, "No URLs found")

    def check_required_sections(
        self, output: str, required: list[str]
    ) -> EvalResult:
        """Verify that required sections or keywords appear in the output."""
        missing = [s for s in required if s.lower() not in output.lower()]
        if missing:
            score = 1.0 - (len(missing) / len(required))
            return EvalResult(
                "required_sections", False, score,
                f"Missing sections: {missing}"
            )
        return EvalResult("required_sections", True, 1.0, "All sections present")

    def check_no_refusal(self, output: str) -> EvalResult:
        """Detect if the model refused to answer when it should not have."""
        refusal_phrases = [
            "i cannot", "i can't", "i'm unable to", "as an ai",
            "i don't have access", "i'm not able to"
        ]
        output_lower = output.lower()
        for phrase in refusal_phrases:
            if phrase in output_lower:
                return EvalResult(
                    "no_refusal", False, 0.0,
                    f"Possible refusal detected: '{phrase}'"
                )
        return EvalResult("no_refusal", True, 1.0, "No refusal detected")

    def run_all(self, output: str, config: dict = None) -> list[EvalResult]:
        """Run all deterministic checks and return results."""
        config = config or {}
        results = [
            self.check_length_bounds(
                output,
                config.get("min_chars", 10),
                config.get("max_chars", 5000)
            ),
            self.check_no_hallucinated_links(output),
            self.check_no_refusal(output),
        ]
        if config.get("expect_json"):
            results.append(self.check_json_validity(output))
        if config.get("required_sections"):
            results.append(
                self.check_required_sections(output, config["required_sections"])
            )
        return results


if __name__ == "__main__":
    evaluator = DeterministicEvaluator()

    # Test with a normal output
    good_output = "Python is a high-level programming language known for its readability."
    results = evaluator.run_all(good_output)
    for r in results:
        print(f"  {r.check_name}: {'PASS' if r.passed else 'FAIL'} ({r.details})")

    # Test with a suspicious output
    bad_output = "Visit https://fake-docs.example.com/api for more details."
    results = evaluator.run_all(bad_output)
    for r in results:
        print(f"  {r.check_name}: {'PASS' if r.passed else 'FAIL'} ({r.details})")

Every one of these checks runs in under a millisecond and they cost nothing. But don't let the simplicity fool you because the hallucinated links check alone has saved me from shipping fabricated documentation URLs to users more times than I'd like to admit.

Also one thing worth stressing is that these are starting points. The generic checks above work for any LLM app. But the biggest wins come from domain-specific ones. If your app generates SQL, add a syntax parser. If it drafts emails, verify that there's a subject line and a greeting. If it outputs code, try running it through a linter.

Every check you add here is one fewer bad output that reaches the expensive layers downstream or worse, your users.

How to Build Layer 2: LLM-as-Judge Evaluation

Alright, so your output passes the sanity checks: it's valid JSON, reasonable length, no fabricated links. But here's a question Layer 1 can't answer: is the response actually helpful?

An output can be perfectly structured, pass every deterministic check, and still be completely useless to the person reading it. "The capital of France is Berlin" is valid text, correct length, no hallucinated URLs...but it's also wrong.

This is where things get a little meta. The idea behind LLM-as-judge is that you make a separate LLM call whose only job is to read your main model's output and score it. Yes, you're using AI to grade AI. It sounds like asking one student to grade another student's homework. But it actually works surprisingly well, and research from labs like Anthropic and Google have shown that LLM judges correlate strongly with human evaluators when you give them clear scoring criteria.

The key phrase there is "clear scoring criteria." Without that, this whole approach falls apart.

How to Design Scoring Rubrics

If you tell an LLM "rate this from 1 to 10," you'll get back scores that are all over the place. A 7 on one run becomes a 5 on the next. The scores are essentially meaningless because the model has no shared definition of what each number means.

The fix is a rubric with concrete anchor descriptions. Here's one for helpfulness.

Score 1 - The response is completely irrelevant, incorrect, or harmful.
Score 2 - The response addresses the topic but contains major errors or omissions.
Score 3 - The response is partially correct but misses key information.
Score 4 - The response is correct and helpful with minor issues.
Score 5 - The response is comprehensive, accurate, and directly addresses the question.

Now notice how each level describes something you could point to in the output, not a vibe. "Completely off-topic" is observable. "Kind of bad" is not. That specificity is what makes the judge consistent across runs.

How to Implement the Judge

Here's the full LLMJudge class. I'll walk through the important design decisions after.

import json
import os
from openai import OpenAI
from dataclasses import dataclass

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))


@dataclass
class JudgeResult:
    """Holds the result of an LLM judge evaluation."""
    criterion: str
    score: int
    max_score: int
    reasoning: str


RUBRICS = {
    "relevance": {
        "description": "Does the response directly address the user's question?",
        "levels": {
            1: "Completely off-topic or addresses a different question entirely.",
            2: "Tangentially related but misses the core question.",
            3: "Addresses the question but includes significant irrelevant content.",
            4: "Directly addresses the question with minor tangents.",
            5: "Precisely and completely addresses the question asked.",
        },
    },
    "accuracy": {
        "description": "Is the factual content of the response correct?",
        "levels": {
            1: "Contains critical factual errors that would mislead the reader.",
            2: "Multiple factual errors on important points.",
            3: "Mostly accurate but contains one notable error.",
            4: "Accurate with only trivial imprecisions.",
            5: "Completely accurate with no factual errors.",
        },
    },
    "completeness": {
        "description": "Does the response cover all important aspects of the question?",
        "levels": {
            1: "Addresses less than 20 percent of what the question requires.",
            2: "Covers some aspects but misses major required components.",
            3: "Covers the basics but lacks depth on important points.",
            4: "Comprehensive coverage with minor gaps.",
            5: "Thoroughly covers all aspects the question requires.",
        },
    },
}


class LLMJudge:
    """Layer 2: Uses a separate LLM to evaluate response quality."""

    def __init__(self, model: str = "gpt-4o-mini"):
        self.model = model

    def evaluate(
        self, question: str, response: str, criterion: str
    ) -> JudgeResult:
        """Evaluate a single response on a single criterion."""
        rubric = RUBRICS[criterion]
        levels_text = "\n".join(
            f"Score {score}: {desc}"
            for score, desc in rubric["levels"].items()
        )

        judge_prompt = f"""You are an expert evaluator. Your job is to score an AI assistant's response.

CRITERION: {rubric['description']}

SCORING RUBRIC:
{levels_text}

USER QUESTION:
{question}

AI RESPONSE:
{response}

Evaluate the response on the criterion above. You must respond with valid JSON only:
{{"score": <integer 1-5>, "reasoning": "<2-3 sentence explanation>"}}"""

        judge_response = client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": judge_prompt}],
            temperature=0.0,
            response_format={"type": "json_object"},
        )

        result = json.loads(judge_response.choices[0].message.content)
        return JudgeResult(
            criterion=criterion,
            score=result["score"],
            max_score=5,
            reasoning=result["reasoning"],
        )

    def evaluate_all(
        self, question: str, response: str, criteria: list[str] = None
    ) -> list[JudgeResult]:
        """Evaluate a response across all specified criteria."""
        criteria = criteria or list(RUBRICS.keys())
        return [self.evaluate(question, response, c) for c in criteria]


if __name__ == "__main__":
    judge = LLMJudge()

    question = "What is a Python decorator and when should you use one?"
    good_response = (
        "A Python decorator is a function that takes another function as input "
        "and extends its behavior without modifying it. You define a decorator "
        "with the @decorator_name syntax above a function definition. Use "
        "decorators when you need to add cross-cutting concerns like logging, "
        "authentication checks, or caching to multiple functions without "
        "duplicating code in each one."
    )

    results = judge.evaluate_all(question, good_response)
    for r in results:
        print(f"  {r.criterion}: {r.score}/{r.max_score} - {r.reasoning}")

There are a few things worth calling out in this code.

1. Temperature is zero: You're not asking the judge to be creative. You want the same input to produce the same score every time, or as close to it as possible.

2. The output is structured JSON: I learned this one the hard way. If you let the judge respond in free text, you end up writing fragile parsing code to extract the score. Force JSON output and your life gets much easier.

3. The rubric is baked into every prompt: The judge never uses its own idea of what good means. It always scores against your rubric and that's what makes it reproducible.

How to Handle Judge Reliability

Even with all of that, a single judge call can be noisy. I've seen the same response score a 4 on one call and a 3 on the next. If you're making decisions based on those scores, that variance matters. Two things can help you with that.

The first is multi-judge consensus. This means you run the same evaluation three times and take the median. Yes, it costs 3x as much. But the scores become much more stable, and for CI/CD gating decisions, stability matters more than saving a few cents.

The second is calibration sets. You keep a small set of responses (maybe 20-30) where you already have reliable human scores. Run your judge on these periodically. If the judge starts disagreeing with the humans, something changed and you need to investigate.

We can look at this consensus implementation that shows how to handle that:

import numpy as np


def evaluate_with_consensus(
    judge: LLMJudge,
    question: str,
    response: str,
    criterion: str,
    num_judges: int = 3,
) -> JudgeResult:
    """Run multiple judge evaluations and return the median."""
    results = [
        judge.evaluate(question, response, criterion)
        for _ in range(num_judges)
    ]
    scores = [r.score for r in results]
    median_score = int(np.median(scores))
    median_result = min(results, key=lambda r: abs(r.score - median_score))
    return JudgeResult(
        criterion=criterion,
        score=median_score,
        max_score=5,
        reasoning=f"Consensus ({scores}): {median_result.reasoning}",
    )

How to Build Layer 3: Human Evaluation Loops

I once had an LLM judge giving a response 5/5 on accuracy, 5/5 on relevance, 4/5 on completeness. The scores looked perfect until a colleague actually read the response and said, "This is technically correct but it would confuse the hell out of anyone who isn't already an expert." And he was right.

The answer used jargon the user wouldn't know, buried the key point three paragraphs deep, and read like a textbook instead of a helpful reply.

That's the ceiling of automated evaluation. LLM judges are great at detecting factual errors and structural problems, but they have blind spots around tone, clarity for a specific audience, and the subtle difference between "correct" and "actually helpful." Those blind spots are where human evaluation comes in.

Now, to be clear, this doesn't mean hiring a team to review every single response. That doesn't scale and you don't need it. The goal is narrower: get a small batch of human scores on a regular schedule and use those scores as a reality check on your automated layers.

How to Build a Lightweight Annotation Interface

You really don't need Label Studio or some fancy annotation platform for this. You only need a Python script that shows a response and asks for a score.

Here's how this works at a high level: the script takes a question-response pair, displays it in the terminal, asks the reviewer to score it on a 1-5 scale, and saves the result to a file. Each annotation gets stored as a single line of JSON called JSONL format which makes it easy to load back later, run analysis on, or feed into a dashboard.

import json
import random
from pathlib import Path
from dataclasses import dataclass, asdict


@dataclass
class Annotation:
    """A single human annotation for an LLM response."""
    question: str
    response: str
    annotator: str
    score: int
    notes: str


class AnnotationCollector:
    """Collects and stores human evaluations."""

    def __init__(self, output_file: str = "annotations.jsonl"):
        self.output_path = Path(output_file)

    def collect_annotation(
        self, question: str, response: str, annotator: str
    ) -> Annotation:
        """Present a question-response pair and collect a human score."""
        print("\n" + "=" * 60)
        print(f"QUESTION: {question}")
        print("-" * 60)
        print(f"RESPONSE: {response}")
        print("-" * 60)
        print("Score this response (1-5):")
        print("  1 = Terrible  2 = Poor  3 = Acceptable  4 = Good  5 = Excellent")

        while True:
            try:
                score = int(input("Score: "))
                if 1 <= score <= 5:
                    break
                print("Please enter a number between 1 and 5.")
            except ValueError:
                print("Please enter a valid number.")

        notes = input("Notes (optional, press Enter to skip): ").strip()

        annotation = Annotation(
            question=question,
            response=response,
            annotator=annotator,
            score=score,
            notes=notes,
        )
        self.save(annotation)
        return annotation

    def save(self, annotation: Annotation) -> None:
        """Append annotation to JSONL file."""
        with open(self.output_path, "a") as f:
            f.write(json.dumps(asdict(annotation)) + "\n")

    def load_all(self) -> list[Annotation]:
        """Load all saved annotations."""
        annotations = []
        if self.output_path.exists():
            with open(self.output_path) as f:
                for line in f:
                    data = json.loads(line)
                    annotations.append(Annotation(**data))
        return annotations

Let me walk through what's happening in this script.

The Annotation dataclass is just a container that holds everything about a single review, the original question, the model's response, who reviewed it, the score they gave, and any notes they added. Nothing fancy, but having a structured format means you can easily compare scores across reviewers later.

The collect_annotation method is where the actual review happens. It prints the question and response to the terminal with some visual separators so the reviewer can read them clearly then prompts for a score.

The while true loop with input validation is important here. It keeps asking until the reviewer gives a valid number between 1 and 5 so you don't end up with garbage data in your annotations file.

The save method appends each annotation as a single JSON line to an annotations.jsonl file. I'm using JSONL (one JSON object per line) instead of a regular JSON array because it's append-friendly. You can add new annotations without reading and rewriting the entire file, which matters when you're collecting hundreds of reviews over time.

And load_all reads everything back, parsing each line into an Annotation object. This is what you'd call when you want to analyze your annotations, compare them to your LLM judge scores, or calculate agreement between reviewers.

In practice, you'd use this by feeding it a batch of question-response pairs from your production logs or golden dataset. You might run it during a weekly review session where a team member spends 30 minutes scoring 20-30 responses. That small investment gives you a reliable ground truth to calibrate your automated layers against.

How to Calculate Inter-Annotator Agreement

Now here's a problem you'll hit quickly: you ask two people to score the same response and they give it different scores. Is the response ambiguous or is your rubric ambiguous?

You need a way to measure this, and Cohen's Kappa is the standard tool for that. It basically tells you how much two annotators agree, adjusted for the amount of agreement you'd expect just by chance.

from sklearn.metrics import cohen_kappa_score


def measure_agreement(
    scores_annotator_1: list[int], scores_annotator_2: list[int]
) -> dict:
    """Calculate inter-annotator agreement using Cohen's Kappa."""
    kappa = cohen_kappa_score(scores_annotator_1, scores_annotator_2)

    interpretation = "poor"
    if kappa > 0.8:
        interpretation = "almost perfect"
    elif kappa > 0.6:
        interpretation = "substantial"
    elif kappa > 0.4:
        interpretation = "moderate"
    elif kappa > 0.2:
        interpretation = "fair"

    exact_agreement = sum(
        a == b for a, b in zip(scores_annotator_1, scores_annotator_2)
    ) / len(scores_annotator_1)

    return {
        "cohens_kappa": round(kappa, 3),
        "interpretation": interpretation,
        "exact_agreement": round(exact_agreement, 3),
    }


if __name__ == "__main__":
    # two annotators scored the same 10 responses
    annotator_a = [5, 4, 3, 4, 5, 2, 3, 4, 5, 4]
    annotator_b = [5, 4, 4, 4, 5, 3, 3, 4, 5, 3]

    agreement = measure_agreement(annotator_a, annotator_b)
    print(f"Cohen's Kappa: {agreement['cohens_kappa']}")
    print(f"Interpretation: {agreement['interpretation']}")
    print(f"Exact Agreement: {agreement['exact_agreement']:.0%}")

You would want a Kappa above 0.6. Anything below that and your rubric is the problem, not your annotators. Go back and add more concrete examples to each score level. Keep refining until people consistently agree. It usually takes two or three rounds of iteration.

How to Build the Regression Testing Pipeline

We can look at a scenario that's probably happened to you or other people you know: you tweak a prompt to fix one bad output you noticed. It works and that specific output is better now. You later ship it and week later, you find out the change broke three other responses you never thought to check.

This is incredibly common. The only way out is regression testing. If you've done traditional software development, you might already know what regression testing means. It's the practice of re-running a fixed set of tests every time you make a change, specifically to make sure you didn't break something that was already working.

The word regression literally means going backwards: your system was handling a question correctly and now after your change, it isn't.

In regular software, regression tests are usually unit tests or integration tests. For LLM applications, it works a bit differently. Instead of checking for exact outputs, you're scoring a batch of responses and comparing those scores against a previous run. If the scores drop, something regressed. The idea is the same but the mechanism is built around scoring rather than pass/fail assertions.

How to Create Golden Datasets

A golden dataset is just a curated list of questions that represent what your app actually needs to handle. You run your system against this list every time something changes (new prompt, new model, or updated retrieval logic) and compare the scores to your last run.

import json
from pathlib import Path
from dataclasses import dataclass, asdict


@dataclass
class GoldenExample:
    """A single test case in the golden dataset."""
    id: str
    question: str
    reference_answer: str
    category: str
    difficulty: str  # "easy", "medium", "hard"
    criteria: list[str]  # which criteria to evaluate


class GoldenDataset:
    """Manages a curated evaluation dataset."""

    def __init__(self, filepath: str = "golden_dataset.json"):
        self.filepath = Path(filepath)
        self.examples: list[GoldenExample] = []
        if self.filepath.exists():
            self.load()

    def add(self, example: GoldenExample) -> None:
        """Add a new example to the dataset."""
        self.examples.append(example)
        self.save()

    def get_by_category(self, category: str) -> list[GoldenExample]:
        """Filter examples by category."""
        return [e for e in self.examples if e.category == category]

    def save(self) -> None:
        """Persist dataset to disk."""
        data = [asdict(e) for e in self.examples]
        with open(self.filepath, "w") as f:
            json.dump(data, f, indent=2)

    def load(self) -> None:
        """Load dataset from disk."""
        with open(self.filepath) as f:
            data = json.load(f)
            self.examples = [GoldenExample(**item) for item in data]

    def summary(self) -> dict:
        """Return dataset statistics."""
        categories = {}
        for e in self.examples:
            categories[e.category] = categories.get(e.category, 0) + 1
        return {
            "total_examples": len(self.examples),
            "categories": categories,
        }

Some things I've learned about building these is that you should start with 50 to 100 examples. That's enough to catch meaningful regressions without making each eval run take forever.

Also, make sure you include edge cases – those weird questions that tripped up your model before. If your dataset is 90% easy questions, you won't notice when hard questions start failing.

Finally, treat this as a living document. Every time something breaks in production, turn it into a golden dataset example. Over a few months, your dataset evolves from generic test questions into a detailed map of exactly where your app is fragile.

How to Run Evaluations in CI/CD

Now let's wire everything together. This RegressionPipeline class runs your system against the golden dataset, scores every response, and compares the results to a previous run.

import json
from datetime import datetime, timezone
from dataclasses import dataclass, asdict


@dataclass
class EvalRun:
    """Records the results of one full evaluation run."""
    run_id: str
    timestamp: str
    model: str
    prompt_version: str
    total_examples: int
    avg_scores: dict  # criterion -> average score
    pass_rate: float  # percentage of examples above threshold
    failures: list[dict]  # examples that scored below threshold


class RegressionPipeline:
    """Runs evaluation against golden dataset and detects regressions."""

    def __init__(
        self,
        deterministic_eval: "DeterministicEvaluator",
        llm_judge: "LLMJudge",
        threshold: float = 3.5,
    ):
        self.det_eval = deterministic_eval
        self.judge = llm_judge
        self.threshold = threshold

    def run(
        self,
        golden_dataset: "GoldenDataset",
        generate_fn: callable,
        model_name: str,
        prompt_version: str,
    ) -> EvalRun:
        """Run full evaluation pipeline against golden dataset.

        Args:
            golden_dataset: The dataset to evaluate against.
            generate_fn: A function that takes a question string and
                         returns the model's response string.
            model_name: Identifier for the model being tested.
            prompt_version: Identifier for the prompt version.
        """
        all_scores = {}
        failures = []

        for example in golden_dataset.examples:
            # Generate response
            response = generate_fn(example.question)

            # Layer 1: Deterministic checks
            det_results = self.det_eval.run_all(response)
            det_failures = [r for r in det_results if not r.passed]

            if det_failures:
                failures.append({
                    "id": example.id,
                    "question": example.question,
                    "layer": "deterministic",
                    "details": [r.details for r in det_failures],
                })
                continue

            # Layer 2: LLM judge
            judge_results = self.judge.evaluate_all(
                example.question, response, example.criteria
            )

            for result in judge_results:
                if result.criterion not in all_scores:
                    all_scores[result.criterion] = []
                all_scores[result.criterion].append(result.score)

                if result.score < self.threshold:
                    failures.append({
                        "id": example.id,
                        "question": example.question,
                        "layer": "llm_judge",
                        "criterion": result.criterion,
                        "score": result.score,
                        "reasoning": result.reasoning,
                    })

        avg_scores = {
            criterion: sum(scores) / len(scores)
            for criterion, scores in all_scores.items()
        }

        total_evaluated = len(golden_dataset.examples)
        pass_count = total_evaluated - len(failures)

        return EvalRun(
            run_id=f"eval_{datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')}",
            timestamp=datetime.now(timezone.utc).isoformat(),
            model=model_name,
            prompt_version=prompt_version,
            total_examples=total_evaluated,
            avg_scores=avg_scores,
            pass_rate=pass_count / total_evaluated if total_evaluated else 0,
            failures=failures,
        )

    def compare_runs(self, baseline: EvalRun, current: EvalRun) -> dict:
        """Compare two evaluation runs to detect regressions."""
        regressions = {}
        improvements = {}

        for criterion in current.avg_scores:
            if criterion in baseline.avg_scores:
                diff = current.avg_scores[criterion] - baseline.avg_scores[criterion]
                if diff < -0.2:  # Score dropped by more than 0.2
                    regressions[criterion] = {
                        "baseline": baseline.avg_scores[criterion],
                        "current": current.avg_scores[criterion],
                        "change": round(diff, 3),
                    }
                elif diff > 0.2:
                    improvements[criterion] = {
                        "baseline": baseline.avg_scores[criterion],
                        "current": current.avg_scores[criterion],
                        "change": round(diff, 3),
                    }

        return {
            "verdict": "REGRESSION" if regressions else "PASS",
            "regressions": regressions,
            "improvements": improvements,
            "pass_rate_change": current.pass_rate - baseline.pass_rate,
        }

Now you can hook this into your CI/CD pipeline so it runs whenever someone changes a prompt or model config. If compare_runs returns REGRESSION, the build fails. No one deploys until they figure out what went wrong.

How to Know If Your AI Actually Got Better: Statistical Significance

So you tweaked your prompt and the average score went from 3.8 to 4.0. Time to celebrate, right? Maybe. Or maybe that 0.2 improvement is just random noise.

With a golden dataset of 50-100 examples, variance alone can easily produce score differences that big. You need an actual statistical test to know if the change is real.

A quick primer if you haven't done statistics in a while. A paired t-test is a way to compare two sets of measurements that are linked together. In our case, each pair is the same question scored under two different versions of your system: the old prompt and the new prompt.

The test looks at every pair, calculates how much the score changed for each question, and then asks: "Are these changes consistently in one direction or are they scattered randomly?"

If the changes are consistent (most questions scored higher with the new prompt), the test gives you a low p-value which means the improvement is likely real. If the changes are all over the place (some questions got better, some got worse, no clear pattern), the p-value will be high which means you can't be confident that the new version is actually better.

The reason we use a paired t-test instead of a regular one is that it accounts for question difficulty. Some questions are inherently harder than others, and pairing ensures we're measuring the change per question rather than just comparing two unrelated batches of scores.

Here's how to implement this:

from scipy import stats
import numpy as np


def is_improvement_significant(
    scores_before: list[float],
    scores_after: list[float],
    alpha: float = 0.05,
) -> dict:
    """Test whether a score improvement is statistically significant.

    Uses a paired t-test since the same questions are evaluated in both runs.
    """
    t_stat, p_value = stats.ttest_rel(scores_after, scores_before)
    mean_diff = np.mean(scores_after) - np.mean(scores_before)

    return {
        "mean_before": round(np.mean(scores_before), 3),
        "mean_after": round(np.mean(scores_after), 3),
        "mean_difference": round(mean_diff, 3),
        "p_value": round(p_value, 4),
        "is_significant": p_value < alpha,
        "direction": "improvement" if mean_diff > 0 else "regression",
        "recommendation": (
            "Safe to deploy"
            if p_value < alpha and mean_diff > 0
            else "Do not deploy - change is not a significant improvement"
        ),
    }


if __name__ == "__main__":
    # scores on 20 golden examples, before and after a prompt change
    before = [3, 4, 3, 5, 4, 3, 4, 4, 3, 5, 4, 3, 4, 3, 4, 5, 3, 4, 4, 3]
    after =  [4, 4, 4, 5, 5, 3, 4, 5, 4, 5, 4, 4, 4, 4, 5, 5, 4, 4, 5, 4]

    result = is_improvement_significant(before, after)
    print(f"Mean: {result['mean_before']} -> {result['mean_after']}")
    print(f"p-value: {result['p_value']}")
    print(f"Significant: {result['is_significant']}")
    print(f"Recommendation: {result['recommendation']}")

If the p-value comes back below 0.05, there's less than a 5% chance the improvement is just luck. That's when you ship. Anything above that and your improvement might just be noise, so don't deploy it no matter how good the averages look.

How to Put It All Together: The Complete Evaluation Architecture

Let's connect all three layers into a single orchestrator. This is the class that ties everything together. It runs deterministic checks first, escalates to LLM judging if those pass, and optionally brings in human evaluation for calibration.

class EvaluationOrchestrator:
    """Coordinates all three evaluation layers into a single pipeline."""

    def __init__(self):
        self.det_eval = DeterministicEvaluator()
        self.llm_judge = LLMJudge()
        self.annotation_collector = AnnotationCollector()

    def evaluate_response(
        self,
        question: str,
        response: str,
        run_human_eval: bool = False,
    ) -> dict:
        """Run the complete evaluation pipeline on a single response."""

        # Layer 1: Deterministic (always runs, every request)
        det_results = self.det_eval.run_all(response)
        det_passed = all(r.passed for r in det_results)

        if not det_passed:
            return {
                "status": "FAIL",
                "layer": "deterministic",
                "details": [r for r in det_results if not r.passed],
                "recommendation": "Fix structural issues before deeper eval",
            }

        # Layer 2: LLM Judge (runs on sample or in CI)
        judge_results = self.llm_judge.evaluate_all(question, response)
        avg_score = sum(r.score for r in judge_results) / len(judge_results)

        if avg_score < 3.5:
            return {
                "status": "FAIL",
                "layer": "llm_judge",
                "avg_score": avg_score,
                "details": judge_results,
                "recommendation": "Response quality below threshold",
            }

        # Layer 3: Human eval (periodic calibration)
        if run_human_eval:
            annotation = self.annotation_collector.collect_annotation(
                question, response, annotator="reviewer"
            )
            return {
                "status": "PASS" if annotation.score >= 4 else "REVIEW",
                "layer": "human",
                "automated_score": avg_score,
                "human_score": annotation.score,
            }

        return {
            "status": "PASS",
            "layer": "llm_judge",
            "avg_score": avg_score,
            "details": judge_results,
        }

What I Wish I Knew Earlier

I want to close with some things I wish someone had told me before I started building eval systems.

First, don't build all three layers at once. Start with just the deterministic checks, and then ship them. You'll be surprised how many issues they catch on their own, and the process of writing them forces you to actually define what correct output means for your app. Add the LLM judge when you need it and then add human eval later.

Second, check your judge against humans once a month. Run your LLM judge on 20-30 responses that already have human scores. If the judge has drifted more than 0.5 points on average, something changed: maybe the judge model was updated, or maybe your rubric doesn't cover a new failure mode. Either way, you need to recalibrate.

Third, every production failure becomes a test case. This is maybe the most useful habit. Something breaks? Great, that's a new golden dataset example. Over a few months, your dataset stops being a generic test suite and becomes a detailed map of every way your app has ever failed.

And finally, don't chase perfect eval scores. I've seen teams tweak prompts endlessly to push their eval scores from 4.2 to 4.5 only to discover that their rubric had a blind spot and users were still unhappy. The scores are a tool, not a goal, so human evaluation exists to catch what the numbers miss.

Wrapping Up

We covered a lot of ground in this article, so let me bring it all together. The core problem is that LLM applications fail differently from traditional software. There's no crash, no error log, and no stack trace. Just a confident, well-formatted, wrong answer.

And because the outputs aren't deterministic, you can't test them with simple assertions. You need a different approach entirely.

That approach is a layered evaluation pipeline:

• Layer 1 (Deterministic Checks) handles the basics: is the output valid, the right length, and free of hallucinated URLs? These are fast, free, and catch more problems than you'd expect.

• Layer 2 (LLM-as-Judge) brings in semantic evaluation: is the response actually relevant, accurate, and complete? By giving a judge model a clear rubric with concrete scoring criteria, you get surprisingly reliable and automated quality scores.

• Layer 3 (Human Evaluation) keeps the whole system calibrated. A small batch of human reviews on a regular schedule catches the subtle issues that automated scoring misses, like tone, clarity, and the difference between "correct" and "genuinely helpful."

On top of those three layers, you learned how to build a regression testing pipeline with golden datasets so you can catch quality drops before they reach production. You also learned how to use statistical significance testing to make sure your improvements are real and not just noise.

If there's one thing I'd want you to take away, it's this: start small. Don't try to build all of this in a weekend. Drop the DeterministicEvaluator class into your project today: that takes five minutes and it'll immediately start catching things you're currently missing. Then add the LLM judge when you're ready for deeper evaluation. Then layer in human review and regression testing as your app matures.

The teams that ship reliable AI products aren't the ones with the fanciest models. They're the ones who built the scaffolding to know when those models are failing and who catch it before their users do.

Nine months earlier, a solo Austrian engineer named Mario Zechner shipped something close to the opposite: a coding agent called Pi with four built-in tools and almost nothing else. Pi took roughly a year of organic growth to cross 91,600 stars, without the launch spike. Just a slow, compounding climb from engineers who tried it and stayed (earendil-works/pi).

So here we have two wildly different growth curves, with two wildly different design philosophies. And underneath both of them, we have the same word: harness.

If you build with AI agents in any capacity, that word is now unavoidable, and most explanations of it are either marketing copy or a diagram with too many arrows.

This article defines what an agent harness is, then compares ten of the most popular agent harnesses to date, from Claude Code to DeepSeek Harness to Pi, against the same five-part architecture.

By the end, you'll understand why the term replaced "framework" in developer conversation this year and how the loudest 2026 harnesses differ underneath their branding. You'll also have a 60-line Python harness to run yourself along with a breakdown of the stack layers around it (MCP, orchestration, observability), plus a decision guide for picking one for your team.

Table of contents

• What is an Agent Harness?

• From Agent Frameworks to Agent Harnesses: What Changed

• The Agent Harness Solutions at a Glance

• Three Competing Philosophies for How a Harness Should Work

• Build a Minimal Harness in Under 60 Lines of Python

• The Agent Harness Solution Stack

• Why the Hype Curve and the Adoption Curve Diverge

• How to Choose a Harness for Your Team

• What Transfers No Matter Which Harness Wins

• Conclusion

• What to Explore Next

What is an Agent Harness?

A harness is the runtime shell wrapped around an LLM model. The model itself only does one thing: given a stream of text and a list of available tools, it predicts what to say or which tool to call next. The harness handles everything else.

This unglamorous, boring plumbing includes the loop that calls the model, the code that executes tools, the memory that manages context over 40 turns, and the sandbox that protects your filesystem. It's the infrastructure that decides if an agent recovers from a failed tool call or just hangs.

Figure 1: The five parts every agent harness has to implement, drawn as a loop around a model core. The model sits at the center and predicts only the next message or tool call. Around it: a tool router that dispatches calls to the filesystem, shell, or external APIs, a memory layer that decides what context survives into the next turn, a planning layer that breaks a large task into steps before execution starts, and a sandbox boundary that constrains what the tools are allowed to touch.

The loop arrow shows the model's output feeding back in as the next turn's input, which is what turns a single prediction into an agent that keeps working until the task is done.

One practitioner definition captures the same shape from a different angle. A harness supplies everything a model doesn't do on its own: the loop that carries a goal from plan into action, access to tools like the terminal or file system, a memory layer that survives across turns, coordination for any subagents it spins up, and the permission rules that bound what it's allowed to touch (CellCog).

Concretely, when you type a request into Claude Code, Cursor, or Aider, here's what happens, in order:

1. The harness assembles a prompt: your request, the system instructions, and a list of tool schemas the model can call.

2. The model responds, usually with a mix of reasoning text and one or more tool calls (read_file, run_bash, edit, or whatever the harness exposes).

3. The harness executes each tool call, ideally inside a sandbox, and captures the output.

4. The harness appends the tool output back into the conversation and calls the model again.

5. The loop repeats, sometimes for dozens of turns, until the model produces a final answer, or until the harness hits a turn limit, a cost limit, or a human interrupts it.

That five-step loop, sometimes called the agent loop or the ReAct loop (after the 2022 paper that first described reasoning and acting as one interleaved process: Yao et al.), is the part every harness on the market shares.

What varies, and what determines whether a given harness is good at its job, is everything wrapped around step 3 and step 4: how good the planning is before execution starts, how the memory decides what to keep and what to drop as the context fills up, how isolated the sandbox is, and whether the harness can spin up a second, smaller version of itself to handle a sub-task without polluting the main conversation.

When any one of those four goes wrong, the symptoms look identical from the outside: the agent stalls, forgets what it was doing, or burns through your context window on a task that should take five turns.

From Agent Frameworks to Agent Harnesses: What Changed

The word "framework" dominated agent conversation from 2023 through 2025: tools like LangChain, AutoGen, and CrewAI. Frameworks in that era were libraries. You imported components, chose your own model calls, and wrote the orchestration logic yourself. They gave you building blocks.

A harness is a different kind of product. It ships the loop already built, and that loop is opinionated about memory, planning, and safety. You then interact with it by running a command.

Anthropic's Claude Code made this shift undeniable through 2025: a terminal-native agent that plans, edits files, runs tests, and commits code without you writing any orchestration logic.

By 2026, the ship-the-loop pattern showed up across the ten harnesses profiled in the table below, from Claude Code to DeepSeek Harness to Cline, and "harness" became the word everyone started using to describe that shape, distinct from a framework you assemble yourself.

You can see it in the naming: DeepSeek's own repository is called deepseek-harness, echoing the same framework-to-harness shift that Claude Code introduced.

LangChain's Deep Agents shows that the industry now treats "harness" as its own architectural layer, released as an attempt to reverse-engineer what made Claude Code's harness effective and rebuild it as an open, model-agnostic library.

LangChain's own account of the project traces it back to one question, in Harrison Chase's words: "What about Claude Code made it general purpose, and could we abstract out and generalize those characteristics?"

LangChain has an obvious incentive here too: it's pitching an alternative to the tool it's studying, and the four mechanisms it names still hold up regardless of who names them.

Deep Agents packages four specific mechanisms that Claude Code's harness relies on:

• A planning tool that forces the model to write out its steps before touching any files. This cuts down on the model quietly drifting off task over a long session.

• A virtual filesystem and sandbox that gives the agent structured, isolated read and write access to a repository.

• Subagent delegation, where the main agent spins up a smaller agent with its own clean context window to handle an isolated piece of work, then reports back a summary.

• Context and memory management, including middleware that compresses conversation history and offloads large tool outputs so a long session doesn't blow through the model's context window (LangChain).

That list is worth memorizing because those four mechanisms (planning, sandboxing, delegation, and context management) are the engineering problems every serious harness has to solve, whether or not Deep Agents remains the harness people point to. Everything else is branding.

The Agent Harness Solutions at a Glance

The table below covers the harnesses pulling the most developer attention as of August 2026 and what each one bets its architecture on.

A few of these are coding-specific, and a few (such as CellCog and Hermes Agent especially) are trying to generalize the harness pattern past code and into broader knowledge work.

A harness built for coding can assume a repository, a test suite, and a diff as its unit of work. A harness built for general knowledge work has to invent an equivalent structure for research, writing, and multi-step business tasks, which is a harder, less standardized problem.

If you're evaluating a harness for anything beyond code, ask first: what's its unit of work, and did anyone build the equivalent of a diff for it, or just assume one exists?

Three Competing Philosophies for How a Harness Should Work

Strip away the marketing, and three different engineering bets sit underneath the 2026 agent harness boom.

Figure 2: Three bets on how to build a harness, shown as three parallel columns. Column one, DeepSeek Harness, centers on a plugin kernel where models, sandboxes, memory, and the UI are all interchangeable modules.

Column two, Claude Code and Deep Agents, centers on four fixed mechanisms: planning, virtual filesystem, subagents, and context compression.

Column three, Hermes Agent, centers on a compounding skill library that grows every time the agent solves something new. The three columns share only the base loop from Figure 1. Everything above that loop is a different bet on what makes an agent reliable over long sessions.

Bet One: Everything is a Plugin.

DeepSeek Harness is built on a meta-framework called Cordis, whose design is described in DeepSeek's own paper "A Programming Paradigm for Spatiotemporal Composability," which boils down to one idea: everything can be swapped at runtime (deepseek-ai/deepseek-harness).

In practice, that means the model, sandbox, session storage, scheduling loop, and even the UI theme are all swappable modules. The harness also ships a "creator mode" for inspecting the running system, testing Cordis plugins in memory, and combining them into new configurations (DeepSeek).

The bet here: no single architecture wins forever, so the winning move is to make architecture itself a configuration file.

Bet Two: a Small, Fixed Set of Mechanisms, Executed Well.

Claude Code and, following it, LangChain's Deep Agents bet the opposite way: pick four mechanisms (planning, sandboxed filesystem access, subagent delegation, and context compression) and invest in making each one reliable.

Every mechanism on this list is familiar enough that rivals borrow it wholesale: the table above credits Claude Code with popularizing the planning-plus-subagent pattern other harnesses now copy. The bet works because all four run together on every task. Skip one, and the others cover for it, for a while, until a long session finds the gap.

Bet Three: Memory That Compounds.

Hermes Agent bets that the biggest unsolved problem is what happens between sessions. Most harnesses reset to a blank context on every new conversation. Hermes instead offers to save the approach as a reusable skill when it solves something non-trivial. It then checks that skill library before reasoning from scratch on a similar future request so it can get faster at recurring tasks the longer you use it (Nous Research).

That's an advantage, as well as a risk: a skill library that grows unchecked can turn into debt that outlives the reason it was written. Paired with native scheduling and channel integrations across platforms like Telegram, Slack, and Discord, the design goal is closer to a standing assistant that lives on a server than a tool you open for one session and close.

A fourth bet sits underneath all three: Pi and Oh-My-Pi argue that most of what the other harnesses build in is unnecessary weight, and that four tools plus an extension system beat a feature-complete platform for engineers who know what they want.

Pi's climb past 91,600 GitHub stars, driven by organic word of mouth rather than a launch campaign, suggests that bet has staying power.

All four bets are defensible. They optimize against different failure modes: DeepSeek Harness optimizes against architectural lock-in, Claude Code and Deep Agents optimize against unreliable long-session behavior, Hermes optimizes against repeated work across sessions, and Pi optimizes against bloat.

So before you pick one, ask which failure mode costs you time today. The answer will help you choose the correct agent harness.

Build a Minimal Harness in Under 60 Lines of Python

The example below builds the five-step loop from Figure 1 with Anthropic's Messages API: a model, three tools, and a loop that keeps calling the model until it stops asking for tool calls. You'll see every failure mode this section talks about waiting inside these 60 lines.

import subprocess
from anthropic import Anthropic

client = Anthropic()

TOOLS = [
    {
        "name": "read_file",
        "description": "Read a UTF-8 text file from the working directory.",
        "input_schema": {
            "type": "object",
            "properties": {"path": {"type": "string"}},
            "required": ["path"],
        },
    },
    {
        "name": "write_file",
        "description": "Write content to a file, overwriting it if it exists.",
        "input_schema": {
            "type": "object",
            "properties": {
                "path": {"type": "string"},
                "content": {"type": "string"},
            },
            "required": ["path", "content"],
        },
    },
    {
        "name": "run_bash",
        "description": "Run a shell command inside the sandbox directory and return its output.",
        "input_schema": {
            "type": "object",
            "properties": {"command": {"type": "string"}},
            "required": ["command"],
        },
    },
]

def execute_tool(name, tool_input):
    if name == "read_file":
        return open(tool_input["path"]).read()
    if name == "write_file":
        with open(tool_input["path"], "w") as f:
            f.write(tool_input["content"])
        return f"wrote {len(tool_input['content'])} bytes to {tool_input['path']}"
    if name == "run_bash":
        result = subprocess.run(
            tool_input["command"],
            shell=True,
            cwd="./sandbox",
            capture_output=True,
            text=True,
            timeout=30,
        )
        return result.stdout + result.stderr
    raise ValueError(f"unknown tool: {name}")

def run_harness(task, max_turns=15):
    messages = [{"role": "user", "content": task}]

    for _ in range(max_turns):
        response = client.messages.create(
            model="claude-sonnet-5",
            max_tokens=4096,
            tools=TOOLS,
            messages=messages,
        )
        messages.append({"role": "assistant", "content": response.content})

        if response.stop_reason != "tool_use":
            return response.content[0].text

        tool_results = []
        for block in response.content:
            if block.type == "tool_use":
                output = execute_tool(block.name, block.input)
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": output,
                })
        messages.append({"role": "user", "content": tool_results})

    return "stopped: hit max_turns without a final answer"

Run run_harness("Write a Python script in sandbox/hello.py that prints the first 10 Fibonacci numbers, then run it and show me the output.") and watch the turns unfold: the model writes the file, calls run_bash to execute it, reads the output, and only then produces a final text answer. Every production harness in the tables above is a more engineered version of this same shape.

Claude Code adds a planning step before turn one and a permission gate before every run_bash equivalent. Deep Agents adds a virtual filesystem, plus a middleware layer that compresses messages before it grows past the model's context window. DeepSeek Harness makes the TOOLS list and the model client themselves swappable at runtime.

The gap between this toy loop and a serious one sits entirely in reliability engineering: what happens when a tool call fails, what happens at turn 50, and what stops the sandbox from touching anything outside ./sandbox.

Nothing in execute_tool catches a malformed response or a tool that errors out, so a single bad tool call can loop the model back onto the same broken result turn after turn. Add a retry path yourself, or the harness keeps doing this by default.

Two things in this example deserve a closer look. First, cwd="./sandbox" is a load-bearing safety boundary: without it, run_bash can execute anything the host user can, which is why every serious harness runs tool execution inside a container or a restricted directory. It's an easy line to delete by accident during a refactor, and a dangerous one to lose.

Second, max_turns=15 exists because nothing here tells the model to stop on its own. If you skip it, a harness with no turn limit and no cost limit will keep looping and keep spending tokens for as long as the model keeps asking for tools. If you forget that line during a refactor, the failure looks identical from the outside: a job that never returns, and a token bill that keeps climbing until someone kills the process by hand.

The Agent Harness Solution Stack

A harness doesn't run alone. Three adjacent layers show up in almost every production agent deployment, and knowing where each one starts and stops keeps you from asking a harness to solve a problem that belongs one layer over. Skip that mapping, and you'll spend a week debugging the harness for a bug that lives in the sandbox instead.

Figure 3: Four horizontal layers, stacked bottom to top. The bottom layer, the protocol layer, is the Model Context Protocol (MCP). This is the shared standard that lets any harness talk to any external tool or data source the same way. The second layer up is the harness itself, the loop from Figure 1.

The third layer, orchestration frameworks, sits above single-agent harnesses and coordinates multiple agents or long-running stateful workflows: LangGraph, CrewAI, AG2, Mastra, and DSPy live here.

The top layer, drawn as two side panels rather than a fourth horizontal band, is observability and sandboxing: tools like Langfuse and LangSmith watch all the layers below them, and E2B and Modal provide the isolated execution environment the harness's sandbox runs inside.

The Protocol Layer: MCP

The Model Context Protocol is an open standard, originally introduced by Anthropic in November 2024, for connecting a model to external tools, files, and data sources in one consistent way (Anthropic). By late 2025, it had moved to the Agentic AI Foundation under the Linux Foundation, backed by Anthropic, OpenAI, and Block (Wikipedia).

A harness typically loads its tool list from an MCP config, not from code you write by hand:

{
  "mcpServers": {
    "filesystem": {
      "command": "npx",
      "args": ["-y", "@modelcontextprotocol/server-filesystem", "/Users/you/project"]
    },
    "postgres": {
      "command": "npx",
      "args": ["-y", "@modelcontextprotocol/server-postgres", "postgresql://localhost/mydb"]
    }
  }
}

Every MCP server you add here becomes available as tools inside TOOLS, without you writing a single new execute_tool branch. Write the integration once, and any MCP-compatible harness like Claude Code, Deep Agents, DeepSeek Harness, or one you build yourself can use it. That's the entire argument for the protocol layer in one sentence.

The Orchestration Layer

A harness runs one agent through a single loop. The moment you need multiple agents cooperating on a stateful, long-running workflow with dedicated roles like a planner, researcher, and reviewer, you enter orchestration framework territory. Choosing the wrong framework here will cost you months instead of a few lines of code.

1. LangGraph, which models a multi-agent workflow as a graph with checkpointing and time-travel debugging, and is widely used for stateful production workflows at regulated companies (GitHub)

2. CrewAI, built around defining agents by role and letting them collaborate on a shared task

3. AG2, the community-maintained successor to Microsoft's original AutoGen project, which moved in 2026 to an async, event-driven runtime built around composable middleware (GitHub, pickaxe.co)

4. Mastra, a TypeScript-first agent framework that crossed 22,000 GitHub stars and 300,000 weekly npm downloads after reaching version 1.0 in January 2026 (pickaxe.co, GitHub)

5. DSPy from Stanford NLP, which treats prompt engineering as something closer to compilation than hand-authorship, optimizing prompts against a metric (GitHub)

Observability and Sandboxing

Once an agent makes tool calls on its own, you need to see what it did and where it did it, to avoid debugging blindly. Langfuse and LangSmith trace every model call, tool call, and token cost across a session, which is how you debug a harness that failed on turn 34 (GitHub, LangChain).

Braintrust and Arize Phoenix add rigorous evaluation on top of that tracing, so you can regression-test a harness's behavior the same way you'd test a codebase (Braintrust, Arize-ai/phoenix). And for the sandbox itself, the isolated environment where run_bash-style tool calls execute, E2B and Modal provide disposable micro-VMs that let a harness run untrusted code without touching the host machine (GitHub, Modal).

Why the Hype Curve and the Adoption Curve Diverge

Figure 4: Two GitHub star growth curves plotted on the same axes over roughly 400 days. The DeepSeek Harness curve is nearly vertical: flat at zero, then a near-instant spike to 95,000-plus stars within the first two days after its August 13, 2026 launch, then flattening out.

The Pi curve is the opposite shape: a shallow, steady, almost straight-line climb from its August 2025 release to over 91,600 stars a year later, with no single spike anywhere on the line. Both curves end up in roughly the same place.

The point of putting these two curves on one chart is that the shape getting there is different for each: one curve reflects a coordinated launch and a well-timed announcement. The other reflects a year of engineers individually deciding, one at a time, that the tool was worth keeping installed.

A launch spike tells you a project generated attention. Sustained use tells you whether the tool is still open in a terminal six months later, and those are different questions with different causes.

DeepSeek Harness's 95,000 stars in two days is a verifiable number (Flowtivity), but it's also driven largely by timing, distribution, and a well-known model lab's existing audience.

Pi's climb to a similar star count carries a different kind of signal: nobody coordinated a launch for it a year in. It accumulated through word of mouth among engineers who tried a four-tool coding agent, kept using it, and told other engineers.

A tool picked off a launch-week spike can look just as capable on day one and still leave a team stranded three months later if the maintainers move on to the next announcement. Neither number outweighs the other, but if you're choosing a harness to bet a team's workflow on, research the curve's shape, not just its current height.

A steep spike with a flattening tail tells you a project has an active community forming, worth watching before you commit production workflows to it. A long, shallow, unbroken climb tells you engineers kept it installed after the excitement wore off, which is a stronger, if slower, signal.

How to Choose a Harness for Your Team

Match the harness to the failure mode in front of you, not whatever's trending this week. Picking based on stars instead of your bottleneck is the mistake that costs a team weeks of migration work later.

• You need one agent finishing one coding task reliably, end to end: Start with Claude Code, Deep Agents, or Aider if you want the tightest, most reviewable diff-per-commit loop you can get. All three implement the planning-plus-sandbox pattern from Figure 1 well.

• You're worried about vendor or architecture lock-in and expect to swap models frequently: DeepSeek Harness's plugin-everything design and Deep Agents' model-agnosticism both directly target this concern. A harness that hardcodes one provider's SDK into its core is the wrong choice here, regardless of how capable that provider's model is today.

• The same categories of tasks keep recurring across weeks or months, and you want the agent to get faster at them over time by building on what it knows: Hermes Agent's compounding skill library is built specifically for this pattern, especially if you also want it reachable from the chat platforms your team lives in.

• You want the smallest possible audit surface area, and you are comfortable writing your own extensions for anything missing: Pi's four-tool core, or Oh-My-Pi if you specifically want IDE-grade tooling, LSP diagnostics, and a debugger, layered on top of that same minimal foundation.

• You need several agents coordinating on a long-running, stateful process: That question sits a layer above the harness. Move up to LangGraph, CrewAI, AG2, or Mastra.

Whichever you pick, treat the observability layer as non-optional from day one. A harness that fails on turn 30 of an unattended run is a debugging nightmare without a trace. The same failure with Langfuse or LangSmith attached turns into a five-minute fix. Skip this step to save an afternoon of setup, and you'll pay for it the first time an agent fails mid-run, and nobody can say why.

What Transfers No Matter Which Harness Wins

The specific tool names in this article will likely look dated within a year, because the category is moving this fast. What will stay useful is the five-step loop in Figure 1, the four mechanisms LangChain identified inside Claude Code's architecture, and the layered stack in Figure 3.

Read any new harness that shows up next month against those three references, and you'll know within an hour whether it's doing something structurally new or repackaging the same loop under a different plugin system and a louder launch post.

That's the skill worth keeping: reading architecture instead of reading marketing, the one thing this category can't make obsolete no matter how fast the tool names turn over.

Conclusion

An agent harness isn't a mysterious new category of software. It's the runtime shell that turns a model's next-token prediction into an agent that plans, acts, checks its own work, and keeps going until a task is finished.

Harnesses are built from five parts that show up in every implementation: a loop, a tool router, memory, planning, and a sandbox boundary. What changed in 2026 is scale.

Enough teams shipped competing implementations that the architectural differences between them became worth studying. The landscape now ranges from DeepSeek's plugin-everything kernel and Pi's radical minimalism to Hermes Agent's compounding skills and the four fixed mechanisms of Claude Code and Deep Agents.

The numbers from the last section back this up: 95,000 stars in two days and 91,600 stars in a year prove two different routes reach the same conclusion.

Build the 60-line version yourself. Watch it loop. After you do, every harness on the market stops looking like magic and starts looking like an engineering decision you can evaluate on its merits.

What to Explore Next

• DeepSeek Harness on GitHub: read the README for the Cordis plugin architecture in the project's own words.

• LangChain's Deep Agents context-engineering docs: how the automatic compression and offloading middleware referenced above works under the hood.

• The Model Context Protocol specification: the protocol layer every harness in this piece can plug into.

• Hermes Agent's skills documentation: how a compounding skill library gets written and reused.

• LangGraph: the next layer up once one agent stops being enough.

• E2B: a concrete starting point for sandboxing tool execution off your host machine.

Visit my GitHub to explore the 30+ open-source software solutions and developer tools I built and shared using this agentic AI-native engineering process.

Open-weight Mixture-of-Experts models have expanded at a remarkable pace: Mixtral had about 47 billion total parameters, DeepSeek-V3 reached 671 billion, and Kimi K3 entered the trillions. The surprising part isn't simply how large these models became, but how little of each model processes any one token.

Kimi K3 has 2.8 trillion parameters, but it uses only about 104 billion of them for any single token. In almost every layer, a small router picks 16 of 896 specialized feed-forward networks, called experts, while two shared experts process every token.

This article focuses on that width-side design: how a model can offer a large pool of processing capacity without using all of it for every token. This differs from sequence memory, which concerns how the model stores and retrieves information from earlier tokens.

To see how K3 arrived at this design, we'll follow the evolution of Mixture of Experts (MoE) through four architectures.

Mixtral is a clear open-weight example of the basic pattern: route each token to a few full-size experts. DeepSeekMoE divided that work among finer-grained and shared experts. LatentMoE then compressed the routed path so those experts could work in a smaller space. Finally, K3 adopted it as Stable LatentMoE, adding mechanisms for numerical stability and balanced routing across 896 experts per layer.

Along the way, you'll learn how to interpret an MoE model's expert counts and active-parameter numbers, and what they imply for computation and data movement.

Activating only a small subset of experts is what makes MoE attractive, but it creates new bottlenecks. For example, the selected experts' weights still have to be read from GPU memory, and token representations may have to travel between GPUs. The architectures below are best understood as successive attempts to manage these costs.

Prerequisites

This is a conceptual article, so there's nothing to install or run.

• Helpful: familiarity with neural networks and the general shape of a Transformer layer: attention followed by a feed-forward network.

• Not required: prior knowledge of Kimi K3, Mixture-of-Experts routing, or distributed training. Each is introduced here.

• No code or tools needed.

Table of Contents

• Prerequisites

• 1. From One Dense Layer to a Mixture of Experts

• 2. How MoE Evolved Beyond Mixtral

• 3. LatentMoE: Compress the Expert Path

• 4. How Kimi K3 Makes LatentMoE Stable

• Conclusion: Making Sparse Capacity Usable

• References

1. From One Dense Layer to a Mixture of Experts

A Transformer layer performs two different kinds of work.

Attention lets tokens exchange information: the representation of one token (roughly, a word or piece of one) can incorporate information from other positions in the sequence.

The feed-forward network, or FFN, then transforms each token independently. By the time the token reaches the FFN, the relevant context has already been folded into its current numerical representation.

In a dense Transformer, every token goes through the same FFN. The FFN usually contains several large matrices and accounts for a substantial portion of the model's parameters and computation. Making it wider gives the model more capacity, but the added work falls on every token because each one passes through the whole FFN.

A Mixture of Experts, or MoE, changes that arrangement. Instead of one FFN, the layer contains several FFNs with different learned weights. A small router examines the token's current representation, scores the available experts, and selects the top few. Only those selected experts process the token, and their outputs are combined into the layer's result.

The final output is the weighted sum of only the selected experts' outputs, where each selected expert contributes in proportion to a router weight. Experts that aren't selected do no FFN computation for that token.

This routing happens independently in every MoE layer. Experts can and often do specialize, and the router learns which combination best fits a token in its current context. But those roles emerge during training: they may overlap and aren't guaranteed to match clean labels such as "Python" or "history." The same word can therefore take different routes in different contexts, and the same token may select different experts at different depths. Each layer also has its own expert pool, so expert 7 in one layer is unrelated to expert 7 in another.

A widely recognized open-weight example was Mixtral 8x7B. Each Mixtral layer contains eight FFN experts, and its router selects two for every token. The selected pair can change from token to token and layer to layer.

The name 8x7B is easy to misread. Mixtral is one Transformer, not eight complete 7-billion-parameter models. Its attention, embeddings, normalization layers, and other shared components exist only once.

What's repeated eight times inside each layer is the FFN: each expert has its own weights and the same 4,096 -> 14,336 -> 4,096 dimensions as the ordinary FFN in Mistral 7B. The router runs only two of those eight FFNs for each token. Mixtral therefore has about 47 billion parameters in total, not 56 billion, while about 13 billion parameters are active for one token. That active count includes both the two selected experts per layer and the model's shared parameters.

It's worth being precise about what this saves. A token still passes through two full-size FFNs in each layer, so Mixtral performs roughly twice the FFN computation of Mistral 7B's single FFN, but only one quarter of what running all eight experts would require. Its advantage isn't a smaller expert or necessarily less computation than the original dense model. It's access to a much larger pool of parameters without running that entire pool for every token.

Figure 1 (above): A dense layer sends every token through one FFN (left). A Mixture-of-Experts layer keeps many expert FFNs but activates only a few per token (right). Mixtral uses 2 of 8. Total parameters grow while the work per token stays much smaller.

2. How MoE Evolved Beyond Mixtral

Mixtral's eight experts are therefore eight standard-width FFNs, not eight smaller slices of one FFN. Selecting more of them would let a token combine more learned transformations, but each additional Mixtral-sized expert would add substantial computation. DeepSeekMoE asked whether the same compute budget could instead be divided among more, smaller experts.

An FFN usually expands the token vector into a wider internal layer, transforms it there, and then reduces it back to the token vector's original size.

For example, a 2,000-dimensional token vector might be expanded into an 8,000-dimensional internal representation, then projected back to 2,000 dimensions. It must return to 2,000 so its result can continue through the rest of the model.

DeepSeek makes an expert smaller by narrowing that internal layer. If the capacity of one large expert is replaced by four narrower experts, the router can select roughly four times as many while keeping the amount of expert computation similar. The token therefore receives contributions from several smaller FFNs instead of one or two large ones. This doesn't guarantee that every expert learns a clean specialty, but it gives training a finer set of building blocks to work with.

DeepSeekMoE adds a second idea: shared experts. Routed experts process only the tokens that select them, but a shared expert processes every token.

Think of the shared expert as common library code. If several routed experts all need the same general transformation, having each learn its own copy wastes parameters. The shared expert can learn that reusable work once and contribute it to every token, leaving the routed experts more room for transformations that differ across contexts.

DeepSeek calls this capturing common knowledge. The designers don't assign it a skill such as grammar or programming. Training decides what reusable work it learns.

Figure 2 (above): A conceptual illustration of DeepSeekMoE's two changes: replace a few large routed FFNs with more, smaller routed FFNs, and add an always-on shared FFN. The boxes are illustrative rather than DeepSeek-V3's literal expert count. DeepSeek-V3 uses one shared expert and 256 routed experts in each MoE layer, selecting eight routed experts per token.

DeepSeek-V3 scales this pattern up: most of its layers use one shared expert and 256 routed experts, with eight routed experts selected for each token. More experts create more available capacity, but they also make the physical execution of the model harder.

Two costs matter for the next step of our story.

The first cost appears inside a GPU. An expert's learned weights are matrices stored in the GPU's high-bandwidth memory, or HBM. Before the GPU can apply an expert, those matrices must be read by the hardware that performs the multiplications.

When many tokens use the same expert together, the GPU can reuse chunks of the expert's weights across many token calculations. When only a few tokens reach an expert, it must move a large amount of weight data for relatively little arithmetic, so its computing units may spend much of their time waiting for those bytes.

The bottleneck is then not how quickly the GPU can multiply numbers, but how quickly it can deliver the weights to its computing units. That delivery rate is called memory bandwidth, and it can determine the speed of low-latency MoE serving.

The second cost appears between GPUs. A model with hundreds of experts can't usually keep every expert on every GPU, so the expert pool is distributed across them.

Suppose a token is represented by a vector of 7,168 numbers, as in DeepSeek-V3. If its router selects experts stored on other GPUs, the system sends the entire 7,168-number vector to each selected expert's GPU. Each expert returns another vector of the same length, and those results are combined. This exchange of token vectors among many GPUs is called all-to-all communication.

This reveals what fine-graining does and does not solve. Narrowing the internal layer makes each individual expert smaller, but DeepSeek activates proportionally more of them to keep total expert computation roughly unchanged.

The token vector sent to each selected expert also stays the same length. Selecting more experts can therefore mean sending more complete copies between GPUs, even though each expert is smaller inside. Fine-graining creates a more flexible set of building blocks. It doesn't compress the route into and out of them.

That distinction motivates LatentMoE: what if the model compressed the token vector before sending it to the routed experts, performed the expert work in that smaller space, and expanded it again only after the results returned?

3. LatentMoE: Compress the Expert Path

LatentMoE was introduced by an NVIDIA research team and adopted in the Nemotron 3 model family. Kimi K3 didn't invent the underlying architecture. Rather, it adopts LatentMoE and adds the stability changes discussed in the next section.

The central move is straightforward. Before a token is dispatched to the routed experts, LatentMoE projects its full representation into a smaller latent space. The routed experts operate entirely in that smaller space. Their outputs are combined there and projected back to the model's full width afterward. Here, latent means the routed experts' compressed workspace.

The router still examines the original full-width token representation. The shared experts also remain full width. Only the path through the routed experts is compressed.

Figure 3 (above): LatentMoE compresses only the routed path. In Kimi K3, one shared down-projection changes the routed representation from 7,168 dimensions to 3,584 before it reaches the selected experts. Their weighted outputs are combined and normalized once, then one shared up-projection restores 7,168 dimensions. The router and two shared experts continue to use the original 7,168-dimensional representation.

This shorter routed interface cuts two costs at once. First, each routed expert's input and output matrices connect to 3,584 dimensions rather than 7,168, so they contain fewer weights and require less weight data to be read when the expert runs.

Second, when experts are spread across GPUs, the system sends a 3,584-dimensional vector to each selected expert instead of the original 7,168-dimensional one. The selected experts return vectors of the same shorter length, and those results are combined and projected back to 7,168 dimensions. In K3's half-width design, each routed message therefore carries half as many values.

LatentMoE can spend those savings in two ways: keep the same number of active experts and lower inference cost, or increase both the available experts and the number selected per token without letting weight movement and cross-GPU communication grow as they would at the original width.

This differs from DeepSeek's fine-graining, which narrows the middle of each FFN but leaves its entrance and exit unchanged. The two ideas are compatible because they shrink different dimensions.

Compression still has a limit. If the latent representation becomes too short, its down-projection may discard information the experts need, and the shared down- and up-projections add computation of their own.

The latent width is therefore a balance to strike, not a number to minimize: Kimi K3 settles on 3,584 dimensions, half its 7,168-dimensional model width. That makes the routed path cheap enough to widen the expert pool dramatically, which raises the next problem: keeping so many experts stable during training.

4. How Kimi K3 Makes LatentMoE Stable

Kimi K3 has 93 backbone layers. The first uses a dense FFN, while the remaining 92 use Stable LatentMoE. Each of those 92 layers has its own router and its own pool of 896 routed experts. The phrase "16 of 896" therefore describes a separate routing decision at every MoE layer, not one global pool shared across the whole model.

For a token arriving at one of those layers:

1. The router scores all 896 routed experts and selects 16.

2. Two full-width shared experts process the token regardless of that selection.

3. The routed path projects the token from 7,168 dimensions to 3,584.

4. The 16 selected latent experts process that smaller representation.

5. Their weighted outputs are combined, normalized, and projected back to full width.

6. The routed and shared results are added together.

This arrangement helps K3 place 2.8 trillion parameters in the model while activating about 104 billion for one token.

But the scale also magnifies three training problems. Stable LatentMoE adds one targeted mechanism for each.

Figure 4 (above): Stable LatentMoE addresses three separate problems: RMSNorm steadies the routed branch's scale, SiTU-GLU caps unusually large activations, and Quantile Balancing sets selection biases toward an even global load without changing the contribution weights of selected experts.

Problem 1: the combined expert output can vary in scale.

Different tokens select different expert combinations with different routing weights, so the overall magnitude of the combined routed representation can vary before it reaches the up-projection.

K3 inserts RMSNorm after the selected expert outputs are combined and before they're projected from 3,584 dimensions back to 7,168. RMSNorm doesn't make the experts identical or erase what they computed. It rescales their combined result so the up-projection receives an input with a more consistent overall magnitude.

Problem 2: two large internal values can multiply into an activation spike.

LatentMoE first uses a shared projection to compress the token from 7,168 to 3,584 dimensions. Inside each selected expert, two bias-free learned linear projections produce 3,072-dimensional pre-activations g (gate) and v (value). SiTU-GLU transforms the gate into 4 tanh(g/4) sigmoid(g) and the value into 25 tanh(v/25), then multiplies them element-wise. The 3,072-dimensional product passes through a third bias-free linear projection, which returns a 3,584-dimensional expert result for aggregation.

These three projections, the nonlinear transformations, and the multiplication together form one gated expert FFN.

Figure 5 (above): Orange projections belong to the shared LatentMoE wrapper, and blue projections belong to one selected expert. The expert performs its gated calculation in a temporary 3,072-dimensional workspace and returns a 3,584-dimensional result, the common shape required for expert aggregation.

Following either the gate branch or the value branch, a signal passes through exactly four learned matrix multiplications: the shared LatentMoE down-projection, that branch's expert input projection, the expert output projection, and the shared LatentMoE up-projection.

The K3 paper calls this "nearly four consecutive matrix multiplications" because the computation isn't one uninterrupted linear chain: the gate and value projections run in parallel and meet through SiTU and element-wise multiplication, then selected expert outputs are aggregated and normalized before the final projection. So the four matrix operations can't be collapsed into one matrix multiplication.

The K3 authors describe the combined structure as ill-conditioned and report exploding internal activations at their model's scale. In low precision, a large outlier can overflow or force a shared quantization scale to sacrifice accuracy for ordinary values.

The multiplication inside SwiGLU is one source of that unbounded growth. Its value branch produces candidate values, while its gate branch uses Swish to modulate how strongly each value passes through. The linear factor inside the Swish gate and the value branch can both grow without bound, so two large elements can produce a much larger product.

So K3 replaces SwiGLU with SiTU-GLU (Sigmoid Tanh Unit GLU). SiTU smoothly caps the gate's linear factor at magnitude 4 and the value branch at magnitude 25, while retaining the sigmoid gate and matching SwiGLU near zero. Their element-wise product is consequently bounded in magnitude by 4 x 25 = 100 before the expert output projection. That later projection can still change the scale, but the multiplication inside the expert is no longer unbounded.

Figure 6 (above): An illustrative one-dimensional slice in which both branch inputs equal the same scalar x. SiTU-GLU follows SwiGLU near the origin but approaches a magnitude limit of 100, while SwiGLU continues growing. In the real expert, separate learned projections produce the two branch vectors and combine them element by element.

Problem 3: routing can become uneven.

The router gives every expert an affinity score for each token, then selects the 16 highest-scoring experts. Because the router is learned, some experts can attract far more tokens than others. Those experts become hardware bottlenecks, while rarely selected experts receive too little training to become useful.

A common response is to add a balancing loss to the model's training objective, but that makes the optimizer trade prediction quality against even expert use.

DeepSeek-V3 instead made its primary global balancing method auxiliary-loss-free. It maintains a separate selection bias for each expert: after a training step, an underused expert's bias rises by a fixed amount, while an overloaded expert's bias falls by that amount.

The method works, but the update size must be chosen carefully. Too small reacts slowly, while too large can make the load oscillate. (DeepSeek-V3 also retained a small sequence-level balancing loss as a safeguard against extreme imbalance within one sequence.)

K3 keeps the expert-specific selection biases but replaces the fixed adjustment with Quantile Balancing. It examines how an expert's scores are distributed across the global training step and calculates a different adjustment for each expert: a larger correction when the scores indicate that more movement is needed, and a smaller one when the expert is already near its target load. It's called quantile balancing because the update is chosen from a target percentile of that expert's score margins, rather than moving every bias by the same preset amount.

The bias changes which experts are selected, not how strongly their outputs contribute. K3 ranks experts using the biased scores, but derives the selected experts' contribution weights from their original scores without the bias. The newly calculated biases take effect on the next training step, and the final biases are frozen during inference.

Quantile Balancing targets an even aggregate load across the global training step, not equal expert use within every sentence or sequence. Its purpose is narrower: keep training opportunities and distributed computation from concentrating on too small a part of the 896-expert pool.

Conclusion: Making Sparse Capacity Usable

A dense Transformer sends every token through the same FFN. Mixtral showed the basic MoE alternative: keep several full FFNs and route each token to only a few. DeepSeekMoE then divided that work into more, smaller routed experts and added shared experts for transformations used across many contexts.

LatentMoE changes a different dimension. Instead of sending the model's full token representation to every selected expert, it compresses the routed interface, performs the expert computation in that smaller space, and restores the original width afterward. This reduces both routed-expert weight traffic and the amount of token data exchanged between GPUs.

Kimi K3 pushes that design to 896 routed experts per MoE layer, with 16 selected for each token. At that scale, compression alone isn't enough. RMSNorm controls the scale of the combined routed result, SiTU-GLU bounds the multiplicative activation inside each expert, and Quantile Balancing distributes training assignments without adding the balancing bias to the experts' contribution weights.

The central lesson isn't simply that MoE activates fewer parameters. Increasing sparse capacity creates new numerical, routing, and communication constraints, and the architecture must address them together. Stable LatentMoE is K3's answer at the model level.

Extended reading: For a deeper look at the sequence-memory side of Kimi K3, see From GPT-2 to Kimi K3: How Language Models Learned to Manage Memory.

References

1. Jiang et al. (2024). Mixtral of Experts. arXiv:2401.04088

2. Dai et al. (2024). DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. arXiv:2401.06066

3. DeepSeek-AI et al. (2024). DeepSeek-V3 Technical Report. arXiv:2412.19437

4. Elango et al. (2026). LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts. arXiv:2601.18089

5. Kimi Team (2026). Kimi K3: Open Frontier Intelligence. arXiv:2607.24653

6. Wang et al. (2024). Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts. arXiv:2408.15664

If you come from general machine learning, you may think you already know what these words mean. And sometimes you do.

But medical imaging adds a few twists. Some terms have a different meaning, and some are used in more than one way depending on the context.

This article follows a chest X-ray from the moment it's acquired to the point where a model makes a prediction. Along the way, we'll look at the common terms you'll see in medical imaging papers and what they actually mean.

A companion notebook lets you run most of these steps yourself instead of just reading about them.

What We'll Cover:

• What You'll Learn

• From Image to Dataset

  • 1. Acquisition

  • 2. Anonymization, de-identification, and pseudonymization

  • 3. Safe Harbor and Expert Determination

• From Dataset to Model Input

  • 4. Preprocessing

  • 5. Normalization

  • 6. Annotation and label

  • 7. Dataset splitting

• From Model to Prediction

  • 8. Classification, detection, and segmentation

  • 9. Augmentation

  • 10. Postprocessing

• From One Hospital to the Real World

  • 11. Harmonization

  • 12. Retrospective and prospective validation

• Putting it together

• Conclusion

What You'll Learn

• What each stage of a medical imaging pipeline is called, and what actually happens at each stage

• The difference between anonymization, de-identification, and pseudonymization

• Why "annotation" and "normalization" can mean different things depending on the context

• How classification, detection, and segmentation differ on the same image

• What harmonization fixes, and why validation design can matter more than model choice

From Image to Dataset

1. Acquisition

Acquisition is when the image is created. It includes the imaging machine, its settings, and how the patient is positioned.

Two chest X-rays of the same patient may look different if they were taken using different machines. The manufacturer, detector, exposure settings, and image-processing software can all affect the final image.

Patient's position matters too. For example, a standing PA (Posteroanterior) chest X-ray can look very different from a portable AP (Anteroposterior) X-ray taken while a patient is lying in bed.

One important difference is the apparent size of the heart. This can affect what a model learns.

A lot of this information is stored in the DICOM file.

DICOM (Digital Imaging and Communications in Medicine) is the standard format used to store and communicate medical images. A DICOM file contains much more than pixels. It can also contain information about the patient, scanner, study, image orientation, pixel spacing, and other details.

The image used in this tutorial originally arrived as a JPEG, so the original clinical DICOM metadata wasn't available. For demonstration, the notebook wraps the image in a new DICOM file and adds a few useful fields.

ds = pydicom.dcmread("synthetic_cxr.dcm")

for tag in ["Modality", "BodyPartExamined", "ViewPosition",
            "Manufacturer", "ManufacturerModelName", "KVP", "PixelSpacing"]:
    print(f"{tag:24s} {ds[tag].value}")

One field worth paying attention to is PixelSpacing. It tells you the physical size represented by each pixel, usually in millimeters. Two images can both be 512 × 512 pixels but cover different physical areas.

So if you want to measure something in millimeters, you can't simply count pixels. You also need to know the pixel spacing.

2. Anonymization, De-identification, and Pseudonymization

These three terms are often used interchangeably, but they have different meanings.

De-identification

De-identification removes or changes information that could identify a person.

For example:

Patient name
Medical record number
Date of birth
Phone numbers
Other identifying information

The goal is to reduce the chance that the data can be linked back to a person.

Pseudonymization

Pseudonymization replaces an identifier with a code.

For example:

Jane Doe → SUBJ_0041

The important difference is that a separate key can still connect SUBJ_0041 back to Jane Doe.

Hospitals may need this because they sometimes need to find the original patient again.

Anonymization

Anonymization aims to make re-identification no longer reasonably possible.

Unlike pseudonymization, there's no retained key that can be used to reconnect the data to the person.

A simple question can help:

Can this data still be linked back to the person using additional information?

If the answer is yes because a separate key exists, you're generally dealing with pseudonymization rather than true anonymization.

The exact legal meaning depends on the country and regulation, so researchers should be careful when using these terms.

The pixels can contain identifiers, too

Removing information from the DICOM header is only one part of the job.

Sometimes text is actually drawn into the image.

For example:

Patient name
Date
Hospital name
R
PORTABLE

This is called burned-in annotation.

Removing the DICOM metadata does nothing to this text. You have to find the text in the pixels and remove it.

This is harder than it sounds. A simple rule such as "look for bright pixels near the top of the image" will find more than just text. Bones such as the clavicle and ribs are also bright.

That means a real de-identification pipeline usually combines several techniques:

Text detection
OCR
Image-region rules
DICOM metadata removal
Validation

The goal isn't just to remove text. It's to prove that the text was actually removed.

For example:

IDENTIFIERS = [
    "PatientName",
    "PatientID",
    "PatientBirthDate",
    "PatientSex",
    "InstitutionName",
    "ReferringPhysicianName",
    "StudyDate",
    "StudyTime",
    "AccessionNumber"
]

for tag in IDENTIFIERS:
    if tag in ds:
        ds[tag].value = ""

# Records that an identity-removal process was performed.
# This field alone does not make the dataset de-identified.
ds.PatientIdentityRemoved = "YES"

The pixel-level part is where a simple demo and a real production pipeline are very different.

3. Safe Harbor and Expert Determination

If you work with US healthcare data, you'll often see two HIPAA terms: Safe Harbor and Expert Determination.

They are two ways to determine whether protected health information has been de-identified under HIPAA.

Safe Harbor

Safe Harbor is a checklist. It requires removing 18 categories of identifiers, including things such as:

Names
Geographic information smaller than a state
Certain dates
Phone numbers
Medical record numbers
Device identifiers
Full-face photographs

It's relatively straightforward because you can follow a defined list.

Expert Determination

Expert Determination takes a different approach. A qualified expert evaluates the risk of re-identification and documents why the remaining risk is very small. This can allow researchers to keep information that Safe Harbor would require them to remove.

For example, exact dates can be very useful when studying how a disease changes over time.

So there's a trade-off:

Safe Harbor is simpler and more restrictive

Expert Determination is more flexible, but requires a documented risk assessment

These are US HIPAA concepts. Other countries have different privacy laws and definitions.

For example, the GDPR in the EU and India's DPDP Act have their own approaches to anonymous and pseudonymous data.

From Dataset to Model Input

4. Preprocessing

Preprocessing is what you do to an image before giving it to a model.

Common preprocessing steps include:

• Resizing

• Windowing

• Changing orientation

• Normalizing pixel values

Resizing

Models usually expect images of a fixed size. Medical images can come in many different sizes, so they're often resized.

But resizing changes the relationship between pixels and physical space.

For example, if you resize an image from 1024 × 1024 to 512 × 512, each pixel now represents a different physical area.

So if you need physical measurements, such as distances or areas in millimeters, you need to account for the original or resampled pixel spacing.

Windowing

Windowing selects a range of intensity values and maps that range to the display range. Values outside the range are clipped.

Radiologists commonly use different window settings when looking at CT images because different tissues become easier to see under different windows.

Windowing changes how the image is represented or displayed. It doesn't mean the original image data has been changed.

The notebook demonstrates windowing using the 2nd and 98th percentiles:

lo, hi = np.percentile(img, [2, 98])
windowed = np.clip(
    (img.astype(float) - lo) / (hi - lo),
    0,
    1
)

Orientation

Medical images also contain information about orientation. Getting this wrong means left and right are swapped, which, for a chest X-ray, is a serious mistake.

5. Normalization

Here's the first word with dual meanings.

In machine learning, normalization usually means transforming numerical values into a more consistent scale or distribution so that training behaves more predictably. Two common versions:

• Min-max: rescale so the lowest value becomes 0 and the highest becomes 1.

f = img.astype(np.float32)

minmax = (f - f.min()) / (f.max() - f.min())

• Z-score: subtract the mean and divide by the standard deviation, giving a mean of 0 and a standard deviation of 1.

  zscore = (f - f.mean()) / f.std()
  

Min-max normalization has one weakness. It depends directly on the minimum and maximum values. A very bright pixel, such as one caused by a device or noise, can change the range and squeeze most of the image into a smaller part of the scale.

Z-score normalization is less dependent on the minimum and maximum, although extreme values can still affect the mean and standard deviation.

In medical imaging, normalization can also be applied consistently across images or datasets, but making images from different scanners or institutions comparable is more specifically described as harmonization. The distinction becomes important when you are dealing with multiple sites, as explained in section 10.

6. Annotation and Label

Both describe information associated with an image, but they differ in what that information represents.

A label usually describes the image as a whole.

For example:

Image → Pneumonia

An annotation usually tells you where something is in the image.

It could be:

• A point

• Bounding box

• Polygon

• Contour

• Pixel-level mask

For example:

Image → Lung mask

The key difference is location.

A label tells you what is present while an annotation can tell you what is present and where it is.

Annotations also take much more effort to create.

A large dataset may have millions of image-level labels extracted from radiology reports, but only a small subset may have detailed masks created by clinicians.

That shortcut has a cost. A report may say "pneumonia," but that doesn't necessarily tell you exactly which pixels show pneumonia. This can lead to noisy labels.

7. Dataset Splitting

One of the most important decisions in a medical imaging study is how you split the data.

The split should usually happen at the patient level, not the image level.

Imagine a patient has five X-rays. If you put three images into training and two into testing, the model has already seen images from that patient during training. Patient-specific characteristics can appear in both sets. This is a form of data leakage.

The same idea applies to:

• Multiple scans from the same patient

• Multiple images from the same study

• Images derived from the same original scan

The goal is simple: the test set should contain patients the model didn't see during training.

So instead of:

Split images and assign patients

do this:

Split patients and assign their images

This is one of those details that can have a bigger effect on your results than changing the model architecture.

From Model to Prediction

8. Classification, Detection, and Segmentation

Take the same chest X-ray and ask three different questions.

Classification

Is there pneumonia?

The model returns a label or probability.

Pneumonia: 0.92

It tells you what is in the image. It doesn't tell you where.

Detection

Where is the abnormality? The model returns a bounding box.

[x, y, width, height]

It tells you roughly where the finding is. It doesn't describe its exact shape.

Segmentation

Which pixels belong to the abnormality?

The model returns a pixel-level mask. This gives you the shape and location of the finding.

A simple way to remember it:

Classification = What?

Detection = Where?

Segmentation = Which pixels?

The amount of annotation work usually increases as you move from classification to detection to segmentation.

So choose the simplest task that answers your question. If you only need to know whether a scan is abnormal, you probably don't need segmentation.

9. Augmentation

Augmentation creates new training examples by transforming existing images.

Common transformations include:

• Rotation

• Translation

• Zoom

• Brightness changes

This can be useful when medical datasets are small.

But medical images have an important constraint: the transformed image should still look like something that could realistically happen to a patient.

For example, a horizontal flip is common in natural-image machine learning.

Flip a photo of a cat and you still have a cat.

But flip a chest X-ray and the heart moves to the other side. You may have just created an image that looks like dextrocardia, where the heart is on the right side.

The same problem happens with markers. A right-side marker can suddenly appear on the left. The letter itself is also mirrored. A model doesn't automatically understand that this is anatomically wrong.

Other augmentations can cause problems, too. A large brightness change might hide important findings. An aggressive crop might remove part of the lungs.

So before using an augmentation, ask: could this image realistically come from a real scanner and a real patient?

If not, don't use it.

# Reasonable: small rotation
M = cv2.getRotationMatrix2D((cx, cy), 7, 1.0)
rotated = cv2.warpAffine(
    img,
    M,
    (w, h),
    borderMode=cv2.BORDER_REPLICATE
)

# Potentially problematic for a chest X-ray:
flipped = img[:, ::-1]

10. Postprocessing

Postprocessing happens after the model produces its output.

A segmentation model usually produces a probability for every pixel. After applying a threshold, the mask may contain small unwanted regions or holes.

For example, the notebook's raw mask contains:

• Two large lung regions

• Fourteen small unwanted regions

A common cleanup step is to keep only the largest connected components.

Since we expect two lungs, we can keep the two largest regions. We can also fill small holes.

labelled, n = ndimage.label(mask > 0)

sizes = ndimage.sum(
    mask > 0,
    labelled,
    range(1, n + 1)
)

keep = np.isin(
    labelled,
    np.argsort(sizes)[-2:] + 1
)

clean = ndimage.binary_fill_holes(keep)

In this example, the cleanup reduces the number of connected components from 16 to 2 while changing less than 5% of the pixels.

That illustrates an important point about evaluation: pixel-overlap metric might barely change, even though the structure of the prediction has changed significantly.

For some applications, the number and shape of connected regions matter more than a small change in pixel overlap. So choose metrics that match the errors you actually care about.

Also, be careful with postprocessing. If your model needs a lot of cleanup before the result looks good, the cleanup may be hiding problems in the model.

Always look at the raw output, too.

From One Hospital to the Real World

11. Harmonization

Imagine two hospitals use different scanners.

The images may have different:

• Brightness

• Contrast

• Noise

• Resolution

• Image-processing characteristics

A model trained mostly on Hospital A may learn some of these differences instead of learning features related to the disease. It may perform well on Hospital A but fail on Hospital B.

This is where harmonization can help.

Harmonization tries to make data from different sources more comparable while preserving the information that matters.

One simple example is histogram matching.

It adjusts the pixel-value distribution of one image to look more like a reference image.

from skimage.exposure import match_histograms

harmonized = match_histograms(
    image_from_hospital_b,
    reference_from_hospital_a
)

This can help with brightness and contrast differences, but it can't solve everything.

If two scanners have different resolution, noise characteristics, or image-processing pipelines, histogram matching alone isn't enough.

And if information was lost during image acquisition or clipping, harmonization can't magically recover it.

A useful rule of thumb is:

Normalization makes numerical values more consistent.

Harmonization addresses systematic differences between data sources.

12. Retrospective and Prospective Validation

This section is about study design.

Retrospective

A retrospective study uses data that already exists.

For example:

We collected 4,000 chest X-rays from the hospital archive and tested our model on them.

This is common in medical AI because it's relatively fast and inexpensive.

But it also creates opportunities for bias. Researchers may make decisions about which patients to include, which scans to exclude, which hospital to use, and which threshold to choose.

These decisions can unintentionally make the results look better.

Prospective

In a prospective study, you define the study plan first and then collect data going forward.

For example:

We define the patient population, evaluation criteria, and success metrics before collecting the images.

This can reduce some sources of bias because important decisions are made before seeing the results.

External validation

You'll also see the term external validation.

This means testing the model on data that is independent of the data used to develop it.

For example:

Hospital A → Training
Hospital A → Internal test
Hospital B → External validation

An even stronger test might be:

Hospital A + B → Development
Hospital C → External validation

A model that performs well on its own hospital's data but poorly at another hospital may have learned site-specific patterns instead of general disease features.

External validation is therefore often much more informative than simply creating another random split from the same dataset.

Putting it Together

Now consider this sentence from a hypothetical medical imaging paper:

We retrospectively collected 4,120 de-identified frontal chest radiographs from two institutions. Images were resampled to 512 × 512, intensity-normalized, and harmonized across sites by histogram matching. Lung fields were manually annotated by two radiologists; segmentation output was postprocessed by largest-component selection. The model was externally validated on 890 studies from a third site.

That paragraph contains a lot of information, but now you can decode it:

1. Retrospectively collected: The researchers used images that already existed.

2. De-identified: Identifying information was removed or changed.

3. Two institutions: The dataset comes from more than one source.

4. Resampled to 512 × 512: Images were converted to a common image size.

5. Intensity-normalized: Pixel values were transformed to a more consistent numerical scale.

6. Harmonized: The researchers tried to reduce systematic differences between the two sites.

7. Manually annotated by two radiologists: Experts created spatial information showing where the lung fields are.

8. Postprocessed: The model's raw segmentation output was cleaned up.

9. Externally validated: The model was tested on independent data from another site.

And that last part is especially important. A model that works well only on the data it was developed on tells you much less about how it will perform in the real world.

Conclusion

A medical imaging paper can describe an entire data pipeline in just a few sentences.

Once you understand the terminology, you can start reading those sentences differently.

Instead of just looking at the model and its accuracy, you can ask:

• Where did the images come from?

• What scanner was used?

• Were patients kept separate between training and testing?

• How were identifiers removed?

• Could there be burned-in text?

• Who created the labels or annotations?

• What preprocessing was performed?

• How were pixel values normalized?

• Were different hospitals or scanners harmonized?

• Was the model externally validated?

• Was the test data truly independent?

• What happened to the model's output after prediction?

The model is only one part of a medical imaging pipeline. Very often, the more important questions come before the model ever sees an image.

Your watch doesn't actually see a "step" or directly measure "sleep." Instead, tiny sensors continuously detect things such as movement, changes in blood flow, electrical activity, and temperature. Those signals are converted into digital data, processed to remove noise, and analyzed by algorithms that look for meaningful patterns.

This process happens quietly in the background. Every movement of your wrist can become a stream of numbers. A change in reflected light can become a heart-rate reading. Several different signals can be combined to estimate what you were doing or how your body was responding.

By the time that information reaches the screen, the original signal has already gone through several layers of processing, turning something the sensor can detect into something you can understand.

What We'll Cover:

• What Is a Wearable Sensor Actually Measuring?

• The Tiny Sensors Doing All the Work

  • Accelerometer

  • Gyroscope

  • Photoplethysmography (PPG)

  • Electrocardiogram (ECG)

  • Temperature Sensor

• How a Physical Signal Becomes Data

• Raw Sensor Data Is Messier Than It Looks

• How Algorithms Turn Messy Signals Into Useful Information

  • Digital Filtering

  • Feature Extraction

  • Metric Calculation

  • Why Wearables Combine Multiple Sensors

• Where Does All This Data Go?

• A Sensor Can Be Accurate and the Final Result Can Still Be Wrong

• Wrap Up

When you check your daily summary, there's a fundamental gap between what the user interface displays and what the underlying hardware actually captured.

Consumer health trackers don't observe abstract concepts like "recovery" or "cardio strain." Instead, they observe physical, mechanical, and optical properties occurring right at the surface of your skin.

The sensor’s sole job is to capture raw physical reality without bias, while software carries the burden of interpretation. Because biological signals are dynamic and influenced by countless environmental variables, converting physical values into physiological insights requires comprehensive mathematical modeling.

A comprehensive review in Nature Digital Medicine on wearable sensing and analytics highlights that wearable health tracking is fundamentally a signal-processing challenge rather than a simple hardware readout.

The Tiny Sensors Doing All the Work

To capture physical signals accurately within a compact form factor, modern wearables rely on a cluster of miniaturized electromechanical and optical modules.

Accelerometer

The accelerometer detects linear acceleration and inertial forces across three spatial axes (X, Y and Z). Built using Micro-Electro-Mechanical Systems (MEMS), it contains microscopic suspended masses that deflect during movement, altering local electrical capacitance.

When you walk, your arm swings in a predictable, periodic pattern. The accelerometer records these repetitive acceleration peaks, allowing software to distinguish rhythmic locomotion from random gestures like typing or drinking water.

Gyroscope

While the accelerometer detects linear movement and gravity, the gyroscope measures angular velocity and rotational motion. It monitors how quickly and along which axis the device rotates in space.

By pairing a gyroscope with an accelerometer, the device can accurately determine its spatial orientation, ensuring that a simple wrist roll to view the screen isn't mistakenly categorized as an exercise rep or a walking stride.

Photoplethysmography (PPG)

PPG sensors use light to monitor changes in microvascular blood volume. Green light-emitting diodes (LEDs) illuminate the capillary bed beneath the skin, while adjacent photodetectors measure the light reflected back.

Because hemoglobin naturally absorbs green light, each ventricular contraction of the heart expands arterial volume, briefly increasing light absorption and lowering the reflected signal. The time between these reflection dips corresponds directly to individual pulse events.

Electrocardiogram (ECG)

While PPG relies on optical reflection, an ECG sensor detects the direct bioelectrical impulses driving the cardiac muscle.

When the heart beats, electrical currents spread across the myocardium, creating subtle voltage fluctuations across your body. By placing a finger on a dedicated case electrode while the back of the watch rests against your wrist, you complete a circuit that lets differential amplifiers measure the heart's depolarization and repolarization waves directly.

Temperature Sensor

Wearable temperature sensors employ thermistors or dedicated resistance temperature detectors (RTDs) resting against the skin surface.

It's worth noting that peripheral skin temperature isn't identical to core body temperature. Skin temperature fluctuates significantly based on ambient air, blood vessel dilation, and peripheral blood circulation. This makes it most valuable for identifying relative baseline deviations, such as sleep-phase cooling or early illness markers, rather than absolute clinical readings.

Comprehensive engineering overviews, such as this IEEE review of wearable physiological sensors, emphasize that these diverse hardware components must operate in close harmony to continuously reconstruct a clear picture of bodily activity.

How a Physical Signal Becomes Data

Before computational logic can make sense of physical phenomena, continuous analog events must be converted into discrete numerical data.

When an optical photodiode detects fluctuating light levels, it outputs a continuous, smooth electrical voltage. Computers, however, can't compute infinite continuous curves. They operate exclusively on discrete numbers. This transition is handled by an Analog-to-Digital Converter (ADC).

The ADC periodically samples the continuous voltage wave and quantizes it into a discrete digital value:

• Sampling Rate: Expressed in Hertz (Hz), this defines how many times per second the ADC records a value.

  Fast, electrically complex signals like ECG require sampling rates of 250 Hz to 500 Hz to capture sharp wave morphology without losing critical cardiac peaks. Conversely, skin temperature changes slowly and can be accurately tracked at a fraction of a single Hertz (such as one sample every few seconds), preserving battery life and system storage.

• Resolution: Typically measured in bits (such as 12-bit, 16-bit or 24-bit depth), resolution determines how finely the converter quantizes the electrical signal. A higher bit-depth allows the system to resolve tiny physiological variations, such as shallow pulse signals on darker skin tones or during cold weather, without the waveform clipping or flattening.

As explored in signal processing literature on wearable biometric data acquisition, selecting appropriate sampling frequencies and quantization ranges balances the need for high signal fidelity with power consumption constraints.

Raw Sensor Data Is Messier Than It Looks

In controlled clinical environments, diagnostic tools are firmly attached to stationary patients. Consumer wearables, by contrast, must gather physiological data during dynamic, everyday movements. Consequently, raw sensor output rarely resembles textbook physiological waveforms.

Everyday wear introduces severe real-world interference:

• Motion Artifacts: When you run, type, or grip objects, muscle contractions and sudden impacts physically rattle the device, creating massive inertial spikes that obscure subtle cardiac pulses.

• Sensor Displacement: A loose strap causes the device chassis to bounce against the epidermis, changing the optical path length between the LEDs and the photodetector, which introduces sharp baseline drifts.

• Environmental & Physiological Noise: Ambient sunlight leaking beneath the device edges can overwhelm sensitive photodiodes, while cold environments trigger peripheral vasoconstriction, drastically reducing blood volume in the wrist capillaries.

Because of this constant interference, wearable firmware includes automated Signal Quality Indices (SQIs). Before handing raw data to downstream algorithms, the system evaluates signal-to-noise ratios (SNR). If a specific data window is completely distorted by motion, the algorithm flags it as unreliable and discards it rather than generating an inaccurate reading. The dynamics of real-time artifact suppression are thoroughly analyzed in wearable artifact removal research.

How Algorithms Turn Messy Signals Into Useful Information

Once the signal is digitized and validated for basic quality, deterministic digital signal processing and algorithmic modeling convert the raw numerical stream into actionable human metrics.

Digital Filtering

Raw data first passes through digital bandpass filters configured to discard frequencies that fall outside the bounds of human physiology.

For an optical heart-rate signal, an algorithm suppresses frequencies below 0.5 Hz (30 BPM) and above 4.0 Hz (240 BPM), filtering out slow baseline drift and high-frequency electrical hum.

Feature Extraction

Instead of continuously processing thousands of raw digital samples, the software extracts concise statistical and morphological markers:

• Peak-to-Peak Intervals (△ t): The precise time duration between consecutive pulse crests.

• Signal Variance: The degree of dispersion in acceleration values across a rolling 5-second window.

• Dominant Frequency: The primary harmonic component identified through Fast Fourier Transforms (FFT).

Metric Calculation

For heart rate, the algorithm identifies valid systolic peaks, measures the inter-beat interval, eliminates mathematical outliers and computes the instantaneous beats per minute (60 / △ t).

For step detection, the algorithm processes 3-axis accelerometer arrays:

The software monitors this composite acceleration value for rhythmic threshold crossings and frequency signatures typical of a human gait, ignoring non-cyclical vibrations like riding a car over a bumpy road.

Machine learning classifiers, trained on large labeled movement datasets, help classify these feature profiles into specific activities such as cycling, swimming or sleeping.

Why Wearables Combine Multiple Sensors

A single physical sensor often lacks the context needed to accurately understand what your body is doing. To resolve ambiguity, devices use Sensor Fusion, combining data from multiple distinct sensors to generate more accurate inferences than any single sensor could provide alone.

Consider a sudden rise in heart rate from 65 BPM to 145 BPM:

• If the accelerometer detects no concurrent body movement, the algorithm may interpret the event as psychological stress, caffeine intake, or a cardiac anomaly.

• If the accelerometer simultaneously registers a sustained, high-cadence rhythmic movement signature, the system identifies the elevated heart rate as a normal physiological response to running.

By pairing optical, thermal, and inertial data points simultaneously, the system constructs a detailed picture of the user's metabolic state.

As detailed in the Biomedical Engineering survey on multimodal sensor fusion, combining complementary sensor streams helps eliminate false positives and balances out individual hardware limitations.

Where Does All This Data Go?

The data pipeline extends beyond the physical device on your wrist. Processing tasks are distributed across local hardware, your paired mobile phone and remote cloud infrastructure.

• On the Wearable (Edge Computing): Time-sensitive tasks run directly on low-power microcontrollers embedded inside the wearable. Filtering raw voltages, detecting steps and monitoring safety alerts (such as fall detection) happen locally, ensuring low latency, lower power consumption and better privacy.

• On the Smartphone: Because smartphones have faster multi-core processors and larger batteries, they handle heavy machine learning tasks, data visualization and the fusion of GPS traces with wrist kinematics.

• In the Cloud: Aggregated summaries are periodically uploaded to remote data centers for long-term historical tracking, deep longitudinal comparisons and training the next generation of algorithmic models across anonymized populations.

A Sensor Can Be Accurate and the Final Result Can Still Be Wrong

A common misconception is that an inaccurate health metric points directly to a broken sensor. In reality, a physical sensor can function with micro-voltage precision while the final displayed metric remains fundamentally incorrect.

There's a distinct difference between direct physical measurement and algorithmic estimation:

• The photodiode may accurately record light absorption.

• The ADC may convert those currents into digital samples without losing precision.

• Yet, if the user experiences severe vascular constriction from cold air or if an unpredictable arm movement mimics a pulse frequency, the peak-detection logic may latch onto the wrong frequency peak.

A metric can fail at multiple points along the pipeline: poor contact mechanics, edge-case physiology that falls outside the training dataset, or mathematical assumptions that break down during specific sports. Recognizing that wearables provide informed physiological estimations rather than direct clinical measurements is key to interpreting everyday health data properly.

Wrap Up

Wearable data goes through much more than a sensor before it becomes the numbers you see on your screen. Sensors capture physical signals, hardware converts them into digital data, and algorithms filter, combine, and interpret those signals to produce useful metrics.

Understanding this process also makes one thing clear: wearable measurements aren't always direct readings. They're often estimates built from several layers of sensing and computation. The better we understand that pipeline, the better we can understand what our wearable data is actually telling us.

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