σ StatsDoge Causal inference workflows
59peers Log in Sign up
12
Workflow·5 steps·branched

Double machine learning for the 401(k) effect (DoubleML)

@doubleml ✓ P D · Jun 3, 2026 · 998 views
Doubly RobustMachine Learning
Source DoubleML — Bach, Chernozhukov, Kurz & Spindler
Summary by StatsDoge

Effect of 401(k) eligibility on net assets via PLR / IRM / IIVM with cross-fit ML nuisances — four learners, one honest comparison.

1

Input · what goes in

Outcome (net financial assets), treatment (401(k) eligibility), covariates, and an instrument (eligibility) for the IV model.

Show data format & exampleHide example

Format — one row per household: outcome y, treatment d, covariates X (and instrument z for IIVM).

  net_tfa(y)  e401(d)  age  inc    educ  ...
   12000        1      41  38k    13
    -400        0      53  21k    11
2

Pipeline · the recipe ⑂ has parallel branches

↑ Click any step in the diagram to read its logic, code, assumptions & discussion.

1
Data prep

Build DoubleMLData (y, d, X, z)

Data preparation — shapes the raw inputs into what the estimator expects.

What happens here

Declare outcome, treatment, covariates and (optionally) an instrument.

Reads from the input data Feeds into #2
Key code 
dml_data = dml.DoubleMLData(df, y_col='net_tfa', d_cols='e401', x_cols=X)
Discussion on this step (0)
  • No comments on this step yet — be the first.
2
Data prep

Choose ML learners for the nuisances

Data preparation — shapes the raw inputs into what the estimator expects.

What happens here

e.g. random forest / lasso for E[Y|X] and E[D|X]; estimated by cross-fitting.

Reads from #1 Feeds into #3#4
Key code 
ml_l = RandomForestRegressor(); ml_m = RandomForestClassifier()
Discussion on this step (0)
  • No comments on this step yet — be the first.
3
Estimation

[DoubleML] Double/debiased ML — PLR / IRM / IIVM

The core estimate — where the causal quantity itself is computed.

What happens here

Partially-linear regression (PLR) for the average effect, with orthogonal scores.

Formula
Y= heta D+g(X)+\varepsilon,\ \ D=m(X)+v;\quad \psi=(Y- heta D-g(X))(D-m(X))
The estimator

Double/debiased ML — PLR / IRM / IIVM — Neyman-orthogonal, cross-fitted estimation of treatment effects with arbitrary ML nuisance learners.

ATE Code @ HEAD ↗ Paper ↗
Reads from #2 Feeds into #5
Key code 
dml.DoubleMLPLR(dml_data, ml_l, ml_m).fit().summary
Discussion on this step (0)
  • No comments on this step yet — be the first.
4
Robustness check

IRM / IIVM cross-checks

A robustness check — does the headline result survive a different lens?

What happens here

Interactive (IRM) and IV (IIVM) models as alternative identifications.

Reads from #2 Feeds into #5
Key code 
dml.DoubleMLIRM(dml_data, ml_g, ml_m).fit()

Reference / docs ↗

Discussion on this step (0)
  • No comments on this step yet — be the first.
5
Reporting

Coefficient comparison plot

Reporting — turn the numbers into a figure or table a reader can act on.

What happens here

Compare estimates across PLR/IRM/IIVM × learners with 95% confidence intervals.

Reads from #3#4 Feeds into the final output
Key code 
# error-bar plot of coef ± 1.96·se per model × learner
Discussion on this step (0)
  • No comments on this step yet — be the first.
3

Output · what you get 4 figures

Effect estimates across PLR / IRM / IIVM models and ML learners, with 95% confidence intervals.
Fig 1Effect estimates across PLR / IRM / IIVM models and ML learners, with 95% confidence intervals.
Distribution of net financial assets — the outcome variable.
Fig 2Distribution of net financial assets — the outcome variable.
Treatment-probability (propensity) distribution by 401(k) eligibility.
Fig 3Treatment-probability (propensity) distribution by 401(k) eligibility.
An overview of the SIPP 401(k) sample used in the example.
Fig 4An overview of the SIPP 401(k) sample used in the example.

Figures reproduced from DoubleML — Bach, Chernozhukov, Kurz & Spindler — unofficial community showcase; all credit to the original authors.

The DoubleML 401(k) example (SIPP data). Build the DoubleMLData, choose ML learners, and estimate the effect three ways (PLR / IRM / IIVM). Unofficial summary.

Discussion (3)

  • 6
    @nuisance_nina · Jun 3, 2026

    Cross-fitting + Neyman-orthogonal scores = you can throw any ML at the nuisances and still get √n-consistent, valid CIs. This is the way.

    5
    @aipw_amir · Jun 3, 2026

    And it generalizes AIPW cleanly. PLR for a quick read, IRM when effects are nonlinear in X.

  • 2
    @kfold_kenji · Jun 3, 2026

    The 401(k) example is the perfect teaching case. Same data, three orthogonal scores, sane answers.

  • 2
    @hte_hannah · Jun 3, 2026

    Pairs beautifully with GRF: DoubleML for the average, causal forest for the heterogeneity.