Causal machine learning
Introduction
Correlation \(\neq\) causation
We often see variables move together. But this correlation does not necessarily imply that one variable causes the other.
A third factor\(—\)called confounder\(—\)can increase/decrease both, creating a non-causal association.
Identification vs. estimation
Identification and estimation
Quantifying causal effects requires disentangling
- Identification, the question of what can be learned about causal effects, from
- Estimation, which concerns how these effects can be learned efficiently from data.
Often, units can strategically choose whether to “take the treatment.” This behavior is called self-selection and might lead to selection bias.
-
→ Identification kills such bias by leveraging different data structures and assumptions\(—\)thus ensuring valid apples-to-apples comparisons.
Once identification is deemed credible, we move to estimation. This is where ML kicks in.
-
→ Most estimation problems reduce to prediction tasks, which we want to solve flexibly.
Identification strategies
Several identification strategies exist:
- Randomized controlled trials
- Selection-on-observables
- Instrumental variables
- Regression discontinuity
- Difference-in-differences
Exercise
Exercise\(—\)Understanding DAGs (\(\approx 5\) minutes)
Discuss and answer the following questions.
- Q1. In each DAG, is the naive difference in means an unbiased ATE estimator?
- Q2. For each DAG, which variable(s) act as confounders?
- Q3. In the top DAG, do we need to control for \(Z\)?
- Q3. In the bottom DAG, should we control for \(C\)?
Solution.
- The difference in means is biased in both DAGs, as \(X\) affects both \(D\) and \(Y\).
- \(X\) is a confounder in both DAGs.
- No; \(Z\) directly affects only \(D\) and impacts \(Y\) only through \(D\), so it is not a confounder.
→ We call \(Z\) an instrument.
- No; doing so “opens the path” and creates a spurious association between \(D\) and \(Y\).
→ We call \(C\) a collider.
Exercise\(—\)Selection-on-observables trade-offs (\(\approx 5\) minutes)
Can you think of any trade-off between Assumptions 1 and 2? Think about what happens if \(\text{dim}(\boldsymbol{X}_i) \uparrow\).
Solution. As \(\text{dim}(\boldsymbol{X}_i) \uparrow\), we (may) strengthen the credibility of unconfoundedness; however, we simultaneously make positivity less plausible.
→ Curse of dimensionality: the covariate space becomes sparse, and treated and control units may no longer overlap well.
Building blocks for ATE estimation
ATE identification under selection-on-observables
Under Assumptions 1\(\text{-}\)2, the ATE can be written in three equivalent ways:
\[ \tau = \mathbb{E}[\mu(1, \boldsymbol{X}_i) - \mu(0, \boldsymbol{X}_i)]. \] \[ \tau = \mathbb{E} \left[\frac{D_i Y_i}{\pi ( \boldsymbol{X}_i)} - \frac{(1 - D_i) Y_i}{1 - \pi ( \boldsymbol{X}_i)} \right]. \] \[ \tau = \mathbb{E} \left[\mu(1, \boldsymbol{X}_i) - \mu(0, \boldsymbol{X}_i) + \frac{D_i (Y_i - \mu(1, \boldsymbol{X}_i))}{\pi ( \boldsymbol{X}_i)} - \frac{(1 - D_i) (Y_i - \mu(0, \boldsymbol{X}_i))}{1 - \pi ( \boldsymbol{X}_i)} \right]. \]
These equalities point to different estimation strategies, all relying on learning unknown nuisance functions.
→ We might need a high-dimensional \(\boldsymbol{X}_i\) to satisfy Assumption 1. Machine learning is thus attractive.
Cross-fitting
Naive plug-in ML of nuisance functions can introduce overfitting bias.
- Split data into \(I_1, \dots, I_K\) (\(\approx\) equal size).
- Let \(I_k^c = \{1, \dots, n\} \setminus I_k\) and \(k(i)\) be \(i\)’s fold.
- For each \(k\):
- Train \(\hat\mu_{-k}, \hat\pi_{-k}\) on \(I_k^c\);
- Predict on \(I_k\) using \(\hat\mu_{-k}, \hat\pi_{-k}\).
- Train \(\hat\mu_{-k}, \hat\pi_{-k}\) on \(I_k^c\);
- Stack all out-of-fold scores; average for the final estimator.
Exercise
Exercise\(—\)Suitable estimators (\(\approx 3\) minutes)
Discuss and list suitable estimators for \(\mu(\cdot, \cdot)\) and \(\pi(\cdot)\).
Solution. Notice that \(\pi(\boldsymbol{X}_i) = \mathbb{P} (D_i = 1 | \boldsymbol{X}_i)\)\(—\)the propensity score equals the “probability of success” of \(D_i\).
→ We already learnt how LPM/Logit/Probit can help us here.
→ ML provides flexible alternatives (e.g., classification trees/forests/boosting, penalized regressions).
The estimation of \(\mu(\cdot, \cdot)\) can be carried out differently according to the outcome’s nature.
→ (Penalized) linear models or regression trees/forests/boosting if \(Y_i\) is continuous.
→ Classification algorithms if \(Y_i\) is binary.
→ Ordered Logit/Probit or Ordered Correlation Forest (Di Francesco (2025)) if \(Y_i\) is ordered.
Exercise\(—\)Cross-fitting trade-offs (\(\approx 3\) minutes)
Discuss trade-offs when changing the number of folds \(K\).
Solution. As \(K\) increases, we expect more stable nuisance fits, as the training size increases. Carrying the argument to its extreme, one could implement a “leave-one-out” approach by setting \(K = n\).
However, computational time increases in \(K\), as nuisances must be fit \(K\) times; this could lead to infeasibility if \(n\) is large and we use “heavy” ML methodologies.
Toy simulation
Data-generating process
We run a simple Monte Carlo to illustrate the benefits of cross-fitting for AIPW estimation of the ATE.
- Generate \(p\) i.i.d. covariates \(X_{ij} \sim \mathcal N(0,1)\).
- Propensity score defined as \(\pi(\boldsymbol X_i) = \Lambda (0.6 X_{i1} + 0.8 \sin(X_{i3}))\).
- Treatment assignment follows a Bernoulli process \(D_i \sim \mathrm{Bernoulli}(\pi(\boldsymbol{X}_i))\).
- Control regression function is \(\mu(0, \boldsymbol{X}_i) = 2\sin(X_{i1}) + 0.5 X_{i2}^2\).
- CATE function specified as \(\tau(\boldsymbol{X}_i) = 1 + 0.5 \cdot \mathbf{1} \{X_{i1} > 0\}\).
- Generate observed outcomes as \(Y_i = \mu(0, \boldsymbol{X}_i) + D_i \tau(\boldsymbol{X}_i) + \mathcal N(0,1)\).
We construct \(\hat\tau^{AIPW}\) using regression/classification trees for nuisance estimation.
Exercise\(—\)Understanding the DPG (\(\approx\) 5 minutes)
Discuss and answer the following questions.
- Q1. Is this DGP confounded?
- Q2. Which variable(s) are confounders?
- Q3. Is there treatment-effect heterogeneity?
- Q4. What’s the model of \(\mu(1, \boldsymbol{X}_i)\)?
Solution.
- Yes.
→ \(X_{i1}\) enters the models of both \(Y_i\) and \(D_i\).
- \(X_{i1}\) is the only confounder.
→ Controlling for \(X_{i1}\) is enough for identification.
- Yes.
→ Effects are larger when \(X_{i1} > 0\).
- Notice that \(\mu(1, \boldsymbol{X}_i) = \tau(\boldsymbol{X}_i) + \mu(0, \boldsymbol{X}_i)\).
How to form groups
Data-driven groups
So far, we assumed groups were pre-specified. Pre-specifying groups is simple but risks missing unexpected heterogeneity.
Researchers often want to “let the data speak” and discover data-driven heterogeneous subgroups.
→ Tree-based approaches are popular because of their piece-wise constant structure.
However, data-driven groups are endogenous, and naive reuse of the same sample for grouping and estimation biases results.
A remedy is to combine tree-based approaches with honesty (Athey & Imbens (2016)).
- Split the data into training and honest subsamples.
- Construct a tree using training data.
- Estimate leaf-ATEs using honest data.
Causal trees (Athey & Imbens (2016))
Grow tree by adapting splitting criterion to target effect heterogeneity, and estimate GATEs by
\[ Y_i = \sum_{l = 1}^L L_{il} \gamma_l + \sum_{l = 1}^L L_{il} D_i \beta_l + \epsilon_i. \]
Aggregation trees (Di Francesco (2024))
Grow tree by applying standard CART to estimated CATEs, and estimate GATEs by \[ \hat\psi_{-k(i)}(\boldsymbol{X}_i) = \sum_{l = 1}^L L_{il} \beta_l + \epsilon_i. \]
Toy simulation\(—\)Results
CATE validation
High variance of CATE estimates \(\neq\) heterogeneity
Looking at the distribution of the estimated CATEs is not an effective strategy for assessing effect heterogeneity.
→ High variation in predictions due to estimation noise does not necessarily imply heterogeneous effects.
Emerging literature on CATE validation
There is an emerging literature on CATE validation that introduces methodologies to quantify how much signal a given CATE estimator is actually capturing.
- BLP, GATES, and CLAN analyses based on repeated sample splitting (Chernozhukov et al. (2017)).
- GATES analysis under Neyman’s repeated-sampling framework (Imai & Li (2025)).
- RATEs (Yadlowsky et al. (2025)).