Causal Ml

Causal Machine Learning

Reference for semiparametric ML estimators: DML with cross-fitting, generalized random forests, debiased regularization, and nuisance function approximation. Covers Neyman-orthogonal moment conditions, sample splitting, plug-in bias correction, and heterogeneous treatment effects.

When to Use This Skill

Use when the user is:

• Estimating treatment effects with high-dimensional controls (p large relative to n)
• Interested in heterogeneous treatment effects (CATE) as a primary estimand
• Applying ML for flexible nuisance function estimation within a causal framework
• Implementing cross-fitting, sample splitting, or Neyman-orthogonal estimators
• Using econml, DoubleML, or grf packages

Skip when:

• Sample is small (n < 500 — ML nuisance models need data)
• A well-specified parametric model is available and defensible
• The task is standard IV/DiD/RDD without high-dimensional controls (use causal-inference skill)
• Structural modeling is needed (use structural-modeling skill)
• The task needs formal identification proof (use identification-proofs skill)

Where to Start

• Choosing a method? Jump to Method Selection Guide
• ATE with many controls? See references/dml.md
• Heterogeneous treatment effects? See references/grf-meta-learners.md
• Variable selection for controls? See references/high-dim-cross-fitting.md
• Reporting HTE results? See references/hte-inference.md
• Connecting to traditional methods? See references/connections-traditional.md

---

Causal ML vs Traditional Methods

Critical point: Causal ML does not relax identification assumptions. If you need a valid instrument, parallel trends, or no unmeasured confounding, those must still hold.

---

Double Machine Learning (DML)

DML (Chernozhukov et al. 2018) fixes regularization bias in naive ML-in-regression. Partial out controls X from both Y and D using separate ML nuisance models, then regress residuals. Two properties: Neyman orthogonality (moment condition locally insensitive to nuisance error) and cross-fitting (prevents overfitting bias).

PLR (Partially Linear Regression): $Y = \theta D + g(X) + \varepsilon$. Workhorse for continuous or binary D with ATE under selection on observables. IRM (Interactive Regression Model): relaxes additive separability for binary D with heterogeneous effects.

Full implementation (Python/R code, cross-fitting from scratch, diagnostics) in references/dml.md.

Causal Forests

Causal forests (Wager-Athey 2018; Athey-Tibshirani-Wager 2019) estimate CATE $\tau(x) = E[Y(1)-Y(0)|X=x]$ using honest forests (structure learned on one subsample, effects estimated on another). Use when CATE is the primary estimand and n $\geq$ 2,000. Always run the calibration test before reporting heterogeneity.

R (grf) and Python (econml) implementations, ATE/ATT extraction, BLP projections in references/grf-meta-learners.md.

Meta-Learners

Decompose CATE estimation into supervised learning sub-problems. DR-Learner (Kennedy 2023): best properties when both nuisance models are well-specified. T-Learner: simplest baseline. X-Learner: designed for imbalanced treatment. For applied work: DR-Learner primary, T-Learner benchmark. Large disagreement signals nuisance model problems.

All implementations in references/grf-meta-learners.md.

High-Dimensional Controls

PDS-LASSO (Belloni-Chernozhukov-Hansen 2014): separate LASSOes of Y on X and D on X, union of selected variables, then OLS. Works at moderate n (~200 with sparse confounders). See references/high-dim-cross-fitting.md.

HTE Inference

Before reporting CATE, test for genuine heterogeneity using BLP calibration test. Do not report heterogeneous effects if calibration test fails (p > 0.10). See references/hte-inference.md.

---

Method Selection Guide

Decision Heuristic

1. n < 500? → Use standard methods (causal-inference skill)
2. High-dim controls (p > 20), want ATE? → PDS-LASSO or DML-PLR; binary D → DML-IRM
3. CATE is primary estimand? → Causal Forest (large n) or DR-Learner (doubly robust)
4. Endogenous treatment with instrument? → DML-PLIV
5. Treatment is rare/imbalanced? → X-Learner
6. Quick benchmark? → Always compute T-Learner as baseline

Full Method Comparison

Limitations to State Explicitly

• ML needs data: Causal forests need n $\geq$ 2,000; DML needs n $\geq$ 500. Below these, use parametric methods.
• Identification is not relaxed: ML is better nuisance estimation, not weaker assumptions.
• CATE inference is hard: Individual-level CIs are conservative; policy targeting requires care.
• Publication: DML and causal forests are mainstream in top applied micro journals. Compare to traditional estimators.

---

Connections to Traditional Methods

Causal ML nests traditional estimators: DML with linear nuisance = OLS (Frisch-Waugh), DML + IV = PLIV, causal forests + instrument = heterogeneous LATE (grf::instrumental_forest), post-LASSO + many instruments = sparse instrument selection then 2SLS. Details in references/connections-traditional.md.

---

Integration with Plugin

Agents: econometric-reviewer (post-estimation review, table/code consistency), identification-critic (IV/PLIV assumptions), numerical-auditor (convergence, seeding, Monte Carlo validation).

Cross-references: empirical-playbook skill → sensitivity-analysis.md (specification curve over ML choices), empirical-playbook skill → diagnostic-battery.md (nuisance R², overlap, calibration), numerical-auditor agent (synthetic data with known CATE).

Relationship to `causal-inference` skill: Use causal-inference to establish identification; use causal-ml for implementation with high-dimensional controls or when heterogeneity is primary. Complements, not substitutes.

Reference Files

• references/dml.md — Full DML implementation: PLR, IRM, PLIV with econml/DoubleML, cross-fitting, diagnostics
• references/grf-meta-learners.md — Causal forests (grf/econml), DR/T/S/X-Learner, calibration tests
• references/high-dim-cross-fitting.md — PDS-LASSO, Belloni-Chernozhukov-Hansen, cross-fitting protocols
• references/hte-inference.md — Calibration tests, individual CATE CIs, BLP projections, subgroup analysis
• references/connections-traditional.md — DML-OLS equivalence, PLIV, instrumental forests, post-LASSO

Connections to Traditional Methods — Reference

Understanding how causal ML relates to traditional methods helps build intuition and credibility with traditional audiences.

DML Reduces to IV When Nuisance Models Are Linear

If E[Y|X] and E[D|X] are both estimated by OLS (linear projections), then DML's partialling out is numerically identical to the Frisch-Waugh-Lovell theorem. The DML θ̂ equals the OLS coefficient on D in a regression of Y on D and X.

This means: DML with linear nuisance models = standard OLS. DML adds value precisely when the nuisance functions are nonlinear — it allows flexible control for X while maintaining √n-inference on θ.

Practical check: Run DML with linear nuisance models (OLS) and compare to OLS with all controls. They should match. If not, there is a coding error.

IV with DML Nuisance Models

DML extends naturally to IV. The partially linear IV model:

Y = θ₀ D + g₀(X) + ε
D = m₀(X) + v
Z: instrument with E[Z · ε | X] = 0

Cross-fitted IV moment: regress residualized Y on residualized D, instrumenting with residualized Z.

# DoubleML: Partially linear IV
pliv = dml.DoubleMLPLIV(
    obj_dml_data=data,   # data must include Z (instrument)
    ml_g=ml_g,           # learner for E[Y|X]
    ml_m=ml_m,           # learner for E[D|X]
    ml_r=ml_r,           # learner for E[Z|X]  — partialling out Z
    n_folds=5,
)
pliv.fit()
print(pliv.summary)

Causal Forests Generalize Local ATE

Standard IV/2SLS at a single instrument value (e.g., an RDD cutoff) gives LATE for compliers at that point. A causal forest with an instrument generalizes this to heterogeneous LATE across the covariate space:

# R: grf — instrumental forest
iv_forest <- instrumental_forest(
  X = X_matrix,
  Y = Y_vector,
  W = W_treatment,   # endogenous treatment
  Z = Z_instrument,  # instrument
  seed = 42
)
tau_late_hat <- predict(iv_forest)$predictions
ate_late <- average_treatment_effect(iv_forest)

Post-LASSO Generalizes 2SLS with Many Instruments

The many-instruments problem (Bekker 1994) causes 2SLS to be inconsistent when the number of instruments grows with n. Post-LASSO selects a sparse set of strong instruments, then runs standard 2SLS on the selected instruments. This connects to LIML and jackknife IV estimators.

# R: hdm — LASSO for many instruments
# First, select relevant instruments using LASSO
lasso_z <- rlasso(D_vector ~ Z_matrix)  # regress D on instruments
selected_z <- which(lasso_z$coef != 0)

# Then run 2SLS with selected instruments
library(fixest)
iv_formula <- as.formula(
  paste("Y ~", paste(X_names, collapse = "+"),
        "| D ~ ", paste(Z_names[selected_z], collapse = "+"))
)
result_iv <- feols(iv_formula, data = df, vcov = "hetero")
print(result_iv)

Double Machine Learning — Implementation Reference

Reference: Chernozhukov, Chetverikov, Demirer, Duflo, Hansen, Newey, Robins (2018), "Double/debiased machine learning for treatment and structural parameters," Econometrics Journal.

Core Idea

Naive approach: regress Y on D and X with ML. This fails because ML regularization (LASSO shrinkage, random forest bias) contaminates the coefficient on D. The bias does not vanish even as n → ∞.

DML fix: partial out X from both Y and D using separate ML models, then regress the residuals on each other. The key properties that make this work:

1. Neyman orthogonality: The moment condition is locally insensitive to perturbations in the nuisance parameters. Small errors in nuisance estimates have second-order (not first-order) effects on the target parameter. 2. Cross-fitting: Estimate nuisance models on a held-out fold to avoid overfitting bias contaminating the main estimate.

Partially Linear Model (PLR)

The PLR is the workhorse DML specification:

Y = θ₀ D + g₀(X) + ε,   E[ε | D, X] = 0
D = m₀(X) + v,           E[v | X] = 0

where g₀(X) is an unknown function of controls X, and θ₀ is the ATE of interest. The nuisance functions are g₀ and m₀.

Identification assumption: After conditioning on X, D is as good as randomly assigned. This is selection on observables — the same assumption as standard regression, but allowing the functional form of X to be flexible.

Interactive Regression Model (IRM)

When treatment D is binary and the effect may be heterogeneous:

Y = g₀(D, X) + ε,   E[ε | D, X] = 0
D ~ Bernoulli(m₀(X))

The IRM estimates the ATE by averaging individual-level predictions:

θ₀ = E[g₀(1, X) - g₀(0, X)]

Use IRM when:

• D is binary and you suspect treatment effect heterogeneity
• You want ATE rather than a single θ coefficient
• The partially linear assumption (additive separability) seems too strong

Cross-Fitting Procedure

Cross-fitting prevents overfitting bias from contaminating inference. The K-fold procedure (K=5 is standard):

import numpy as np
from sklearn.model_selection import KFold

def cross_fit_residuals(Y, D, X, ml_model_y, ml_model_d, n_splits=5, random_state=42):
    """
    Cross-fitting step for DML partially linear model.

    Returns:
        W: residuals Y - E[Y|X] (partialled-out Y)
        V: residuals D - E[D|X] (partialled-out D)
    """
    n = len(Y)
    W = np.zeros(n)  # Y residuals
    V = np.zeros(n)  # D residuals

    kf = KFold(n_splits=n_splits, shuffle=True, random_state=random_state)

    for train_idx, test_idx in kf.split(X):
        # Train nuisance models on training fold
        ml_model_y.fit(X[train_idx], Y[train_idx])
        ml_model_d.fit(X[train_idx], D[train_idx])

        # Predict and residualize on held-out test fold
        W[test_idx] = Y[test_idx] - ml_model_y.predict(X[test_idx])
        V[test_idx] = D[test_idx] - ml_model_d.predict(X[test_idx])

    return W, V

def dml_plr_estimate(W, V):
    """
    DML estimate from partialled-out residuals.
    theta_hat = (V'W) / (V'V)  — OLS of W on V (no intercept)
    Standard errors via influence function.
    """
    n = len(W)
    theta_hat = np.dot(V, W) / np.dot(V, V)

    # Influence function: psi_i = V_i * (W_i - theta_hat * V_i)
    psi = V * (W - theta_hat * V)

    # Sandwich variance
    J = -np.mean(V ** 2)
    var_hat = np.mean(psi ** 2) / (J ** 2)
    se = np.sqrt(var_hat / n)

    return theta_hat, se

Using the DoubleML Package (Python)

import doubleml as dml
import numpy as np
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import LassoCV

# Setup: data object
# Y: outcome (1D array), D: treatment (1D array), X: controls (2D array)
data = dml.DoubleMLData.from_arrays(X=X, y=Y, d=D)

# Choose learners for nuisance functions
# For continuous D: two regression learners
ml_g = RandomForestRegressor(n_estimators=100, max_depth=5, random_state=42)
ml_m = RandomForestRegressor(n_estimators=100, max_depth=5, random_state=42)

# Partially Linear Regression model
plr = dml.DoubleMLPLR(
    obj_dml_data=data,
    ml_g=ml_g,      # learner for E[Y|X]
    ml_m=ml_m,      # learner for E[D|X]
    n_folds=5,
    score='partialling out',
)
plr.fit()
print(plr.summary)

# For binary D: use classification learner for propensity
ml_m_binary = RandomForestClassifier(n_estimators=100, max_depth=5, random_state=42)
irm = dml.DoubleMLIRM(
    obj_dml_data=data,
    ml_g=ml_g,
    ml_m=ml_m_binary,
    n_folds=5,
    score='ATE',
)
irm.fit()
print(irm.summary)

# Cluster-robust standard errors
plr_clustered = dml.DoubleMLPLR(data, ml_g, ml_m, n_folds=5)
plr_clustered.fit()
# Pass cluster variable:
# data = dml.DoubleMLData.from_arrays(X=X, y=Y, d=D, cluster_cols=cluster_ids)

Using the DoubleML Package (R)

library(DoubleML)
library(mlr3)
library(mlr3learners)

# Create DoubleML data object
dml_data <- DoubleMLData$new(
  data = df,
  y_col = "outcome",
  d_cols = "treatment",
  x_cols = c("x1", "x2", "x3")  # control variables
)

# Specify learners (mlr3 ecosystem)
learner_g <- lrn("regr.ranger", num.trees = 100, max.depth = 5)
learner_m <- lrn("regr.ranger", num.trees = 100, max.depth = 5)

# Partially linear regression
plr <- DoubleMLPLR$new(
  data = dml_data,
  ml_g = learner_g,
  ml_m = learner_m,
  n_folds = 5
)
plr$fit()
plr$summary()

# For binary treatment (IRM)
learner_m_cls <- lrn("classif.ranger", num.trees = 100, max.depth = 5,
                     predict_type = "prob")
irm <- DoubleMLIRM$new(
  data = dml_data,
  ml_g = learner_g,
  ml_m = learner_m_cls,
  n_folds = 5,
  score = "ATE"
)
irm$fit()
irm$summary()

DML Diagnostic Checklist

• [ ] Nuisance fit quality: Report R² (or classification accuracy) for both nuisance models (E[Y|X] and E[D|X]). Low R² on E[D|X] implies weak "first stage" — the controls barely explain treatment variation.
• [ ] Residual balance: After partialling out, regress V (D residuals) on X — coefficients should be near zero. If not, the ML model for E[D|X] is misspecified.
• [ ] Cross-fitting fold stability: Repeat with different random seeds. Estimates should be stable across seeds. Large variation implies insufficient sample size for the chosen ML method.
• [ ] Compare K=5 vs K=10: If estimates differ substantially, sample size may be too small for cross-fitting to work well.
• [ ] Neyman orthogonality check: Perturb nuisance estimates slightly — the main estimate should be insensitive. Large sensitivity suggests the score is not sufficiently orthogonal.
• [ ] Trim extreme propensity scores: For binary D, trim observations where E[D|X] is near 0 or 1 (e.g., below 0.01 or above 0.99). Extreme values inflate variance.

Common DML Pitfalls

Causal Forests and Meta-Learners — Implementation Reference

Causal Forests (Generalized Random Forests)

Reference: Athey, Tibshirani, Wager (2019), "Generalized random forests," Annals of Statistics. Wager and Athey (2018), "Estimation and inference of heterogeneous treatment effects using random forests," JASA.

Core Idea

Causal forests estimate the CATE τ(x) = E[Y(1) - Y(0) | X = x] at any point x. The key innovation over standard random forests is honesty: the tree structure is learned on one subsample, and the leaf-level treatment effect is estimated on a separate subsample. This prevents overfitting from conflating the splitting criterion with the estimation.

Honesty is necessary for valid confidence intervals. Without it, leaf estimates are biased and confidence intervals have poor coverage.

Intuition: Local ATE via Weighted Neighbors

Causal forests solve:

τ̂(x) = argmin_τ Σᵢ αᵢ(x) · [Yᵢ - m̂(Xᵢ) - τ · (Dᵢ - ê(Xᵢ))]²

where αᵢ(x) are forest weights (how much unit i's neighborhood contributes to τ(x)), and m̂(X), ê(X) are residualized outcomes and propensities. Units that are neighbors of x in feature space get high weight. Units far away get low weight.

This is local ATE estimation, where "local" is defined by proximity in the feature space learned by the forest.

Python: CausalForestDML via econml

from econml.dml import CausalForestDML
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
import numpy as np

# X: features for CATE heterogeneity (can differ from controls W)
# T: binary or continuous treatment
# Y: outcome
# W: high-dimensional controls to partial out (optional, separate from X)

cf = CausalForestDML(
    model_y=RandomForestRegressor(n_estimators=100, random_state=42),
    model_t=RandomForestClassifier(n_estimators=100, random_state=42),
    n_estimators=500,
    min_samples_leaf=5,
    max_depth=None,
    random_state=42,
    cv=5,           # cross-fitting folds
    honest=True,    # always use honest splitting
)
cf.fit(Y=y, T=t, X=X, W=W)  # W is additional controls not in X

# Estimate CATE for each observation
tau_hat = cf.effect(X)

# Estimate ATE and confidence interval
ate_result = cf.ate(X, T0=0, T1=1)
print(f"ATE: {ate_result:.4f}")

# Confidence intervals for each unit (conservative)
tau_lb, tau_ub = cf.effect_interval(X, alpha=0.05)

# Best linear projection of CATE onto features
from econml.inference import LinearModelFinalInference
blp = cf.const_marginal_effect_inference(X)
print(blp.summary_frame())

R: grf Package

library(grf)

# Prepare data
# X: matrix of features, Y: outcome vector, W: treatment vector

cf <- causal_forest(
  X = X_matrix,
  Y = Y_vector,
  W = W_treatment,
  num.trees = 2000,
  honesty = TRUE,          # required for valid inference
  min.node.size = 5,
  seed = 42
)

# Average treatment effect
ate <- average_treatment_effect(cf, target.sample = "all")
cat("ATE:", ate["estimate"], "+/-", 1.96 * ate["std.err"], "\n")

# ATT
att <- average_treatment_effect(cf, target.sample = "treated")

# CATE predictions with confidence intervals
tau_hat <- predict(cf, estimate.variance = TRUE)
tau_vals <- tau_hat$predictions
tau_se   <- sqrt(tau_hat$variance.estimates)

# Test for heterogeneity: BLP calibration test
# Chernozhukov, Demirer, Duflo, Fernandez-Val (2022)
calibration <- test_calibration(cf)
print(calibration)
# mean.forest.prediction: should be ~1 if CATE is well-calibrated
# differential.forest.prediction: should be >0 if heterogeneity is real

# Best linear projection of CATE on covariates
blp <- best_linear_projection(cf, A = X_matrix)
print(blp)

CATE Heterogeneity: Calibration Test

The calibration test (Chernozhukov, Demirer, Duflo, Fernandez-Val 2022) estimates a linear projection:

τᵢ = α₀ + α₁ · τ̂ᵢ + εᵢ

using an AIPW-based approach. Interpretation:

• α₁ ≈ 1: forest predictions are well-calibrated on average
• α₁ > 0 and significant: there is real heterogeneity (the forest is detecting genuine variation, not noise)
• α₁ = 0: forest's heterogeneity is indistinguishable from noise

Report this test whenever presenting CATE estimates.

# R (grf)
cal <- test_calibration(cf)
# Examine: mean.forest.prediction coefficient and its p-value
# Examine: differential.forest.prediction coefficient and its p-value

# Python (econml): use the best_linear_projection API

Causal Forest Diagnostic Checklist

• [ ] Honesty enabled: Always set honest = TRUE (R) or honest=True (Python). Without honesty, confidence intervals are invalid.
• [ ] Calibration test: Run test_calibration(). Report both coefficients. Significant differential coefficient supports real heterogeneity.
• [ ] ATE recovery: Compare forest ATE to a standard doubly-robust ATE estimator. They should agree closely. Large discrepancy suggests a problem with nuisance models.
• [ ] Overlap / positivity: Check that propensity scores ê(X) are bounded away from 0 and 1. Forest fails when treatment assignment is deterministic given X.
• [ ] Variable importance: Examine variable_importance(cf) (R) or cf.feature_importances_ (Python). Dominant variables driving heterogeneity should be interpretable.
• [ ] Minimum leaf size: Default min.node.size=5 is a starting point. Increase for small samples; the forest should not have near-empty leaves.
• [ ] Number of trees: Use at least 2,000 trees for stable variance estimates. More trees reduce Monte Carlo error in τ̂(x).
• [ ] Subgroup analysis: Pre-specify subgroups before running the forest. Post-hoc "we found heterogeneity along dimension k" inflates false discovery rates.

Common Causal Forest Pitfalls

---

DR-Learner and Meta-Learners

Meta-learners decompose the CATE estimation problem into standard supervised learning sub-problems. The choice of meta-learner determines the statistical properties of τ̂(x).

Reference: Kennedy (2023), "Towards optimal doubly robust estimation of heterogeneous causal effects," Electronic Journal of Statistics. Künzel et al. (2019), "Meta-learners for estimating heterogeneous treatment effects using machine learning," PNAS.

Overview of Meta-Learners

Recommendation for applied work: DR-Learner when sample is large enough for nuisance estimation. T-Learner as a simple benchmark. Report both.

DR-Learner: Doubly Robust CATE

The DR-Learner constructs pseudo-outcomes:

ψᵢ = μ̂₁(Xᵢ) - μ̂₀(Xᵢ)
    + Dᵢ(Yᵢ - μ̂₁(Xᵢ)) / ê(Xᵢ)
    - (1-Dᵢ)(Yᵢ - μ̂₀(Xᵢ)) / (1 - ê(Xᵢ))

Then regresses these pseudo-outcomes on X to get τ̂(x). The pseudo-outcomes are doubly robust: if either the outcome model or propensity model is correct, the pseudo-outcome has the correct expectation.

from econml.dr import DRLearner
from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier
from sklearn.linear_model import RidgeCV

# DR-Learner
dr = DRLearner(
    model_propensity=RandomForestClassifier(n_estimators=100, random_state=42),
    model_regression=RandomForestRegressor(n_estimators=100, random_state=42),
    model_final=RidgeCV(),       # final CATE model (can be any regressor)
    cv=5,
    random_state=42,
)
dr.fit(Y=y, T=t, X=X, W=W)

# CATE estimates
tau_dr = dr.effect(X)

# ATE from DR-Learner
ate_dr = dr.ate(X)
print(f"ATE (DR-Learner): {ate_dr:.4f}")

# Confidence intervals
ate_interval = dr.ate_interval(X, alpha=0.05)
print(f"95% CI: [{ate_interval[0]:.4f}, {ate_interval[1]:.4f}]")

# T-Learner for comparison
from econml.metalearners import TLearner

tl = TLearner(
    models=RandomForestRegressor(n_estimators=100, random_state=42)
)
tl.fit(y, t, X=X)
tau_tl = tl.effect(X)

When to Use Each Meta-Learner

• T-Learner: Quick baseline; when treatment groups are roughly balanced; when you have very large n and flexible models
• S-Learner: When treatment effect is expected to be small or zero for most units (LASSO/tree won't shrink effect to zero unlike T-Learner)
• X-Learner: When treatment is rare or imbalanced (many control, few treated); designed specifically for this case
• DR-Learner: When you want the best-calibrated CATE estimates with valid inference; default for serious empirical work

DR-Learner Diagnostic Checklist

• [ ] Propensity model quality: Check AUC and calibration of propensity score. Miscalibrated propensities inflate DR pseudo-outcome variance.
• [ ] Outcome model quality: Report R² for μ̂₁(X) and μ̂₀(X) separately. Low R² reduces efficiency but does not invalidate doubly-robust property (as long as propensity is right).
• [ ] Compare T-Learner and DR-Learner: If they agree closely, results are likely robust. Large disagreements suggest a nuisance specification problem.
• [ ] ATE vs. mean of CATE: np.mean(tau_dr) should match dr.ate(X) — if not, there is a weighting issue.
• [ ] Calibration: Apply Chernozhukov calibration test logic: project CATE onto a low-dimensional summary; check that projection coefficient is not zero.

High-Dimensional Controls and Cross-Fitting — Implementation Reference

High-Dimensional Controls

Reference: Belloni, Chernozhukov, Hansen (2014), "Inference on treatment effects after selection among high-dimensional controls," Review of Economic Studies.

When You Need This

You have many candidate control variables (p large relative to n) and want to: 1. Avoid overfitting by selecting controls automatically 2. Maintain valid inference on a treatment effect after selection 3. Avoid the Leeb-Pötscher problem: you cannot do inference on θ after LASSO selection of X unless you account for selection

Key insight: Running LASSO to predict Y and then doing OLS on the selected variables gives biased inference on D. You need post-double selection (PDS-LASSO) to avoid this.

Post-Double Selection LASSO (PDS-LASSO)

The Belloni-Chernozhukov-Hansen (2014) procedure:

1. Run LASSO of Y on X → select variables S₁ 2. Run LASSO of D on X → select variables S₂ 3. Union: S = S₁ ∪ S₂ 4. Run OLS of Y on D and all variables in S

The union step is critical: including variables that predict D (even if they don't add predictive power for Y) controls for confounders. Including variables that predict Y (even if they don't add for D) improves efficiency.

Python: Manual PDS-LASSO

import numpy as np
from sklearn.linear_model import LassoCV
import statsmodels.api as sm

def pds_lasso(Y, D, X, n_splits=5, random_state=42):
    """
    Post-double selection LASSO.

    Returns OLS estimate of treatment effect and standard error,
    controlling for selected variables.
    """
    # Step 1: LASSO of Y on X
    lasso_y = LassoCV(cv=n_splits, random_state=random_state)
    lasso_y.fit(X, Y)
    selected_y = np.where(np.abs(lasso_y.coef_) > 0)[0]

    # Step 2: LASSO of D on X
    lasso_d = LassoCV(cv=n_splits, random_state=random_state)
    lasso_d.fit(X, D)
    selected_d = np.where(np.abs(lasso_d.coef_) > 0)[0]

    # Step 3: Union of selected variables
    selected = np.union1d(selected_y, selected_d)
    print(f"Variables selected by Y-LASSO: {len(selected_y)}")
    print(f"Variables selected by D-LASSO: {len(selected_d)}")
    print(f"Union size: {len(selected)}")

    # Step 4: OLS with selected controls
    if len(selected) > 0:
        controls = X[:, selected]
        regressors = np.column_stack([D, controls])
    else:
        regressors = D.reshape(-1, 1)

    regressors_with_const = sm.add_constant(regressors)
    ols = sm.OLS(Y, regressors_with_const).fit(cov_type='HC3')

    # Treatment effect is the coefficient on D (index 1 after constant)
    theta_hat = ols.params[1]
    se = ols.bse[1]
    ci = ols.conf_int().iloc[1]

    return {
        'theta': theta_hat,
        'se': se,
        'ci_lo': ci[0],
        'ci_hi': ci[1],
        'n_selected': len(selected),
        'selected_idx': selected,
    }

R: hdm Package

library(hdm)

# Post-double selection LASSO via hdm
# Single treatment variable
pds <- rlassoEffect(
  x = X_matrix,        # control variables (n x p matrix)
  y = Y_vector,        # outcome
  d = D_vector,        # treatment
  method = "double selection"
)
print(pds)

# Inference on multiple treatment variables simultaneously
pds_multi <- rlassoEffects(
  x = X_matrix,
  y = Y_vector,
  d = D_matrix,        # multiple treatment variables
  method = "double selection"
)
summary(pds_multi)

# LASSO for variable selection only (then examine selected set)
lasso_y <- rlasso(Y_vector ~ X_matrix)
lasso_d <- rlasso(D_vector ~ X_matrix)

selected_y <- which(lasso_y$coef != 0)
selected_d <- which(lasso_d$coef != 0)
selected_union <- union(selected_y, selected_d)
cat("Union size:", length(selected_union), "\n")

Practical Guidance for High-Dimensional Controls

Choosing the LASSO penalty:

• Use theory-based (Belloni-Chernozhukov) penalty: λ = 2c · σ̂ · √(n log p) for some constant c. This is what hdm::rlasso uses by default.
• Cross-validation (LassoCV) is common in practice but does not have the same theoretical guarantees for post-selection inference. Prefer hdm for formal inference.

When p > n:

• PDS-LASSO still works if the true model is sparse (few controls truly matter)
• If the true model is dense (many controls each with small effect), consider ridge or elastic net nuisance models within DML instead

Interactions and polynomials:

• You may want to include interactions D × X in the control set for the Y-LASSO step (to detect effect modifiers)
• But do NOT include D × X in the D-LASSO step — these are endogenous by construction

PDS-LASSO Diagnostic Checklist

• [ ] Selected variable set is interpretable: Review which controls were selected. Variables strongly correlated with both Y and D should appear in the union.
• [ ] First-stage effective F-stat: After union selection, check that D is not partialled out (residual variance is not too small). Compute F from regression of D on selected controls.
• [ ] Sensitivity to LASSO penalty: Vary λ by factor of 0.5 and 2. Selected set should not change dramatically.
• [ ] Compare PDS to OLS with all controls: If PDS estimate differs substantially from OLS with full X, either the high-dimensional OLS is overfitting or there is important nonlinearity.
• [ ] Sparsity assumption: PDS-LASSO requires that few controls truly matter. If you expect dense effects (all controls matter a little), DML with ridge/elastic net is more appropriate.
• [ ] Post-selection F-stat on treatment: After union selection, report the partial F-statistic on D in the first-stage regression — confirms that the selected controls do not absorb all variation in D.

Common PDS-LASSO Pitfalls

---

Sample Splitting and Cross-Fitting

Why Naive ML-in-Regression Fails

Consider fitting Y ~ θD + g(X) with LASSO. The LASSO regularizer penalizes θ just as it penalizes the coefficients on X. Even with large n, θ̂ is biased toward zero by the regularization — this bias does not vanish.

More generally, if you use the same data to (a) learn the nuisance function g(X) and (b) estimate θ, the estimation error in (a) contaminates (b) at first order. The result is that √n-convergence of θ̂ breaks down.

Solution: Cross-fitting separates these two estimation tasks across data folds.

K-Fold Cross-Fitting Step by Step

1. Partition {1, ..., n} into K folds I₁, ..., I_K of approximately equal size.

2. For k = 1, ..., K:
   a. Training set: I^c_k = {1,...,n} \ I_k  (all folds except fold k)
   b. Fit nuisance model ĝ_k on training set I^c_k
   c. Compute residuals Ŵ_i = Y_i - ĝ_k(X_i) for all i ∈ I_k

3. Each observation gets one residual Ŵ_i from the fold in which it was held out.

4. Same procedure for D: Ṽ_i = D_i - m̂_k(X_i) for i ∈ I_k

5. Pool all residuals: use (Ŵ₁, ..., Ŵ_n) and (Ṽ₁, ..., Ṽ_n) for final inference.

6. θ̂ = (Σ Ṽ_i Ŵ_i) / (Σ Ṽ_i²) — OLS of Ŵ on Ṽ (no intercept)

Practical Choices for K

Tip: With K=5, each nuisance model is trained on 80% of the data. This is large enough for random forests and LASSO to be well-fit on most empirically realistic samples (n > 1,000).

Cross-Fitting from Scratch (Illustrative)

import numpy as np
from sklearn.model_selection import KFold
from sklearn.ensemble import RandomForestRegressor
import statsmodels.api as sm

def dml_crossfit(Y, D, X, n_splits=5, seed=42):
    """
    Full DML cross-fitting with inference.
    Assumes partially linear model: Y = theta*D + g(X) + eps
    """
    n = len(Y)
    W_hat = np.zeros(n)  # Y - E[Y|X]
    V_hat = np.zeros(n)  # D - E[D|X]

    kf = KFold(n_splits=n_splits, shuffle=True, random_state=seed)

    r2_y_list, r2_d_list = [], []

    for fold_idx, (train_idx, test_idx) in enumerate(kf.split(X)):
        # Nuisance models
        rf_y = RandomForestRegressor(n_estimators=200, max_depth=5, random_state=seed)
        rf_d = RandomForestRegressor(n_estimators=200, max_depth=5, random_state=seed)

        rf_y.fit(X[train_idx], Y[train_idx])
        rf_d.fit(X[train_idx], D[train_idx])

        Y_pred = rf_y.predict(X[test_idx])
        D_pred = rf_d.predict(X[test_idx])

        W_hat[test_idx] = Y[test_idx] - Y_pred
        V_hat[test_idx] = D[test_idx] - D_pred

        # Nuisance fit diagnostics
        ss_res_y = np.sum((Y[test_idx] - Y_pred) ** 2)
        ss_tot_y = np.sum((Y[test_idx] - Y[test_idx].mean()) ** 2)
        r2_y_list.append(1 - ss_res_y / ss_tot_y)

        ss_res_d = np.sum((D[test_idx] - D_pred) ** 2)
        ss_tot_d = np.sum((D[test_idx] - D[test_idx].mean()) ** 2)
        r2_d_list.append(1 - ss_res_d / ss_tot_d)

    print(f"Mean R2 (Y nuisance): {np.mean(r2_y_list):.3f}")
    print(f"Mean R2 (D nuisance): {np.mean(r2_d_list):.3f}")

    # DML estimate
    theta_hat = np.dot(V_hat, W_hat) / np.dot(V_hat, V_hat)

    # Influence function SE
    psi = V_hat * (W_hat - theta_hat * V_hat)
    J = np.mean(V_hat ** 2)
    var = np.mean(psi ** 2) / J ** 2
    se = np.sqrt(var / n)

    ci_lo = theta_hat - 1.96 * se
    ci_hi = theta_hat + 1.96 * se

    print(f"\nDML Estimate: {theta_hat:.4f}")
    print(f"SE:           {se:.4f}")
    print(f"95% CI:       [{ci_lo:.4f}, {ci_hi:.4f}]")

    return theta_hat, se

# Usage
theta, se = dml_crossfit(Y=y, D=d, X=X_controls)

Aggregating Estimates Across Folds

When implementing DML with repeated cross-fitting (recommended for stability), run the full K-fold procedure M times with different random seeds and aggregate:

def dml_repeated(Y, D, X, n_splits=5, n_reps=5):
    """DML with repeated cross-fitting for stability."""
    estimates = []
    for rep in range(n_reps):
        theta_r, _ = dml_crossfit(Y, D, X, n_splits=n_splits, seed=rep * 42)
        estimates.append(theta_r)

    # Median aggregation (more robust than mean)
    theta_final = np.median(estimates)
    print(f"Median across {n_reps} repetitions: {theta_final:.4f}")
    print(f"Std across repetitions: {np.std(estimates):.4f}")
    return theta_final

A large standard deviation across repetitions signals that the ML models are unstable — either the sample is too small or the models are too complex.

HTE Inference and Reporting — Reference

Global Test for Heterogeneity

Before reporting CATE estimates, test whether there is genuine heterogeneity. The BLP (best linear projection) approach due to Chernozhukov, Demirer, Duflo, Fernandez-Val (2022) is the standard test:

# R: grf
cal <- test_calibration(cf)
# H0: no heterogeneity (differential.forest.prediction = 0)
# Reject H0 → genuine heterogeneity detected

# Python: econml (manual BLP)
from econml.inference import LinearModelFinalInference
# Use cf.effect() predictions and regress on summary statistics

Do not report heterogeneous effects if the calibration test fails to reject at a reasonable level (p > 0.10). Report the calibration test result alongside CATE estimates.

Confidence Intervals on Individual CATE

Individual CATE confidence intervals from causal forests are valid but conservative (they are honest pointwise CIs, not uniform CIs). They should be interpreted as uncertainty about τ(xᵢ), not as evidence that τ(xᵢ) ≠ 0 for that individual.

# econml: point estimates and CIs for each unit
tau_lb, tau_ub = cf.effect_interval(X_test, alpha=0.05)

# R: grf
tau_hat <- predict(cf, estimate.variance = TRUE)
tau_ci_lo <- tau_hat$predictions - 1.96 * sqrt(tau_hat$variance.estimates)
tau_ci_hi <- tau_hat$predictions + 1.96 * sqrt(tau_hat$variance.estimates)

Warning: Do not use individual CIs for policy targeting without accounting for multiple testing. Targeting based on wide CIs that nominally include zero for some units and not others leads to invalid inference.

Best Linear Projection of CATE

The best linear projection (BLP) onto covariates gives a sparse, interpretable summary of heterogeneity:

# econml: summary of CATE heterogeneity
blp = cf.const_marginal_effect_inference(X).summary_frame()
print(blp)  # coefficient on each X variable in BLP of CATE

# R: grf
blp <- best_linear_projection(cf, A = X_matrix)
print(blp)
# Coefficients tell you: which observed characteristics predict larger/smaller CATE

Subgroup Analysis: Pre-Specified vs Data-Driven

Pre-specified subgroups (defined before analysis):

• Report group-specific ATEs using forest-weighted estimators
• Standard: average_treatment_effect(cf, subset = group_indicator) in R

Data-driven subgroups (quartiles of τ̂(x)):

• Compute quartile cutoffs of τ̂ on a held-out sample
• Report ATEs within quartiles — this is exploratory, not confirmatory
• Requires multiplicity correction (Benjamini-Hochberg) if multiple subgroups reported

# R: subgroup ATE using forest
high_effect <- tau_hat$predictions > median(tau_hat$predictions)

ate_high <- average_treatment_effect(cf, subset = high_effect)
ate_low  <- average_treatment_effect(cf, subset = !high_effect)

cat("ATE (high CATE group):", ate_high["estimate"], "\n")
cat("ATE (low CATE group): ", ate_low["estimate"],  "\n")