8  Nonparametric Causal Methods

using DataFrames
using Distributions
using Random
using Statistics
using MLJ
using EvoTrees
using MLJLinearModels
using DecisionTree
using CausalEstimate
import CausalEstimate: estimate, confint, pvalue
using Printf

Related reading: For an R-based TMLE simulation that compares super-learner libraries and stress-tests model misspecification, see the companion blog chapter on TMLE in Topics on Econometrics and Causal Inference.

8.1 Why Do We Need Nonparametric Methods?

• Parametric models are often mis-specified
• Nonparametric models are more flexible; less assumptions
• What kind of nonparametrics are good?

Considerations:

• Plug-in (g-computation) estimators are easy.
• The problem is not the decomposition into \(E[Y(0)]\) and \(E[Y(1)]\) itself: with correctly specified models and regularity, a plug-in estimator can be efficient. The problem appears once the nuisance models are flexible and regularised. Then the error in the nuisance estimate enters the plug-in target linearly, and the plug-in has no term that cancels it. An influence-function estimator adds exactly such a term, and the error that survives is a product of two nuisance errors rather than one. Section 8.2.1 does the algebra.
• So they can be consistent, but the convergence rate inherits the (often slow) nuisance rate, which is especially bad in high dimensions.
• Ideally we remove that first-order sensitivity using an orthogonal influence function, and get \(\sqrt n\) efficiency.
• Influence-function-based (one-step/AIPW/TMLE) estimators are shown to be efficient and orthogonal to first-order nuisance error.

8.2 Influence Function Based Estimators

Estimators based on efficient influence function are often \(\sqrt n\) consistent, meaning they converge to the truth as fast as \(\sqrt n\) rate. This means the bias not only goes to 0 with \(n\) goes to infinity, but also as fast as \(1/ \sqrt n\) goes to 0. Why does this matter? Because we have limited sample size, when we are based on asymptotics, the faster the bias goes to 0 the smaller sample size we need.

If we have an estimator such that \[ \sqrt n (\hat \psi - \psi) = \sqrt n P_n(\phi) + o_P(1) \rightarrow N(0, var(\phi)) \tag{8.1}\]

Then this estimator is \(\sqrt n\) consistent and asymptotically normal; and \(\phi\) is the influence function. The efficient influence function is the one with variance at the lower bound (similar to parametric case for Cramer-Rao lower bound).

Unfortunately there is no universal formula for efficient influence function for every problem. We need to derive it for each problem. For example, ATE has an IF based estimator, but ATT needs a different formula; and LATE has a different formula too, etc. Deriving them is hard!

For the treated potential-outcome mean, \(\psi_1 = E\{E(Y \mid X, A=1)\}\), the efficient influence function is \[\phi_1 (Z; P) = \frac{A}{\pi (X)} \{ Y - \mu_1 (X) \} + \mu_1(X) - \psi_1 \tag{8.2}\] where \(\pi (X) = P(A=1 | X)\) is the propensity score model and \(\mu_1 (X) = E(Y |X, A=1)\) is the treated outcome model. The ATE influence function is the difference between the treated and untreated components, \(\phi_1-\phi_0\).

For example, an IF-based bias-corrected estimator for the treated-arm mean \[\hat \psi_1 = P_n \left[ \frac{A}{\hat \pi (X)} \{ Y - \hat \mu_1 (X) \} + \hat \mu_1(X) \right] \tag{8.3}\] where \(P_n\) means sample mean. This is the familiar AIPW estimator. It is known to be doubly robust.

Since we know the influence function, we can calculate its variance. This variance is the same as variance of \(\hat \psi\) since they differ by a constant \(\psi\). We can use any estimators to estimate the nuisance parameters, then get \(\hat \psi\). Then the sample variance is the variance of the influence function; and the sample mean is the ATE.

In this estimator, both \(\hat \mu\) and \(\hat \pi\) can be done using any estimator, such as machine learning. One caveat: with adaptive learners (random forests, boosting), the nuisance fits should be cross-fit — predicted for each observation from a model trained without it — otherwise own-observation overfitting can bias the estimate. The manual AIPW below reuses full-sample fits to keep the code short; the CausalEstimate.jl calls use crossfit=5.

8.2.1 Why the Plug-In Is Off

The bullets above claimed that the plug-in inherits the nuisance error at first order and that the augmentation term removes it. That is algebra, not a simulation finding, so it is worth doing before we run anything.

Take \(\psi_1 = E[\mu_1(X)]\) with \(\mu_1(x) = E(Y \mid X = x, A = 1)\). The plug-in estimator is \(\hat\psi_1^{\,\text{pi}} = P_n[\hat\mu_1(X)]\). Split its error into a sampling part and a nuisance part:

\[ \hat\psi_1^{\,\text{pi}} - \psi_1 = \underbrace{(P_n - P)[\mu_1]}_{O_P(n^{-1/2})} \; + \; \underbrace{P[\hat\mu_1 - \mu_1]}_{\text{nuisance error, linear}} \; + \; \underbrace{(P_n - P)[\hat\mu_1 - \mu_1]}_{\text{smaller still}} . \tag{8.4}\]

The middle term is \(E[\hat\mu_1(X) - \mu_1(X)]\), and it is linear in the nuisance error. If \(\hat\mu_1\) converges at rate \(n^{-\alpha}\), that term is of order \(n^{-\alpha}\), and a regularised learner in even a moderate number of dimensions gives \(\alpha < 1/2\). The nuisance error then dominates the sampling error, the estimator is not \(\sqrt n\) consistent, and nothing in the plug-in cancels it.

Now do the same for the AIPW estimator of Equation 8.3. Evaluate it under the population distribution, holding the estimated nuisances fixed. Because \(E[A\{Y - \hat\mu_1(X)\} \mid X] = \pi(X)\{\mu_1(X) - \hat\mu_1(X)\}\),

\[ P\Bigl[\frac{A}{\hat\pi(X)}\bigl\{Y - \hat\mu_1(X)\bigr\} + \hat\mu_1(X)\Bigr] = \psi_1 + E\Bigl[\frac{\pi(X) - \hat\pi(X)}{\hat\pi(X)}\, \bigl\{\mu_1(X) - \hat\mu_1(X)\bigr\}\Bigr]. \tag{8.5}\]

The bias is a product of the two nuisance errors, and that is the entire difference between the two estimators. The term linear in \(\hat\mu_1 - \mu_1\) has cancelled: the \(\hat\mu_1\) that the outcome model carries in is subtracted back out by the \(\hat\mu_1\) sitting inside the augmentation, and what survives is multiplied by the propensity-score error. By Cauchy-Schwarz the bias is bounded by \(\lVert \pi - \hat\pi \rVert \, \lVert \mu_1 - \hat\mu_1 \rVert\) up to the overlap constant, so rates of \(o(n^{-1/4})\) on each nuisance suffice to make it \(o(n^{-1/2})\) — slower than \(\sqrt n\) on each piece, fast enough for \(\sqrt n\) on the target.

Two familiar facts fall out of Equation 8.5 directly. Double robustness is immediate: if either factor is zero the whole bias is zero, which is why one correct model is enough. And the \(n^{-1/4}\) requirement is the reason cross-fitting matters, since reusing the same observations to fit and to evaluate \(\hat\mu_1\) makes the product term behave worse than the rate calculation assumes.

8.2.2 Simulation Setup

Random.seed!(1)
n = 1000
w1 = rand(Binomial(1, 0.5), n)
w2 = rand(Binomial(1, 0.5), n)
w3 = round.(rand(Uniform(0, 4), n), digits=3)
w4 = round.(rand(Uniform(0, 5), n), digits=3)

logistic(x) = 1.0 ./ (1.0 .+ exp.(-x))   # inverse logit (sigmoid), not the logit

prob_A = logistic.(-0.4 .+ 0.2.*w2 .+ 0.15.*w3 .+ 0.2.*w4 .+ 0.15.*w2.*w4)
A = rand.(Binomial.(1, prob_A))

prob_Y1 = logistic.(-1.0 .+ 1.0 .- 0.1.*w1 .+ 0.3.*w2 .+ 0.25.*w3 .+ 0.2.*w4 .+ 0.15.*w2.*w4)
prob_Y0 = logistic.(-1.0 .+ 0.0 .- 0.1.*w1 .+ 0.3.*w2 .+ 0.25.*w3 .+ 0.2.*w4 .+ 0.15.*w2.*w4)

Y_1 = rand.(Binomial.(1, prob_Y1))
Y_0 = rand.(Binomial.(1, prob_Y0))
Y = Y_1 .* A .+ Y_0 .* (1 .- A)

W = DataFrame(w1=w1, w2=w2, w3=w3, w4=w4)
df = DataFrame(w1=w1, w2=w2, w3=w3, w4=w4, A=A, Y=Y, Y_1=Y_1, Y_0=Y_0)

True_Psi = mean(df.Y_1 .- df.Y_0)
@printf "True ATE = %.4g\n" True_Psi

True ATE = 0.184

The line above reports mean(Y.1 - Y.0) for this particular draw. That is a legitimate quantity — the sample average treatment effect — but it is not a fixed benchmark, and it is easy to mistake for one. The potential outcomes are binary, so at \(n = 1000\) this realized average has a sampling standard deviation of about 0.020: two runs of the same DGP under different seeds can differ by 0.04 without anything being wrong. (The R and Julia editions of this chapter print 0.21 and 0.184 for exactly this reason.) Estimator standard errors in this chapter are around 0.03, so a noisy target of that size would swamp the comparison we are trying to make.

The population ATE is a definite number, and since the DGP is fully specified we can integrate it out rather than simulate it:

# w1, w2 ~ Bernoulli(0.5); w3 ~ U(0,4); w4 ~ U(0,5). Average the true individual
# effect over that covariate distribution on a trapezoid grid.
function population_ate()
    g3 = range(0, 4; length = 1601)
    g4 = range(0, 5; length = 1601)
    trap(g) = (w = fill(step(g), length(g)); w[1] /= 2; w[end] /= 2; w)
    w3w, w4w = trap(g3), trap(g4)
    total = 0.0
    for w1 in (0, 1), w2 in (0, 1)
        acc = 0.0
        for (j, u4) in enumerate(g4), (i, u3) in enumerate(g3)
            lin0 = -1 - 0.1w1 + 0.3w2 + 0.25u3 + 0.2u4 + 0.15w2 * u4
            acc += w3w[i] * w4w[j] * (logistic(lin0 + 1) - logistic(lin0))
        end
        total += acc / (4 * 5)
    end
    total / 4
end

pop_ate = population_ate()
@printf "Realized-sample ATE (this draw): %.4f\n" True_Psi
@printf "Exact population ATE:           %.4f\n" pop_ate

Realized-sample ATE (this draw): 0.1840
Exact population ATE:           0.2028

Compare the estimators below against 0.2028, not against the realized draw.

# Load MLJ models to form our equivalent of SuperLearner
LogisticClassifier = @load LogisticClassifier pkg=MLJLinearModels verbosity=0
EvoTreeClassifier = @load EvoTreeClassifier pkg=EvoTrees verbosity=0
RandomForestClassifier = @load RandomForestClassifier pkg=DecisionTree verbosity=0

# Define a Stack (SuperLearner)
stack = Stack(
    metalearner = LogisticClassifier(),
    resampling = CV(nfolds=5),
    evo = EvoTreeClassifier(max_depth=4, nrounds=100),
    rf = RandomForestClassifier(n_trees=100, max_depth=5),
    glm = LogisticClassifier()
)

ProbabilisticStack(
  metalearner = LogisticClassifier(
        lambda = 2.220446049250313e-16, 
        gamma = 0.0, 
        penalty = :l2, 
        fit_intercept = true, 
        penalize_intercept = false, 
        scale_penalty_with_samples = true, 
        solver = nothing), 
  resampling = CV(
        nfolds = 5, 
        shuffle = false, 
        rng = TaskLocalRNG()), 
  measures = nothing, 
  cache = true, 
  acceleration = CPU1{Nothing}(nothing), 
  evo = EvoTreeClassifier(
        loss = :mlogloss, 
        metric = :mlogloss, 
        nrounds = 100, 
        bagging_size = 1, 
        early_stopping_rounds = 9223372036854775807, 
        L2 = 1.0, 
        lambda = 0.0, 
        gamma = 0.0, 
        eta = 0.1, 
        max_depth = 4, 
        min_weight = 1.0, 
        rowsample = 1.0, 
        colsample = 1.0, 
        nbins = 64, 
        tree_type = :binary, 
        seed = 123, 
        device = :cpu), 
  rf = RandomForestClassifier(
        max_depth = 5, 
        min_samples_leaf = 1, 
        min_samples_split = 2, 
        min_purity_increase = 0.0, 
        n_subfeatures = -1, 
        n_trees = 100, 
        sampling_fraction = 0.7, 
        feature_importance = :impurity, 
        rng = TaskLocalRNG()), 
  glm = LogisticClassifier(
        lambda = 2.220446049250313e-16, 
        gamma = 0.0, 
        penalty = :l2, 
        fit_intercept = true, 
        penalize_intercept = false, 
        scale_penalty_with_samples = true, 
        solver = nothing))

8.2.3 Plug-in Estimator from Outcome Regression

Random.seed!(1)
X_mu = select(df, :w1, :w2, :w3, :w4, :A)
mach_mu = machine(stack, X_mu, categorical(df.Y))
MLJ.fit!(mach_mu, verbosity=0)

# predict the PO Y^1
X_mu1 = copy(X_mu); X_mu1.A .= 1
X_mu0 = copy(X_mu); X_mu0.A .= 0

# Predict probability of Y=1 (assuming 1 is the positive class)
mu1hat = pdf.(MLJ.predict(mach_mu, X_mu1), 1)
mu0hat = pdf.(MLJ.predict(mach_mu, X_mu0), 1)

plug_in_ate = mean(mu1hat .- mu0hat)
@printf "ATE (plug-in) = %.4g\n" plug_in_ate

ATE (plug-in) = 0.1759

8.2.4 AIPW ATE

Random.seed!(123)
X_pi = select(df, :w1, :w2, :w3, :w4)
mach_pi = machine(stack, X_pi, categorical(df.A))
MLJ.fit!(mach_pi, verbosity=0)

pi1hat = pdf.(MLJ.predict(mach_pi, X_pi), 1)

# Bound propensity scores to avoid extreme weights
pi1hat = max.(min.(pi1hat, 0.999), 0.001)
pi0hat = 1.0 .- pi1hat

# Current prediction on observed data
mu1hat_obs = mu1hat
mu0hat_obs = mu0hat

IF_1 = ((df.A ./ pi1hat) .* (df.Y .- mu1hat_obs) .+ mu1hat_obs)
IF_0 = (((1 .- df.A) ./ pi0hat) .* (df.Y .- mu0hat_obs) .+ mu0hat_obs)
IF = IF_1 .- IF_0

aipw_1 = mean(IF_1); aipw_0 = mean(IF_0)
aipw_manual = aipw_1 - aipw_0

ci_lb = mean(IF) - quantile(Normal(), 0.975) * std(IF) / sqrt(length(IF))
ci_ub = mean(IF) + quantile(Normal(), 0.975) * std(IF) / sqrt(length(IF))
res_manual_aipw = [aipw_manual, ci_lb, ci_ub]
@printf "ATE (AIPW)   = %.4g\n" aipw_manual
@printf "95%% CI       = [%.4g, %.4g]\n" ci_lb ci_ub

ATE (AIPW)   = 0.1918
95% CI       = [0.1311, 0.2525]

8.2.5 AIPW ATE Using CausalEstimate.jl

using CausalEstimate

r(x) = round(x, sigdigits=4)

psi_ate = ATE(outcome=:Y, treatment=:A, confounders=[:w1, :w2, :w3, :w4])
aipw_result = estimate(psi_ate, AIPW(crossfit=5), df)

(
    ATE    = r(estimate(aipw_result)),
    CI     = r.(confint(aipw_result)),
    pvalue = r(pvalue(aipw_result)),
)

(ATE = 0.2072, CI = (0.1206, 0.2939), pvalue = 2.776e-6)

8.2.6 Regression-imputation ATT (manual)

Note that the manual estimator below is evaluated on the treated subset df1, where (1 .- df1.A) is identically zero. The augmentation term ((1-A)/π₀)(Y-μ₀) therefore vanishes and the estimator collapses to the regression-imputation ATT, mean(Y_treated − μ̂₀(X_treated)). That is consistent if the control outcome model is correct, but it is not the doubly-robust / augmented ATT, and the influence-function SE understates uncertainty because it ignores estimation of μ₀ and π. The proper ATT efficient influence function (Hahn 1998) is (1/p)[A(Y−μ₀) − (1−A)·π/(1−π)·(Y−μ₀)] evaluated over the full sample; the CausalEstimate.jl call in the next subsection computes the cross-fit, doubly-robust version.

Random.seed!(1)
df2 = copy(df)
df2.pi0hat = pi0hat
df2.pi1hat = pi1hat

df0 = filter(row -> row.A == 0, df2)
df1 = filter(row -> row.A == 1, df2)

# Fit outcome model only on controls
X_mu0 = select(df0, :w1, :w2, :w3, :w4)
mach_mu0 = machine(stack, X_mu0, categorical(df0.Y))
MLJ.fit!(mach_mu0, verbosity=0)

# Predict counterfactual for treated
X_mu1 = select(df1, :w1, :w2, :w3, :w4)
mu0hat_treated = pdf.(MLJ.predict(mach_mu0, X_mu1), 1)

IF_0_att = (((1 .- df1.A) ./ df1.pi0hat) .* (df1.Y .- mu0hat_treated) .+ mu0hat_treated)
IF_att = df1.Y .- IF_0_att

aipw_manual_att = mean(IF_att)
ci_lb_att = mean(IF_att) - quantile(Normal(), 0.975) * std(IF_att) / sqrt(length(IF_att))
ci_ub_att = mean(IF_att) + quantile(Normal(), 0.975) * std(IF_att) / sqrt(length(IF_att))
res_manual_aipw_att = [aipw_manual_att, ci_lb_att, ci_ub_att]
res_manual_aipw_att

3-element Vector{Float64}:
 0.20145553037952954
 0.1700985774823039
 0.2328124832767552

8.2.7 AIPW ATT Using CausalEstimate.jl

psi_att = ATT(outcome=:Y, treatment=:A, confounders=[:w1, :w2, :w3, :w4])
att_result = estimate(psi_att, AIPW(crossfit=5), df)

(
    ATT    = r(estimate(att_result)),
    CI     = r.(confint(att_result)),
    pvalue = r(pvalue(att_result)),
)

(ATT = 0.189, CI = (0.06404, 0.3139), pvalue = 0.00303)

8.3 TMLE

TMLE is a variant of IF based estimator. It is also a doubly robust estimator. See Mark van der Laan and his coauthors for the recent development of TMLE.

We should expect AIPW and TMLE to give similar answers.

Random.seed!(42)
tmle_result = estimate(psi_ate, TMLE(crossfit=5), df)

(
    TMLE   = r(estimate(tmle_result)),
    CI     = r.(confint(tmle_result)),
    pvalue = r(pvalue(tmle_result)),
)

(TMLE = 0.1775, CI = (0.08638, 0.2687), pvalue = 0.0001349)

@printf "Population ATE (exact):    %.4f\n" pop_ate
@printf "Realized-sample ATE:       %.4f\n" True_Psi

Population ATE (exact):    0.2028
Realized-sample ATE:       0.1840

8.4 DoubleML

DoubleML (Double/Debiased Machine Learning) was proposed by Chernozhukov et al. (2018). An R package DoubleML implements the framework; as of this writing there is no native Julia package, but the algorithm is straightforward to implement directly. The double machine learning framework consists of three key ingredients: Neyman orthogonality, high-quality machine learning estimation, and sample splitting.

They consider a partially linear model:

\[ y_i = \theta d_i + g_0(x_i) + \eta_i \tag{8.6}\]

\[ d_i = m_0(x_i) + v_i \tag{8.7}\]

This model is quite general, except it does not allow interaction of \(d\) and \(x\); therefore no heterogeneous treatment effect across \(x\). But “DoubleML” implements more than partially linear model, it actually allows for heterogeneous treatment effects, in models such as interactive regression model.

The basic idea of doubleML is to use any machine learning algorithm to estimate outcome equation (\(l_0(X) = E(Y | X)\)) and treatment equation (\(m_0(X) = E(D | X)\)). Then get the residuals, namely \(\hat W=Y-\hat l_0(X)\) and \(\hat V = D - \hat m_0(X)\).

Then regress \(\hat W\) on \(\hat V\). Based on FWL theorem, you get \(\hat \theta\).

An important component here is to specify a Neyman-orthogonal score function. The partialling-out form, which is what the residual-on-residual regression below implements, is

\[\psi (W; \theta, \eta) = \big(Y - l(X) - \theta\,(D-m(X))\big)\,(D-m(X)), \tag{8.8}\]

with \(l(X)=E[Y\mid X]\) and \(m(X)=E[D\mid X]\). The estimator \(\hat \theta\) solves the equation that the sample mean of this score function is 0.

The estimator’s variance is the sandwich \(J^{-1}\,E[\psi^2]\,J^{-1}\) with Jacobian \(J = E[\partial\psi/\partial\theta] = -E[(D-m(X))^2]\) — this is what the code computes by dividing the score variance by \(\big(\tfrac1n\sum \hat V_i^2\big)^2\).

using MLJ
using Random

# We implement DoubleML PLR (Partially Linear Regression) manually with cross-fitting
Random.seed!(123)
nsplits = 5
n_samples = nrow(df)
s = shuffle(1:n_samples)
folds = [s[1 + floor(Int, (k-1)*n_samples/nsplits) : floor(Int, k*n_samples/nsplits)] for k in 1:nsplits]

# Base learner (Random Forest for regression)
RandomForestRegressor = @load RandomForestRegressor pkg=DecisionTree verbosity=0
learner = RandomForestRegressor(n_trees=500, max_depth=5)

W_res = zeros(n_samples)
V_res = zeros(n_samples)

for fold in 1:nsplits
    test_idx = folds[fold]
    train_idx = setdiff(1:n_samples, test_idx)
    
    # Fit outcome model l_0(X) = E(Y | X)
    mach_l = machine(learner, W[train_idx, :], df.Y[train_idx])
    MLJ.fit!(mach_l, verbosity=0)
    y_pred = MLJ.predict(mach_l, W[test_idx, :])
    W_res[test_idx] = df.Y[test_idx] .- y_pred
    
    # Fit treatment model m_0(X) = E(D | X)
    mach_m = machine(learner, W[train_idx, :], df.A[train_idx])
    MLJ.fit!(mach_m, verbosity=0)
    a_pred = MLJ.predict(mach_m, W[test_idx, :])
    V_res[test_idx] = df.A[test_idx] .- a_pred
end

# FWL theorem: regress residual W_res on residual V_res
theta_hat = sum(W_res .* V_res) / sum(V_res .^ 2)

# Variance of score function (Neyman-orthogonal score)
psi = (W_res .- theta_hat .* V_res) .* V_res
var_theta = var(psi) / (mean(V_res .^ 2)^2)
se_theta = sqrt(var_theta / n_samples)

ci_lb = theta_hat - 1.96 * se_theta
ci_ub = theta_hat + 1.96 * se_theta
@printf "DML PLR Estimate = %.4g (SE = %.4g)\n" theta_hat se_theta
@printf "95%% CI           = [%.4g, %.4g]\n" ci_lb ci_ub

DML PLR Estimate = 0.192 (SE = 0.03206)
95% CI           = [0.1292, 0.2549]