Causal Inference for Machine Learning

Predicting an outcome from observed features and estimating the effect of changing a treatment require different evidence. A useful predictor can exploit associations that would disappear under intervention. This article develops that distinction through causal diagrams and estimators for an average treatment effect.

Define the intervention before estimating its effect

Let \(T=1\) mean receiving a specified treatment and \(T=0\) mean its specified alternative. \(Y(1)\) and \(Y(0)\) are the outcomes a person would have under these alternatives. The average treatment effect (ATE) in a target population is \(\tau=\mathbb E[Y(1)-Y(0)]\). We observe only one outcome per person, so an individual effect cannot be read directly from that person’s record. Specify the population, treatment versions, outcome, and follow-up time before choosing an estimator.

For example, potential outcomes of 8 and 5 would give an individual effect of 3. An experiment observes different people under the two alternatives; random assignment makes their mean outcomes comparable in expectation. Finite samples still fluctuate, and nonadherence or missing outcomes can complicate the analysis. The effect of assignment is not automatically the effect of treatment received.

With observational data, this article uses recorded pretreatment covariates \(C\) and three identification assumptions. Consistency connects the observed outcome to the potential outcome under the treatment actually received; here treatments are well defined and one person’s treatment does not affect another’s outcome. Conditional exchangeability, \((Y(1),Y(0))\perp T\mid C\), means the groups are comparable in their potential outcomes within covariate strata. Positivity requires \(0<P(T=1\mid C=c)<1\) for relevant strata. Recording many variables does not establish exchangeability.

Under these assumptions, \(\tau=\mathbb E_C[\mathbb E(Y\mid T=1,C)-\mathbb E(Y\mid T=0,C)]\). Compare treated and untreated outcomes within each covariate stratum, then average those differences over the target population. If two equally common strata have differences 1 and 3, this average is 2, regardless of their different treatment frequencies. This step is identification: expressing the causal target using observed quantities. Fitting models and averaging their predictions is the subsequent estimation step.

Choosing an adjustment set

A causal diagram is a directed acyclic graph: an arrow encodes an assumed direct causal influence, and no sequence of arrows returns to its starting node. In the examples below, \(C\to T\) and \(C\to Y\) form a confounding path \(T\leftarrow C\to Y\). Conditioning on \(C\) blocks that path. In \(T\to K\leftarrow U\to Y\), the arrows meet at the collider \(K\); conditioning on it opens an otherwise blocked path. Finally, \(T\to M\to Y\) is a causal path through a mediator. A backdoor adjustment set blocks noncausal paths entering treatment while preserving the causal paths needed for the target effect. The full graph matters; these three roles are not a universal feature-selection checklist.

The first block uses continuous treatments and linear structural equations with independent noise. Its target is the total effect of increasing treatment by one unit, which is 2 in each setup. The helper returns the fitted coefficient on the first input column, treatment.

import numpy as np
from sklearn.linear_model import LinearRegression

def effect(X, y):
    return LinearRegression().fit(X, y).coef_[0]

g = np.random.default_rng(0); n = 200_000
TRUE = 2.0

# 1. CONFOUNDER: C causes both the treatment and the outcome
C = g.normal(size=n)
T = 1.5 * C + g.normal(size=n)
Y = TRUE * T + 3.0 * C + g.normal(size=n)
print(f"confounder   T alone {effect(T.reshape(-1, 1), Y):+.4f}   T and C {effect(np.c_[T, C], Y):+.4f}")

T2, U = g.normal(size=n), g.normal(size=n)
Y2 = TRUE * T2 + 2.0 * U + g.normal(size=n)
K = 1.2 * T2 + 1.2 * U + g.normal(size=n)
print(f"collider     T alone {effect(T2.reshape(-1, 1), Y2):+.4f}   T and K {effect(np.c_[T2, K], Y2):+.4f}")

T3 = g.normal(size=n)
M = 1.0 * T3 + g.normal(size=n)
Y3 = 1.2 * T3 + 0.8 * M + g.normal(size=n)          # total effect 1.2 + 0.8 = 2.0
print(f"mediator     T alone {effect(T3.reshape(-1, 1), Y3):+.4f}   T and M {effect(np.c_[T3, M], Y3):+.4f}")
print(f"the true causal effect is {TRUE:+.4f} in all three")
# confounder   T alone +3.3832   T and C +1.9981
# collider     T alone +2.0102   T and K +0.8227
# mediator     T alone +1.9992   T and M +1.2008
# the true causal effect is +2.0000 in all three

In the collider example, including \(K\) moves the estimate from +2.0102 to +0.8227. Holding \(K\) fixed associates higher \(T\) with lower \(U\), whose effect on \(Y\) then contaminates the treatment coefficient.

In the mediator example, a one-unit treatment increase raises \(M\) by one unit. The total effect is therefore \(1.2+0.8\times1=2\). Holding \(M\) fixed leaves the direct effect 1.2; the fitted coefficient +1.2008 is close to that value. This interpretation relies on the stated additive equations and independent errors. Simply adding a mediator to a regression does not generally identify a direct effect when mediator–outcome confounding or treatment interactions are present.

An adjustment set should be justified using causal assumptions, timing, and domain knowledge. Statistical evidence can help assess proposed graphs, and causal discovery uses additional assumptions to narrow possibilities; associations alone generally do not determine causal direction. Including every available variable can open a collider path or change the target by blocking mediation. Conversely, predictive models can be useful inside a causal estimator after a valid adjustment set is chosen.

Estimators, once the variable set is settled

For binary treatment, write \(e(C)=P(T=1\mid C)\) for the propensity score and \(m_t(C)=\mathbb E[Y\mid T=t,C]\) for an outcome regression. Outcome modeling averages \(\hat m_1(C_i)-\hat m_0(C_i)\). Inverse propensity weighting (IPW) gives treated observations weight \(1/\hat e_i\) and untreated ones \(1/(1-\hat e_i)\); the code normalizes each group’s weighted mean separately (Hájek IPW). A treated row with propensity 0.25 receives weight 4, representing treatment opportunities that occur relatively rarely in that stratum.

The augmented inverse propensity weighted estimate is \(\hat\tau=\frac1n\sum_i[\hat m_1(C_i)-\hat m_0(C_i)+T_i(Y_i-\hat m_1(C_i))/\hat e_i-(1-T_i)(Y_i-\hat m_0(C_i))/(1-\hat e_i)]\). The first term predicts the effect for every row; the remaining terms correct it using observed residuals. Under identification and suitable regularity conditions, consistent estimation of either the propensity or both outcome regressions can make AIPW consistent. This is an asymptotic property, not a promise of an exact finite-sample answer.

The simulation below supplies correct nonlinear feature transformations to one set of models and omits them from another. Three-fold cross-fitting predicts each row using models fitted on other rows. It reduces reuse of observations for fitting and residual correction; it does not fix unmeasured confounding or poor overlap. All four rows use the same folds.

from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold

g = np.random.default_rng(1); n = 120_000
Cv = g.normal(size=(n, 3))
ps_true = 1 / (1 + np.exp(-0.45 * (1.1*Cv[:, 0]**2 - 0.9*Cv[:, 1]**2 + 0.6*Cv[:, 2] - 1.0)))
T = (g.random(n) < ps_true).astype(int)
Y = 2.0*T + 2.5*Cv[:, 0]**2 + 1.8*np.abs(Cv[:, 1]) + Cv[:, 2] + g.normal(size=n)
Xp = np.c_[Cv[:, 0]**2, Cv[:, 1]**2, Cv[:, 2]]
Xm = np.c_[Cv[:, 0]**2, np.abs(Cv[:, 1]), Cv[:, 2]]
folds = list(StratifiedKFold(3, shuffle=True, random_state=0).split(Cv, T))
print(f"naive difference {Y[T == 1].mean() - Y[T == 0].mean():.4f}")
print("propensity outcome IPW outcome-only AIPW clipped-fraction")
for pk, P in (("specified", Xp), ("linear", Cv)):
    raw_ps = np.empty(n)
    for tr, te in folds:
        raw_ps[te] = LogisticRegression(C=1e6, max_iter=2000).fit(P[tr], T[tr]).predict_proba(P[te])[:, 1]
    ps = np.clip(raw_ps, 1e-6, 1 - 1e-6)
    for ok, Q in (("specified", Xm), ("linear", Cv)):
        mu = np.empty((n, 2))
        for tr, te in folds:
            for t in (0, 1):
                rows = tr[T[tr] == t]
                mu[te, t] = LinearRegression().fit(Q[rows], Y[rows]).predict(Q[te])
        mu0, mu1 = mu.T
        w1, w0 = T/ps, (1-T)/(1-ps)
        ipw = np.sum(w1*Y)/w1.sum() - np.sum(w0*Y)/w0.sum()
        score = mu1-mu0 + w1*(Y-mu1) - w0*(Y-mu0)
        print(f"{pk} {ok} {ipw:.4f} {(mu1-mu0).mean():.4f} {score.mean():.4f} {np.mean(ps != raw_ps):.6f}")
    print(pk, "propensity quantiles", np.round(np.quantile(raw_ps, [0, .01, .5, .99, 1]), 4))
# naive difference 3.6871
# propensity outcome IPW outcome-only AIPW clipped-fraction
# specified specified 2.0308 2.0024 2.0055 0.000000
# specified linear 2.0308 3.4778 2.0325 0.000000
# specified propensity quantiles [3.000e-04 5.110e-02 4.008e-01 9.262e-01 1.000e+00]
# linear specified 3.4775 2.0024 2.0026 0.000000
# linear linear 3.4775 3.4778 3.4779 0.000000
# linear propensity quantiles [0.1882 0.2905 0.4177 0.5564 0.6767]

The specified propensity model includes the squared covariates in the generating logit; the specified outcome models include the squared and absolute-value terms in the conditional mean. Thus “specified” has a concrete meaning here. Compare each estimate with the known effect 2. With the specified propensity and linear outcome model, the outcome-only estimate is 3.4778 and AIPW is 2.0325. With the linear propensity and specified outcome models, IPW is 3.4775 and AIPW is 2.0026. Both specified components give AIPW 2.0055; both linear ones give 3.4779. A single table cannot establish consistency, and finite-sample error remains even in correctly specified models.

When both components omit the nonlinear terms, double robustness supplies no guarantee. It does not require failure in every such problem, either: misspecification can sometimes cancel. Prediction accuracy or agreement between estimators cannot establish that their causal assumptions hold.

Positivity excludes probabilities exactly 0 or 1 in relevant strata; probabilities below 0.05 or above 0.95 are not automatic violations. Near-extreme values can produce large, unstable weights. Inspect overlap by treatment group, weight concentrations, and covariate balance after weighting. Here a numerical guard clips at \(10^{-6}\); the printed fraction is zero in this run. The largest propensity rounds to 1.0000 in the display but remains below 1; the displayed precision does not determine overlap. A fixed clipping threshold can introduce asymptotic bias when it changes a correct propensity model. Trimming rows can also change the target population; neither operation creates missing comparisons.

These examples focus on adjustment for an average effect. ATE can hide different effects across covariate groups, and the effect among treated people (ATT) averages over a different population. If exchangeability given recorded covariates is implausible, alternative designs need their own assumptions: an instrumental variable needs relevance, an exclusion restriction, and suitable independence (with monotonicity for a common local-effect interpretation); regression discontinuity uses continuity around a treatment threshold and concerns a local effect; difference-in-differences relies on an appropriate parallel-trends assumption. None is an automatic substitute for randomization.

Report the target effect, why the design identifies it, the adjustment set, overlap diagnostics, and sampling uncertainty. For AIPW, a standard error based on the variability of the cross-fitted scores requires additional rate and regularity conditions; cross-fitting alone does not justify a confidence interval. Such intervals quantify sampling uncertainty under the model and design assumptions, not the possible bias from an omitted confounder.

Exercises

1. An effect that reverses sign when you pool. Build a treatment that helps within every group and appears harmful overall.

You should get: a positive effect in each group and a strongly negative one in the pooled data.

import numpy as np
from sklearn.linear_model import LinearRegression

g = np.random.default_rng(0)
rows = []
# the treatment is given mostly in the group with the worse baseline
for dept, base, n_d, p_treat in (("A (easy)", 0.80, 8000, 0.15),
                                 ("B (hard)", 0.25, 8000, 0.85)):
    T = (g.random(n_d) < p_treat).astype(int)
    Y = (g.random(n_d) < np.clip(base + 0.05 * T, 0, 1)).astype(int)   # true effect +0.05
    rows.append((dept, T, Y))

Tall = np.concatenate([r[1] for r in rows]); Yall = np.concatenate([r[2] for r in rows])
D = np.concatenate([np.full(len(r[1]), i) for i, r in enumerate(rows)])
for dept, T, Y in rows:
    print(f"{dept:10s} treated {Y[T == 1].mean():.4f}  untreated {Y[T == 0].mean():.4f}"
          f"  diff {Y[T == 1].mean() - Y[T == 0].mean():+.4f}")
print(f"{'POOLED':10s} treated {Yall[Tall == 1].mean():.4f}  untreated {Yall[Tall == 0].mean():.4f}"
      f"  diff {Yall[Tall == 1].mean() - Yall[Tall == 0].mean():+.4f}")
print(f"adjusted for department: {LinearRegression().fit(np.c_[Tall, D], Yall).coef_[0]:+.4f}")
# A (easy)   treated 0.8585  untreated 0.7969  diff +0.0616
# B (hard)   treated 0.3063  untreated 0.2465  diff +0.0597
# POOLED     treated 0.3912  untreated 0.7165  diff -0.3253
# adjusted for department: +0.0607

2. What the estimators cannot fix. Vary an unmeasured confounder and compare adjustment using recorded variables with an oracle that also observes the confounder.

You should get: similar biased estimates when the unmeasured confounder affects treatment and outcome.

import numpy as np
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.ensemble import HistGradientBoostingRegressor

g = np.random.default_rng(0); n = 120_000
print(f"{'unmeasured':>11} {'naive':>9} {'adjust for C':>13} {'AIPW on C':>11} {'adjust C and U':>15}")
for gamma in (0.0, 0.5, 1.0, 2.0, 3.0):
    C = g.normal(size=n)
    U = g.normal(size=n)                                   # never recorded
    T = (g.random(n) < 1 / (1 + np.exp(-(0.8 * C + gamma * U)))).astype(int)
    Y = 2.0 * T + 1.5 * C + gamma * U + g.normal(size=n)
    Cc = C.reshape(-1, 1)
    ps = np.clip(LogisticRegression(max_iter=2000).fit(Cc, T).predict_proba(Cc)[:, 1], .02, .98)
    m1 = HistGradientBoostingRegressor(random_state=0).fit(Cc[T == 1], Y[T == 1]).predict(Cc)
    m0 = HistGradientBoostingRegressor(random_state=0).fit(Cc[T == 0], Y[T == 0]).predict(Cc)
    aipw = np.mean(m1 - m0 + T*(Y - m1)/ps - (1 - T)*(Y - m0)/(1 - ps))
    print(f"{gamma:11.1f} {Y[T == 1].mean() - Y[T == 0].mean():9.4f}"
          f" {LinearRegression().fit(np.c_[T, C], Y).coef_[0]:13.4f} {aipw:11.4f}"
          f" {LinearRegression().fit(np.c_[T, C, U], Y).coef_[0]:15.4f}")
print("(the true effect is 2.0000 in every row)")
#  unmeasured     naive  adjust for C   AIPW on C  adjust C and U
#         0.0    3.0394        1.9908      1.9898          1.9908
#         0.5    3.2271        2.2281      2.2301          1.9960
#         1.0    3.6530        2.8281      2.8305          2.0018
#         2.0    4.9934        4.4249      4.4286          1.9861
#         3.0    6.5920        6.1881      6.1915          2.0094

3. A predictor and a cause are not the same thing. Compare test-set R² against causal correctness for two model specifications.

You should get: the model that predicts better being the one whose coefficient is wrong.

import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import r2_score

g = np.random.default_rng(0); n = 60_000
T = g.normal(size=n)
U = g.normal(size=n)
Y = 2.0 * T + 2.0 * U + g.normal(size=n)
K = 1.5 * T + 1.5 * U + g.normal(size=n)          # a COLLIDER: caused by T and U

X_small = T.reshape(-1, 1)
X_big = np.c_[T, K]
for name, X in (("Y ~ T", X_small), ("Y ~ T + K", X_big)):
    Xtr, Xte, ytr, yte = train_test_split(X, Y, test_size=0.4, random_state=0)
    m = LinearRegression().fit(Xtr, ytr)
    print(f"{name:12s} test R2 {r2_score(yte, m.predict(Xte)):.4f}"
          f"   coefficient on T {m.coef_[0]:+.4f}   (true causal effect +2.0000)")
# Y ~ T        test R2 0.4463   coefficient on T +1.9911   (true causal effect +2.0000)
# Y ~ T + K    test R2 0.7523   coefficient on T +0.6044   (true causal effect +2.0000)

References

Hernán and Robins, Causal Inference: What If; Funk et al., Doubly Robust Estimation of Causal Effects; Zivich and Breskin, Machine learning for causal inference: on the use of cross-fit estimators.