Back to Classical ML
AI/ML · Classical ML

Double Machine Learning

Regress the outcome on the covariates, regress the treatment on the covariates, and then regress one set of residuals on the other. What is left after removing everything the covariates explain is the part of the treatment that was effectively random, and its relationship to the leftover outcome is the effect. Using flexible learners for both nuisance models would normally bias the result; the residual-on-residual form plus cross-fitting is what cancels that bias out.

SupervisedRegression

Also known as: DML, Debiased ML, Orthogonal learning

Intuition

What it actually does

Regress the outcome on the covariates, regress the treatment on the covariates, and then regress one set of residuals on the other. What is left after removing everything the covariates explain is the part of the treatment that was effectively random, and its relationship to the leftover outcome is the effect. Using flexible learners for both nuisance models would normally bias the result; the residual-on-residual form plus cross-fitting is what cancels that bias out.

Assumptions it makes
  • Conditional ignorability — the measured covariates block every confounding path
  • Overlap, expressed here as non-degenerate residual variation in the treatment after conditioning
  • Nuisance learners converge fast enough (roughly n^-1/4) for the orthogonality argument to hold
  • Units are exchangeable, which is what makes random cross-fitting valid
Objective

Loss function

The final step is an ordinary least-squares slope of one residual on the other. The two nuisance models are fit by whatever loss they like — boosting, random forest, a neural net — because the estimator is constructed so that small errors in them do not propagate into the effect at first order. That property, Neyman orthogonality, is the entire reason this works with machine learning inside it.

outcome nuisance model — E[Y | X]
treatment nuisance model — E[T | X], the propensity
residuals: what the covariates could NOT explain
the estimated treatment effect
Optimization

How it is minimized

Cross-fitted nuisance estimation, then a residual-on-residual moment condition

the units in fold k
residual computed from a model fit WITHOUT fold k — the cross-fitting
number of folds, typically 2 to 5

Two ingredients, and both are load-bearing. Orthogonality means the moment condition's derivative with respect to the nuisance functions is zero at the truth, so a nuisance error of order n^-1/4 still leaves the effect estimate root-n consistent — which is what licenses using a regularized learner at all. Cross-fitting handles the other half: a nuisance model that has seen a unit will fit its noise, and that overfitting correlates the two residuals and biases the slope. Fitting on other folds and predicting out of fold removes it.

Hyperparameters
  • nuisance learners

    One for outcome, one for treatment; they need not match (GBM, random forest, lasso)

  • cross-fitting folds

    More folds means less bias and more compute (2 to 5)

  • repetitions

    Re-run with different fold splits and median the estimates (1 to 20)

  • final-stage form

    Constant effect (partially linear) or a covariate-dependent one (OLS, or a linear projection onto effect-modifiers)

Convergence & cost

The nuisance fits converge as their learners do; the final stage is a closed-form slope. The theoretical guarantee is asymptotic and the practical failure is finite-sample: when the denominator E[T-tilde squared] is small — meaning the covariates already explain nearly all the treatment variation — the estimate becomes wildly unstable, which is the weak-overlap problem showing up as a near-zero denominator instead of an exploded weight. Estimates that swing across fold splits are the diagnostic, which is why repetition and median aggregation is standard.

Two nuisance fits per fold, so 2K model fits, plus O(n) for the final slope. A 5-fold DML with gradient boosting is 10 boosted fits — the most expensive causal estimator here, and still cheap relative to collecting the data.

Applications

How, where, and why it is used

Time-Series Forecasting

Not applicable

Cross-fitting assumes units are exchangeable so that a model fit on other folds is valid for this one. Consecutive observations of a series are not exchangeable, and random folds leak the future into the past. The time-series analogue exists — sequential or blocked cross-fitting — but it is a different estimator with different guarantees, not this one applied to ordered data.

Anomaly Detection

Not applicable

The residuals here are a device for removing confounding, not a signal about individual units. A large residual means the covariates failed to explain that unit, which is ordinary model error, and reading it as an anomaly score would flag exactly the units the estimator is designed to be robust to.

Optimization

Viable
How

Produces the effect coefficient a downstream decision rule optimizes against, with a valid confidence interval attached — which matters when the allocation decision must account for estimate uncertainty rather than treating a point estimate as fact.

Where
  • Pricing elasticity estimation feeding a revenue optimizer
  • Measuring the return on a spend channel before reallocating budget across channels
Why

The distinguishing feature is a valid interval despite machine learning in the nuisance stage. A boosted model can give you an elasticity, but not one you can put a standard error on; DML gives both, and an optimizer that ignores estimate uncertainty will confidently over-allocate to the noisiest channel.

Featurization
  • Include confounders the design requires rather than the features that improve nuisance fit
  • Continuous treatments work directly — no discretization, which is a genuine advantage here
Evaluation

Stability of the estimate across repeated fold splits, plus the width of the interval relative to the decision threshold.

Pitfalls
  • A small residual-treatment variance makes the estimate unstable while the point value still looks plausible
  • Optimizing against the point estimate discards the interval that was the whole reason to use this method
Also applied in

Causal Inference & Uplift

Primary fit
How

Fit both nuisance models with cross-fitting, residualize, regress residual on residual, and report the coefficient with its robust standard error. Repeat over several fold splits and take the median.

Where
  • Observational effect estimation with high-dimensional or nonlinear confounding
  • Continuous-dose treatments, where propensity weighting is awkward and DML is natural
  • Settings where a linear specification is indefensible but a valid interval is still required
Why

The method of choice when you need flexible confounding adjustment AND valid inference. Plugging a boosted model into an outcome regression gives you neither an unbiased effect nor a usable standard error; DML is the construction that recovers both, and it is why the technique moved from econometrics into general practice.

Featurization
  • Pre-treatment covariates only, as always — a post-treatment control is still a collider here
  • Nuisance learners should be flexible; unlike IPTW, accuracy in the propensity model is genuinely wanted
Evaluation

Estimate stability across fold splits, the residual-treatment variance as an overlap diagnostic, and sensitivity analysis for unmeasured confounding.

Pitfalls
  • Skipping cross-fitting reintroduces exactly the overfitting bias the method exists to remove
  • Weak overlap appears as a near-zero denominator rather than as an extreme weight, so it is easy to miss
  • The orthogonality guarantee is asymptotic and can be optimistic at small n

Recommendation & Ranking

Adapted
How

Estimates the incremental effect of a recommendation slot or position, separating the causal lift of showing an item from the selection bias in which items get shown.

Where
  • Position-bias correction, estimating the effect of rank on click independent of relevance
  • Measuring the true incremental value of a recommendation surface against organic discovery
Why

Ranking data is confounded by construction — the model chose what to show, so exposure and relevance are entangled. DML handles the high-dimensional confounding that a simple propensity adjustment cannot, and the residualization is a natural fit for continuous position.

Featurization
  • The ranker score at serve time is the dominant confounder and must be logged
  • Exploration traffic sharpens overlap dramatically; without it the denominator is tiny
Evaluation

Comparison against a randomized-exposure holdout, plus estimate stability across fold splits.

Pitfalls
  • Without exploration traffic there is almost no residual variation in exposure to work with
  • The feedback loop means the estimate is only valid for the ranking policy that produced the logs
Deployment

What it costs to run

Training

Minutes to an hour. 2K nuisance fits — a 5-fold run with boosting is 10 fits, repeated across splits.

Inference

An offline analysis, not a serving path. The output is a coefficient and an interval, produced per study rather than per request.

Retraining

Re-estimated per analysis and after any policy change. An elasticity measured under one pricing regime does not survive a change to that regime.

Drift & monitoring
  • Track the residual-treatment variance across runs — a shrinking denominator means overlap is degrading
  • Compare estimates across fold splits every run; instability is the signal that the asymptotics have not kicked in
  • Watch nuisance model performance separately, since a collapsed treatment model quietly destroys the residuals
Production gotchas
  • Cross-fitting is easy to omit when refactoring, and its absence is invisible in the output
  • Reported standard errors assume the moment condition is correctly specified; misspecification shows up as unstable estimates rather than as a warning
  • Continuous treatments must be on a meaningful scale — the coefficient carries the units, and rescaling silently changes the reported effect
Trade-offs

Pros and cons, in context

Strengths
  • Valid confidence intervals despite ML in the nuisance stage

    The whole reason the method exists. A boosted outcome regression gives an effect with no usable standard error; this gives both, and inference is usually the deliverable in causal work.

  • Robust to moderate misspecification in either nuisance model

    Orthogonality means first-order nuisance errors cancel, so a single bad model does not sink the estimate the way it does in plain IPTW. That robustness is what makes it the safer default.

  • Handles continuous treatments natively

    No discretization and no weighting scheme needed for dose-response questions — a real advantage where propensity methods get awkward.

  • Scales to high-dimensional confounders

    Hundreds of covariates are fine, because the nuisance models are allowed to be regularized. Overkill when there are five well-understood confounders and a linear model suffices.

Limitations
  • Substantially more compute than the alternatives

    2K nuisance fits, often repeated across splits. Irrelevant on a cohort of thousands; a real constraint when the nuisance learner is expensive and the cohort is millions.

  • The guarantee is asymptotic

    At small n the interval can be optimistic and estimates swing across fold splits. Repetition with median aggregation is a mitigation, not a fix.

  • Weak overlap hides in the denominator

    IPTW announces poor overlap with an exploding weight you cannot miss. Here it appears as a small residual variance, which produces an unstable estimate that still looks like a number.

  • Estimates an average, not per-unit effects

    The partially linear form gives one coefficient. When the decision is whom to target rather than whether to act, a meta-learner is the right tool instead.

Implementation

Intuitive, idiomatic, optimized

The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.

1. Intuitive

The math, transcribed literally. No dependencies.

"""Double ML - residualize both sides, then regress one on the other.

Three steps, and the third is an ordinary least-squares slope. What makes it
causal is that both inputs have already had the covariates removed.
"""


def double_ml(fit, X, treatment, outcome, folds=2):
    n = len(X)
    fold_size = n // folds

    y_residual = [0.0] * n
    t_residual = [0.0] * n

    for k in range(folds):
        # Cross-fitting: hold out fold k, fit on the rest, predict on k.
        # Without this the nuisance models fit each unit's own noise, which
        # correlates the two residuals and biases the slope.
        start = k * fold_size
        end = n if k == folds - 1 else start + fold_size
        test_idx = list(range(start, end))
        train_idx = [i for i in range(n) if i not in set(test_idx)]

        train_X = [X[i] for i in train_idx]

        g = fit(train_X, [outcome[i] for i in train_idx])       # E[Y | X]
        m = fit(train_X, [float(treatment[i]) for i in train_idx])  # E[T | X]

        for i in test_idx:
            y_residual[i] = outcome[i] - g(X[i])
            t_residual[i] = treatment[i] - m(X[i])

    # Step 3: slope of y-residual on t-residual. This IS the effect.
    numerator = sum(t * y for t, y in zip(t_residual, y_residual))
    denominator = sum(t * t for t in t_residual)

    if denominator == 0.0:
        raise ValueError("no residual variation in treatment - overlap has failed")

    return numerator / denominator

O(folds x nuisance fit). The train-index construction is O(n^2) through the set difference.

2. Idiomaticscikit-learn

How this language actually wants it written.

"""Double ML - typed, with the overlap diagnostic that decides trust."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray
from sklearn.base import RegressorMixin, clone
from sklearn.model_selection import KFold

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


@dataclass(frozen=True)
class DmlEstimate:
    theta: float
    std_error: float
    residual_treatment_variance: float
    n_folds: int

    @property
    def confidence_interval(self) -> tuple[float, float]:
        margin = 1.96 * self.std_error
        return self.theta - margin, self.theta + margin

    @property
    def has_adequate_overlap(self) -> bool:
        """Weak overlap shows up HERE, as a vanishing denominator - not as an
        extreme weight the way it does in IPTW. Easy to miss, so name it."""
        return self.residual_treatment_variance > 1e-3


def double_ml(
    outcome_learner: RegressorMixin,
    treatment_learner: RegressorMixin,
    X: Matrix,
    treatment: Vector,
    outcome: Vector,
    folds: int = 5,
    random_state: int = 0,
) -> DmlEstimate:
    if X.ndim != 2:
        raise ValueError(f"X must be 2-D, got shape {X.shape}")
    if not (X.shape[0] == treatment.shape[0] == outcome.shape[0]):
        raise ValueError("X, treatment and outcome must describe the same units")
    if folds < 2:
        raise ValueError(f"cross-fitting needs at least 2 folds, got {folds}")

    n = X.shape[0]
    y_residual = np.empty(n, dtype=np.float64)
    t_residual = np.empty(n, dtype=np.float64)

    splitter = KFold(n_splits=folds, shuffle=True, random_state=random_state)
    for train_idx, test_idx in splitter.split(X):
        # clone() per fold: a refitted estimator can carry warm-start state,
        # which silently leaks the held-out fold back into the nuisance model.
        g = clone(outcome_learner).fit(X[train_idx], outcome[train_idx])
        m = clone(treatment_learner).fit(X[train_idx], treatment[train_idx])

        y_residual[test_idx] = outcome[test_idx] - g.predict(X[test_idx])
        t_residual[test_idx] = treatment[test_idx] - m.predict(X[test_idx])

    denominator = float(np.dot(t_residual, t_residual))
    if denominator <= 0.0:
        raise ValueError("no residual variation in treatment - overlap has failed")

    theta = float(np.dot(t_residual, y_residual)) / denominator

    # Robust (sandwich) standard error for the moment condition.
    moment = t_residual * (y_residual - theta * t_residual)
    std_error = float(np.sqrt(np.sum(moment**2)) / denominator)

    return DmlEstimate(
        theta=theta,
        std_error=std_error,
        residual_treatment_variance=denominator / n,
        n_folds=folds,
    )

What changed: The O(n^2) index construction is replaced by KFold, nuisance estimators are cloned per fold so warm-start state cannot leak the held-out data back in, and the bare float becomes an estimate carrying a robust sandwich standard error plus the residual-treatment variance — the diagnostic that reveals weak overlap, which in this estimator hides in the denominator rather than announcing itself.

Explicit type hints on every public signatureDataclasses or NamedTuples over ad-hoc dictsRaise specific exceptions, never bare exceptGuard clauses over nested conditionals

2 x folds nuisance fits plus O(n) reductions.

3. Optimizedscikit-learn / joblib

Faster, and the justification for each change.

"""Double ML - repeated splits in parallel, residuals reduced without copies."""

import numpy as np
from joblib import Parallel, delayed
from numpy.typing import NDArray
from sklearn.base import clone
from sklearn.model_selection import KFold

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


def _one_split(
    outcome_learner, treatment_learner, X: Matrix, t: Vector, y: Vector, folds: int, seed: int
) -> tuple[float, float]:
    """One complete cross-fitted estimate. Returns (theta, denominator)."""
    n = X.shape[0]
    # Pre-allocated once per split; the fold loop writes into slices of these
    # rather than concatenating per-fold arrays at the end.
    y_res = np.empty(n, dtype=np.float64)
    t_res = np.empty(n, dtype=np.float64)

    for train_idx, test_idx in KFold(n_splits=folds, shuffle=True, random_state=seed).split(X):
        g = clone(outcome_learner).fit(X[train_idx], y[train_idx])
        m = clone(treatment_learner).fit(X[train_idx], t[train_idx])
        # Subtract in place - no intermediate prediction array survives.
        np.subtract(y[test_idx], g.predict(X[test_idx]), out=y_res[test_idx])
        np.subtract(t[test_idx], m.predict(X[test_idx]), out=t_res[test_idx])

    denominator = float(t_res @ t_res)
    return float(t_res @ y_res) / denominator, denominator


def double_ml_repeated(
    outcome_learner,
    treatment_learner,
    X: Matrix,
    treatment: Vector,
    outcome: Vector,
    folds: int = 5,
    repetitions: int = 10,
    n_jobs: int = -1,
) -> tuple[float, float]:
    """Median over repeated fold splits, splits fitted concurrently.

    Repetition is not optional polish: a single split's estimate can swing
    materially at finite n, and the median across splits is what makes the
    number reportable. The splits are independent, so they parallelize freely.
    """
    X = np.ascontiguousarray(X, dtype=np.float64)

    results = Parallel(n_jobs=n_jobs, prefer="processes")(
        delayed(_one_split)(
            outcome_learner, treatment_learner, X, treatment, outcome, folds, seed
        )
        for seed in range(repetitions)
    )

    thetas = np.fromiter((theta for theta, _ in results), dtype=np.float64)
    # Median, not mean: a split with weak overlap produces an outlier estimate,
    # and the mean would let it dominate exactly when it should not.
    return float(np.median(thetas)), float(np.std(thetas, ddof=1))

What changed: Repeated fold splits become the default rather than an afterthought, because a single split can swing materially at finite n — and they are independent, so they run concurrently. Residuals are written into pre-allocated buffers with in-place subtraction, and aggregation uses the median so a weak-overlap split produces an outlier rather than dragging the answer.

  • Batch work to amortize interpreter overhead

    Repeated splits are fully independent, so dispatching them together to worker processes converts the dominant cost — 2 x folds x repetitions nuisance fits — into wall-clock divided by cores.

    Trade-off: Process parallelism pickles the design matrix to every worker; at large n the copies exhaust memory before the cores are saturated.

  • Pre-allocate output arrays and use in-place operations

    np.subtract with out= writes residuals directly into pre-sized buffers, so no per-fold prediction array survives the loop.

    Trade-off: Writing into a fancy-indexed slice is subtle — out= on a non-contiguous view silently allocates anyway, so the win depends on the index being a simple slice.

  • Ensure contiguous memory layout and a single dtype

    ascontiguousarray once up front means every fold's fancy indexing gathers from contiguous memory rather than a strided or mixed-dtype view.

    Trade-off: Forces a full copy of the design matrix when the caller passed a view or a non-float64 array.

2 x folds x repetitions nuisance fits, concurrent. Illustrative, not a measured benchmark.