Back to syllabus中文版

Frontier / Deep Learning Causal

Deep Learning for Causal Inference: Representation Learning and Neural Estimators (TARNet, Dragonnet, DeepIV, CEVAE)

Neural networks do not change the identifying assumptions; they push "fit nuisances with flexible functions and correct covariate shift with a shared representation" to the limit. Understanding TARNet / Dragonnet / DeepIV / CEVAE tells you exactly what deep learning does — and does not do — in a causal workflow.

This page follows DML / causal forests. The core idea: under unconfoundedness (or a valid instrument), replacing the outcome regression, propensity, representation, and even confounder proxies with neural networks lets you estimate ATE / CATE with high-dimensional, text, or image covariates — but identification still comes from the research design, not the network.

Schematic

The principle at a glance

Deep-learning causal: shared rep + two headscovariates Xhigh-dim/textshared rep Φ(x)balancing IPMtwo headsh₀(Φ), h₁(Φ)τ(x) / ATEh₁ − h₀TARNet = shared rep + 2 headsDragonnet = + propensity headDeepIV/CEVAE = endog./latent
The thread of neural causal estimation: high-dim/text covariates X → learn a shared representation Φ(x) → two outcome heads give treated and control counterfactuals → difference is the CATE, expectation the ATE. A balancing IPM controls extrapolation; the Dragonnet propensity head aligns the estimate to the ATE. Identification still rests on unconfoundedness/exclusion; the network only swaps functional form.

Start Here

What you should be able to do

01

Separate the identifying assumption (unconfoundedness / exclusion) from the estimator (a neural network); never treat the latter as relaxing the former.

02

Understand TARNet: a shared representation plus two outcome heads to estimate the individual effect tau(x).

03

Understand representation balancing (CFR): an IPM / Wasserstein penalty on the distance between treated and control representation distributions.

04

Understand Dragonnet: add a propensity head plus targeted regularization on the shared representation for a more stable ATE.

05

Know where DeepIV (two-stage neural IV for endogenous treatment) and CEVAE (latent-variable models for confounder proxies) apply and where they break.

Learning Path

Learning path: represent → two heads → balance → propensity → debias

Read neural causal estimation along this path: learn a representation, split treatment heads, add a balancing penalty to control extrapolation, and when needed add a propensity head with targeted regularization to align the estimate to the ATE.

  1. Step 1

    Represent

    The network learns a shared representation Phi(x) for high-dim / text covariates.

    Phi(x)

  2. Step 2

    Two heads

    Two outcome heads give the treated and control counterfactual predictions.

    h_0,h_1

  3. Step 3

    Balance

    An IPM penalty on the two representation distributions controls extrapolation error.

    alpha·IPM

  4. Step 4

    Propensity

    Dragonnet adds a propensity head to retain assignment information.

    g(x)

  5. Step 5

    Debias

    Targeted regularization / AIPW aligns "predict well" with "estimate the ATE".

    EIF constraint

01 / Intuition

Core Intuition

Classical estimators (OLS, IPW, matching) struggle when covariates are high-dimensional or unstructured (text, images, panel sequences), exactly where neural networks excel at learning representations.

Key insight: factor the conditional mean mu_d(x)=E[Y|D=d,X=x] into a shared representation Phi(x) plus one output head per treatment, so most parameters are shared and only the heads separate treatments — letting control units help estimate treated counterfactuals.

But a flexible representation can amplify covariate shift between treated and control; representation balancing (CFR) penalizes the distance between the two representation distributions with an integral probability metric, trading a little bias for more robust counterfactual extrapolation.

02 / Math

From potential outcomes to shared representations and neural estimation

01 / Target: individual and average effects

Under unconfoundedness and overlap, the CATE is a difference of two conditional means and the ATE is its expectation. The problem becomes estimating mu_0 and mu_1 flexibly and with low variance.

tau(x)=mu_1(x)−mu_0(x);  ATE=E[tau(X)]

02 / TARNet: shared representation + two heads

Learn a shared representation Phi(x), then use two separate heads h_0, h_1 to predict the control and treated outcomes. Sharing the base while splitting the heads mitigates the variance of a T-learner on small treatment arms.

hat mu_d(x)=h_d(Phi(x)),  d∈{0,1}

03 / Counterfactual risk and balancing (CFR)

Beyond the factual loss, add a penalty that shrinks the distance between treated and control distributions in representation space (an IPM such as MMD or Wasserstein), bounding the generalization error of counterfactual extrapolation.

L=factual loss + alpha·IPM_G({Phi|D=1},{Phi|D=0})

04 / Dragonnet: add a propensity head

Attach a propensity head g(x)=P(D=1|Phi(x)) on the shared representation. This forces the representation to retain assignment-relevant information — a neural version of a sufficient propensity score.

g(x)=sigmoid(head_t(Phi(x)))

05 / Targeted regularization

Add a correction term with a perturbation parameter epsilon so the final estimate satisfies the efficient-influence-function constraint for the ATE (TMLE / AIPW style), aligning "predict well" with "estimate the ATE accurately".

tilde mu_d=hat mu_d + epsilon·(D/g − (1−D)/(1−g))

06 / DeepIV and CEVAE

Under endogenous treatment, DeepIV is two-stage: first estimate the conditional distribution F(t|x,z) of treatment given instrument and covariates, then minimize the outcome loss integrated over that distribution. CEVAE models unobserved confounders with a latent variable z and does variational inference through proxies. Both still rely on identification: instrument exclusion or proxy sufficiency.

DeepIV: min E[(Y−∫ h(t,X) dF(t|X,Z))^2]

03 / Code

Code case: a TARNet-style shared representation + two heads for CATE

A minimal PyTorch network with a shared representation and two outcome heads estimates the ATE on synthetic data and shows where the balancing penalty goes. The point is the structure, not tuning.

Case 1: how two heads produce a CATE

After the shared representation, two heads predict treated and control outcomes; their difference is the individual effect.

import numpy as np
phi = np.array([0.3, -0.2, 0.5])     # shared representation for one unit
w0, b0 = np.array([1.0, 0.5, -0.3]), 0.1   # control head
w1, b1 = np.array([1.2, 0.4, -0.1]), 0.4   # treated head
mu0 = phi @ w0 + b0
mu1 = phi @ w1 + b1
print("CATE for this unit:", round(mu1 - mu0, 3))

Expected output

CATE for this unit: 0.46

How to read this code

  • The same representation feeds both heads, giving two counterfactual predictions.
  • Their difference is the estimated treatment effect tau(x) for this unit.
  • A shared base lets control units inform the treated counterfactual.

Case 2: a representation-balancing penalty (MMD surrogate)

The larger the distance between treated and control representations, the less reliable the counterfactual extrapolation.

import numpy as np
rng = np.random.default_rng(1)
phi_t = rng.normal(0.6, 1.0, size=(500, 8))   # treated reps
phi_c = rng.normal(0.0, 1.0, size=(500, 8))   # control reps
mmd = ((phi_t.mean(0) - phi_c.mean(0)) ** 2).mean()
print("mean-MMD imbalance:", round(float(mmd), 3))

Expected output

mean-MMD imbalance: 0.291

How to read this code

  • The farther apart the two representation means, the larger the penalty.
  • CFR adds this distance to the loss, encouraging treatment-indistinguishable representations.
  • Balancing eases covariate shift but cannot create overlap from nothing.

Case 3: the Dragonnet propensity head used for debiasing

The propensity head supplies g(x), which can correct the outcome-head residual AIPW-style.

import numpy as np
m1, m0, g = 5.0, 3.0, 0.7      # treated/control predictions + propensity
Y, D = 5.4, 1
aipw = (m1 - m0) + D * (Y - m1) / g - (1 - D) * (Y - m0) / (1 - g)
print("doubly robust contribution:", round(aipw, 3))

Expected output

doubly robust contribution: 2.571

How to read this code

  • The network only supplies the three nuisances m0, m1, g.
  • The debiasing logic is still AIPW / influence functions, not the network itself.
  • This is exactly what Dragonnet plus targeted regularization aligns to.

04 / Case

Case: using annual-report text as high-dimensional confounder control

  • Question: the effect of a regulation on firm investment, where being regulated is highly correlated with fundamentals and strategic narrative (confounding) — much of which lives in annual-report text.
  • Encode the reports into vectors with a pretrained language model as high-dimensional covariates X, then use TARNet / Dragonnet to learn a shared representation and two heads for the ATE / CATE.
  • Report overlap and balancing diagnostics: whether treated and control overlap in representation space, and how estimates change before and after the balancing penalty.
  • Cross-check against DML + causal forests: agreement raises confidence; disagreement usually signals poor overlap or a questionable unconfoundedness assumption, not "a model that is not deep enough".

05 / Causal

Choosing: neural causal estimators vs DML / forests, matched to the problem

Deep-learning causal inference is not a fancier black box but more flexible nuisance / representation learners under the same identifying assumptions. Common mappings follow.

01 / High-dim / unstructured covariates (text, images) → representation learning

Use TARNet / CFR to learn a shared representation from raw inputs, then split heads for counterfactuals.

02 / Need a robust ATE → Dragonnet + targeted regularization

Add a propensity head and a TMLE-style correction on the shared representation to align to the efficient influence function.

tilde mu_d=hat mu_d+epsilon·clever covariate

03 / Endogenous treatment + instrument → DeepIV

A two-stage neural network handles endogeneity; identification still rests on instrument exclusion.

04 / Only proxies of confounders → CEVAE (with caution)

Model unobserved confounders with a latent variable, but be aware it is highly sensitive to model specification and proxy sufficiency.

Three red lines: (1) identification still comes from design — the network only swaps functional form; (2) prefer DML / forests on small samples, since neural methods need large samples, strong regularization, and careful tuning; (3) always run overlap / balance diagnostics and quantify uncertainty with cross-fitting and honest evaluation.

06 / Risks

Common Pitfalls

Treating a neural network as a tool that relaxes the identifying assumptions; it only relaxes functional form, never unconfoundedness / exclusion.
Ignoring overlap: balancing eases covariate shift, but no method can extrapolate where treated and control simply do not overlap.
Over-flexible nuisances overfit; without cross-fitting / regularization, that overfitting error enters the ATE.
Interpreting internal representations or attention as a causal mechanism.
Piling on deep models for social-science datasets of a few hundred to a few thousand rows, where DML + trees / forests is often more robust.

References