doubly robust | Chen Xing

Notes on Doubly Robust Censoring Unbiased Transformation

Tue, 17 Mar 2026 00:00:00 +0000

Predicting outcomes with right-censored survival data forces a choice: do we model the outcome distribution, or the censoring mechanism? Classical transformations require us to commit to one. Rubin and van der Laan (2007) tell us we don’t have to. Their doubly robust censoring unbiased transformation fuses both approaches, remaining valid as long as at least one of the two nuisance models is correctly specified. This post walks through the setup, the classical transformations, and how the doubly robust version combines them — drawing the analogy to AIPW along the way.

1. Setup

We observe an i.i.d. sample $\{O_i\}_{i=1}^n$ where each observation is

$$O = \bigl(W,\; \Delta = \mathbf{1}(Y \le C),\; \tilde{Y} = Y \wedge C\bigr)$$

$W$: covariates
$Y$: true (possibly unobserved) survival time
$C$: random censoring time
$\tilde{Y} = \min(Y, C)$ is what we actually record
$\Delta = 1$ if the event is observed ($Y \le C$), and $\Delta = 0$ if the outcome is right-censored ($Y > C$)

Our goal is to estimate the regression function

$$m(w) = \mathbb{E}[Y \mid W = w]$$

Two nuisance functions appear throughout:

$\bar{F}(\cdot \mid W)$: conditional survival function of the response $Y$ given $W$
$\bar{G}(\cdot \mid W)$: conditional survival function of the censoring time $C$ given $W$

We maintain the standard assumption $Y \indep C \mid W$.

2. The Challenge: Unidentifiability

When the censoring time $C$ corresponds to a fixed study endpoint, the true response $Y$ may exceed it. Beyond that endpoint, nothing can be learned about the tail of the survival distribution — making $m(W) = \mathbb{E}[Y \mid W]$ unidentifiable in general.

The fix: truncate the response at a known horizon $\tau$,

$$Y \longmapsto Y \wedge \tau = \min(Y, \tau)$$

and estimate $w \mapsto \mathbb{E}[Y \wedge \tau \mid W = w]$ instead.

The Surrogate-Response Strategy

The general approach to prediction with right-censored data is:

Replace the possibly unavailable responses $\{Y_i\}_{i=1}^n$ with surrogate values $\{Y^*(O_i)\}_{i=1}^n$ using an imputation map $Y^*(\cdot)$ built from observed data.
Plug the imputed dataset $\{W_i, Y^*(O_i)\}_{i=1}^n$ into any standard regression algorithm.

The imputation map $Y^*(\cdot)$ is called a censoring unbiased transformation (Fan and Gijbels 1996) if it satisfies

$$\mathbb{E}[Y^*(O) \mid W] = \mathbb{E}[Y \mid W] = m(W)$$

That is, the surrogate is an unbiased proxy for the true (latent) response, conditional on covariates.

3. Two Classical Transformations

A. The Buckley–James Transformation (depends on $\bar{F}$)

The Buckley–James transformation imputes a censored observation with its conditional mean given that it exceeds the censoring time:

$$Y^*(O) = \Delta Y + (1-\Delta) Q_{\bar{F}}(W, C)$$

where $\Delta = 1$ if $Y$ is observed and $\Delta = 0$ if right-censored, and

$$Q_{\bar{F}}(w, y) = \mathbb{E}[Y \mid W = w,\; Y > y] = \frac{1}{\bar{F}(y \mid W=w)} \int_y^{+\infty} u \; dF(u \mid W=w)$$

Intuitively: if we observe the event, we keep $Y$; if censored at $C$, we impute with the expected remaining survival time above $C$.

This transformation requires correctly estimating $\bar{F}(\cdot \mid W)$, i.e., the conditional distribution of the survival time.

B. The IPCW Transformation (depends on $\bar{G}$)

Inverse probability of censoring weighting (IPCW) up-weights the observed events to compensate for the censored ones:

$$Y^*(O) = \frac{Y \Delta}{\bar{G}(Y \mid W)}$$

This requires correctly estimating $\bar{G}(\cdot \mid W)$, the conditional survival function of the censoring time. It is the survival-analysis analogue of IPW in the causal inference literature.

4. The Doubly Robust Censoring Unbiased Transformation

The two classical transformations each stake everything on one nuisance model. The doubly robust approach combines them:

$$ Y^*(O) = \underbrace{\frac{Y\Delta}{\bar{G}(Y \mid W)}}_{\text{1st term}} + \underbrace{\frac{Q_{\bar{F}}(W,C)\,(1-\Delta)}{\bar{G}(C \mid W)}}_{\text{2nd term}} - \underbrace{\int_{-\infty}^{\tilde{Y}} \frac{Q_{\bar{F}}(W,c)}{\bar{G}^2(c \mid W)}\, dG(c \mid W)}_{\text{3rd term (correction)}} $$

where $\tilde{Y} = Y \wedge C = \min(Y,C)$ and $Q_{\bar{F}}(w, c) = \mathbb{E}[Y \mid W=w,; Y > c]$.

Theorem 1.

$\mathbb{E}[Y^*(O) \mid W] = \mathbb{E}[Y \mid W]$ whenever either $\bar{F}(\cdot \mid W)$ or $\bar{G}(\cdot \mid W)$ is correctly specified.

5. Intuition: An AIPW in Disguise

The three-term structure has a clean interpretation. Recall that $Q_{\bar{F}}(W,C) = \mathbb{E}[Y \mid W, Y > C]$ is the outcome regression for censored individuals.

Term	Role
1st: $Y\Delta \,/\, \bar{G}(Y\!\mid\! W)$	For observed outcomes — apply IPCW, the analogue of inverse propensity score weighting
2nd: $Q_{\bar{F}}(W,C)(1-\Delta) \,/\, \bar{G}(C\!\mid\! W)$	For censored outcomes — impute with the outcome regression $Q_{\bar{F}}$, then apply an IPCW-style weight
3rd: $-\int Q_{\bar{F}} / \bar{G}^2 \; dG$	Bias correction term

The first two terms together look exactly like IPW + imputation — the two ingredients of AIPW in the standard (binary treatment) setting. The third term is the survival-analysis counterpart of the augmentation correction in AIPW: it removes the bias that accumulates when both nuisance models are slightly off.

Takeaway

The doubly robust censoring unbiased transformation is a drop-in replacement for the classical Buckley–James or IPCW transformations. Once $Y^*(O_i)$ is computed for each observation, any off-the-shelf regression algorithm can be applied to the pairs $\{W_i, Y^*(O_i)\}$.

Reference

Rubin, Daniel and Mark J. van der Laan (2007), “A Doubly Robust Censoring Unbiased Transformation,” The International Journal of Biostatistics, 3 (1).

Fan, Jianqing and Irène Gijbels (1996), Local Polynomial Modelling and Its Applications, Chapman & Hall.

Calibrate Your Nuisances: A Simple Fix for Doubly Robust Inference

Fri, 09 Jan 2026 00:00:00 +0000

Doubly Robust Inference via Calibration

TL;DR

Van der Laan, Luedtke, and Carone introduce “calibrated debiased machine learning” (calibrated DML), a method that achieves doubly robust asymptotic normality for causal inference estimators by simply adding an isotonic regression calibration step to standard DML pipelines. The key innovation is that valid inference requires only one of two nuisance functions (outcome regression or propensity score) to be estimated well—the other can converge arbitrarily slowly or even inconsistently. This bridges a long-standing gap where consistency was doubly robust but inference was not, and the method can be implemented by adding just a few lines of code to existing workflows.

What is this paper about?

Doubly robust estimators like AIPW are popular for estimating average treatment effects because they remain consistent if either the outcome regression or propensity score is correctly specified. However, there’s a crucial asymmetry: while consistency requires only one nuisance to be correct, valid inference (asymptotic normality, correct confidence intervals) typically requires both nuisances to converge at sufficiently fast rates $(\approx n^{-1/4})$. When one nuisance is estimated poorly—whether due to model misspecification or slow convergence—standard confidence intervals can have incorrect coverage. Prior solutions like the DR-TMLE framework required computationally intensive iterative procedures and case-by-case derivations for each new parameter.

Motivation

The authors ask: can we achieve doubly robust inference using a simple, general-purpose procedure that works with any machine learning estimator?

What do the authors do?

The authors develop calibrated DML, a two-step procedure that:

step 1: takes cross-fitted nuisance estimators from any machine learning algorithm
step 2: calibrates them using isotonic regression before constructing the debiased estimator

The calibration step ensures that nuisance estimates satisfy certain empirical orthogonality conditions that linearize the bias term in the doubly robust expansion. Specifically, they calibrate the outcome regression using squared error loss and the Riesz representer (inverse propensity weights for ATE) using a tailored Riesz loss.

The paper proves that this calibration is sufficient to achieve “doubly robust asymptotic linearity” (DRAL)—meaning the estimator is asymptotically normal whenever at least one nuisance converges at $n^{-1/4}$ rate, even if the other is inconsistent. They also develop bootstrap-assisted confidence intervals that avoid estimating additional nuisance functions. Empirically, they evaluate on simulated data with deliberately misspecified nuisances and on semi-synthetic benchmarks (ACIC, IHDP, Twins, LaLonde), comparing against standard AIPW and the iterative DR-TMLE.

Why is this important?

Applied researchers using flexible ML methods (random forests, neural networks, gradient boosting) for nuisance estimation face a dilemma: these methods provide consistency but may not converge fast enough to guarantee valid inference, especially in moderate dimensions. The standard “product rate” condition requiring both nuisances to converge at $n^{-1/4}$ can fail when even one nuisance is complex or misspecified. Calibrated DML transforms this multiplicative requirement into a requirement on just the better-estimated nuisance.

The practical benefits are substantial: better finite-sample coverage (e.g., improving from 32% to 90% coverage in ACIC-2017 simulations), reduced bias, and a method that integrates into existing pipelines with minimal code changes.

The theoretical contribution is equally significant—it establishes a novel connection between prediction calibration and causal inference validity, showing that calibration of nuisances provides the “debiasing” needed for doubly robust inference.

Who should care?

Applied researchers estimating treatment effects with observational data using ML for nuisance estimation will find immediate practical value—this method provides insurance against nuisance misspecification without computational overhead.

Econometricians and biostatisticians working on semiparametric inference will appreciate the theoretical framework connecting calibration to doubly robust properties.

Methodologists developing new estimands can use this as a general recipe: the approach applies to any linear functional of the outcome regression (counterfactual means, ATE, partial covariance, survival outcomes under missingness), not just the ATE.

Researchers in policy evaluation, epidemiology, marketing, and tech who regularly use AIPW-style estimators should consider adopting calibrated DML as a default given its robustness benefits.

Do we have code?

Yes, the authors provide both R and Python implementations.

An R package calibratedDML and Python code are available on GitHub at https://github.com/Larsvanderlaan/calibratedDML. The paper includes complete code listings for calibrating inverse propensity weights and outcome regressions using xgboost with monotonicity constraints.

The implementation is straightforward: isotonic regression is performed using gradient-boosted trees with monotone_constraints=1 and a single boosting round, making it computationally efficient and easy to integrate into existing DML workflows.

In summary, this paper provides a remarkably practical solution to a longstanding theoretical problem. The insight that isotonic calibration—a standard tool from prediction—can unlock doubly robust inference is both elegant and immediately actionable. For anyone running AIPW or related estimators with ML nuisances, calibrating cross-fitted estimates before debiasing is now the obvious default: it costs almost nothing computationally and provides genuine protection against the scenario where one of your nuisance models is less reliable than you hoped.

Reference

van der Laan, Lars, Alex Luedtke, and Marco Carone (2024), “Doubly robust inference via calibration,” arXiv preprint arXiv:2411.02771.

An R package calibratedDML and Python code are available on GitHub at https://github.com/Larsvanderlaan/calibratedDML.

Balancing Weights for Causal Inference

Wed, 22 Oct 2025 00:00:00 +0000

TL;DR

Cohn et al. (2023) introduces the balancing approach to weighting for causal inference in observational studies. Unlike traditional methods that model the propensity score directly, balancing weights are estimated by solving an optimization problem that directly targets covariate balance between treatment groups. The authors demonstrate that this approach offers protection against model misspecification, connects naturally to bias-variance trade-offs, and can be augmented with outcome modeling for improved performance. Applied to the classic LaLonde job training data, balancing methods achieve better covariate balance than standard propensity score approaches while maintaining reasonable effective sample sizes.

What is this paper about?

Covariate balance is fundamental to causal inference: randomized experiments achieve it by design, while observational studies must adjust for it. This chapter addresses a key challenge in observational causal inference—how to construct weights that remove confounding by balancing observed covariates between treated and control groups. The traditional modeling approach estimates propensity scores (the probability of treatment given covariates) and inverts them to create weights, but this relies heavily on correct model specification. When the propensity score model is wrong, the resulting weights may fail to balance covariates in the sample, leading to biased treatment effect estimates. The chapter explores an alternative: directly finding weights that achieve balance in the observed data, rather than first modeling the propensity score.

What do the authors do?

The authors formalize the balancing approach as an optimization problem that jointly minimizes covariate imbalance and weight dispersion (variance). They show how different choices of the “model class” M—the set of functions of covariates to balance—correspond to different assumptions about the outcome model and lead to different optimization formulations. Using the LaLonde dataset (a constructed observational study where the true treatment effect is known), they compare three designs: balancing main covariate terms only, balancing up to three-way interactions, and balancing an infinite-dimensional reproducing kernel Hilbert space (RKHS). For each design, they evaluate covariate balance using standardized mean differences, examine the effective sample size (a measure of weight dispersion), and explore the bias-variance trade-off by varying regularization parameters. The authors also demonstrate how balancing weights can be augmented with outcome regression to further reduce bias, and they establish asymptotic normality results for inference.

Why is this important?

This work matters because most observational studies include covariates in their analysis, yet practitioners often don’t carefully consider whether their weighting method actually achieves balance on the relevant covariate functions. The balancing approach makes covariate balance a first-order design criterion rather than a post-hoc diagnostic check. It reveals the implicit bias-variance trade-offs in weighting methods and shows that different assumptions about the outcome model (linear, interactive, nonparametric) lead to fundamentally different weighting strategies. The framework unifies many existing methods (entropy balancing, stable weights, kernel balancing) under one optimization structure and clarifies the connection between the balancing and modeling approaches through Lagrangian duality. Importantly, the chapter provides practical guidance on design choices—what to balance, how much dispersion to tolerate, whether to allow negative weights—that applied researchers face but often lack principled ways to resolve.

Who should care?

Applied researchers in economics, epidemiology, public policy, education, and medicine who use inverse propensity weighting or other covariate adjustment methods in observational studies. Methodologists working on causal inference, especially those developing new weighting estimators or studying properties of existing ones. Graduate students learning causal inference who need to understand the trade-offs between different adjustment strategies and the implicit assumptions behind common practices. Policy evaluators who must justify their modeling choices and demonstrate that their treatment effect estimates are robust to covariate imbalance. Anyone who has struggled with poor covariate balance after propensity score weighting or wondered how to choose between competing adjustment methods would benefit from this framework.

Do we have code?

The chapter does not provide standalone replication code or software packages. However, the authors note that many of the specific balancing methods discussed are available in existing R packages: entropy balancing in the ebal package, covariate balancing propensity scores (CBPS) in the CBPS package, and stable balancing weights in the sbw package. The kernel balancing approach can be implemented using standard kernel methods in R or Python. The LaLonde dataset used throughout the examples is publicly available and widely used in the causal inference literature, making it straightforward to reproduce the analyses with these existing tools.

In summary, this chapter reframes propensity score weighting as a balance-optimization problem rather than a pure modeling exercise. By directly targeting the balancing property of inverse propensity weights, the approach offers robustness to propensity score misspecification while making explicit the bias-variance trade-offs inherent in any covariate adjustment. The LaLonde application demonstrates that balancing weights can substantially reduce covariate imbalance compared to standard methods, though at the cost of reduced effective sample size. The framework provides both theoretical insight (connecting balancing to dual regression, establishing asymptotic properties) and practical guidance (how to choose what to balance, when to augment with outcome modeling, whether to allow extrapolation) that fills an important gap in applied causal inference.

Reference

Cohn, Eric R., Eli Ben-Michael, Avi Feller, and José R. Zubizarreta (2023), “Balancing Weights for Causal Inference,” in Handbook of Matching and Weighting Adjustments for Causal Inference, Chapman and Hall/CRC. https://www.taylorfrancis.com/chapters/edit/10.1201/9781003102670-16/balancing-weights-causal-inference-eric-cohn-eli-ben-michael-avi-feller-josé-zubizarreta

Triply Robust Panel Estimators 🛡: When You Don't Know Which Assumptions Hold

Mon, 20 Oct 2025 00:00:00 +0000

TL;DR

Athey, Imbens, Qu, and Viviano introduce the Triply RObust Panel (TROP) estimator, which combines unit weights, time weights, and a flexible low-rank factor model to estimate causal effects in panel data. The estimator achieves triple robustness: it remains unbiased if any one of three components succeeds (unit balance, time balance, or correct outcome modeling). Across 21 empirically calibrated simulation designs, TROP outperforms standard methods like DID, synthetic control, matrix completion, and synthetic difference-in-differences in 20 out of 21 cases, often by substantial margins. The method is particularly valuable when interactive fixed effects violate parallel trends assumptions and when researchers face uncertainty about which modeling assumptions hold.

What is this paper about?

This paper addresses a fundamental challenge in panel data analysis: choosing among competing estimators when the validity of their underlying assumptions is uncertain. Traditional difference-in-differences (DID) methods assume parallel trends, synthetic control (SC) emphasizes pre-treatment outcome alignment, and matrix completion (MC) relies on low-rank factor structures, but researchers rarely know which assumptions hold in practice. A second problem is that existing methods often treat observations uniformly across time and units, ignoring that recent periods may be more informative for predicting counterfactuals than distant ones, and that some control units may be more comparable to treated units than others. Third, most SC-based methods are designed for a single treated unit and do not naturally extend to complex assignment patterns with multiple treated units and staggered adoption. The authors propose a unified framework that addresses all three issues simultaneously while providing formal robustness guarantees.

What do the authors do?

The authors develop the TROP estimator, which minimizes a weighted regression objective that incorporates:

unit-specific weights (to upweight control units similar to treated units)
time-specific weights (to upweight recent periods)
a nuclear-norm-penalized low-rank factor model (to capture interactive fixed effects)

They propose an exponential decay structure for time weights and distance-based weights for units, with all tuning parameters selected via leave-one-out cross-validation.

To evaluate performance, they conduct extensive semi-synthetic simulations calibrated to seven real datasets (CPS wage data, Penn World Table GDP, Germany reunification, Basque country, California smoking, and Mariel boatlift), generating outcomes from rank-4 factor models with realistic autocorrelation and treatment assignments based on estimated propensity scores from actual policies. They systematically vary features of the data-generating process (removing autocorrelation, interactive effects, or fixed effects) and components of the estimator (shutting down unit weights, time weights, or the regression adjustment) to understand what drives performance differences. They provide formal theory showing that TROP’s bias equals the product of (1) unit imbalance, (2) time imbalance, and (3) regression model misspecification, establishing triple robustness.

Figure 1: Theorem 5.1 in the paper

Why is this important?

The paper makes three critical contributions to applied panel data analysis. First, it demonstrates that interactive fixed effects—where unit-specific trends vary over time—are empirically pervasive and cause substantial bias in standard DID estimators, with violations present in over 58% of units across their applications and adding even a single interactive factor reducing RMSE by 10-60%. Second, it shows that no single existing method dominates: DID can have RMSE 900% higher than the best competitor in some settings, SC can be 580% worse, and even modern methods like SDID can be 90% worse, highlighting the risk of mechanically applying any one approach. Third, TROP provides a principled solution by learning from the data which components (unit balance, time balance, or factor modeling) matter most for each application, offering insurance against model misspecification without requiring researchers to know ex ante which assumptions hold. The triple robustness property means that getting any one of three components approximately right suffices for consistency, substantially weakening the conditions needed for valid causal inference.

Who should care?

Applied researchers using panel data methods to evaluate policies, interventions, or natural experiments should care deeply about this paper, especially those working with observational data where parallel trends may fail due to heterogeneous trends across units. Economists studying labor markets, public finance, development, or industrial organization who face settings with staggered treatment adoption or multiple treated units will benefit from the method’s flexibility. Methodologists interested in robust causal inference, doubly robust estimation, or combining machine learning with causal identification will find the theoretical framework valuable. Policy analysts and data scientists in government agencies producing impact evaluations need methods that perform well under uncertainty about modeling assumptions. Anyone who has abandoned a difference-in-differences project because pre-trends failed should reconsider whether allowing for interactive fixed effects and flexible weighting could recover valid identification.

Do we have code?

The paper does not provide replication code, R or Python packages, or links to implementation files. The method requires implementing nuclear-norm-penalized regression combined with weighted least squares and a cross-validation procedure to select three tuning parameters (λ_unit, λ_time, λ_nn), which is computationally involved but feasible for researchers comfortable with optimization. The algorithms are described in detail (Algorithms 1-3 in the paper), and the exponential decay weighting scheme and distance metrics are explicitly specified in equations (3) and surrounding text. Researchers would need to build their own implementation or wait for the authors to release software, though the paper’s clarity makes this more tractable than usual.

In summary, this paper introduces a methodologically sophisticated yet practically motivated estimator that combines the strengths of synthetic control (unit weighting), modern DID approaches (time weighting), and matrix completion (low-rank modeling) while avoiding the brittleness of relying on any single component. The extensive simulations demonstrate that TROP consistently outperforms existing methods across diverse empirically relevant settings, with the advantage most pronounced when interactive fixed effects are present—a common feature of real data that violates standard parallel trends assumptions. The triple robustness property provides formal theoretical insurance: if researchers get unit balance approximately right, or time balance approximately right, or the factor model approximately right, the estimator remains consistent. For applied work, this represents a major advance in practical robustness, moving beyond the “which method should I use?” question to “let the data tell us which features matter most." The lack of available code is the main practical limitation, but the clear exposition of algorithms makes implementation feasible for quantitatively oriented researchers.

Points I enjoyed

TROP vs. SDID

Let’s compare the TROP vs. SDID:

Triple Robustness (TROP): TROP achieves triple robustness, meaning its asymptotic bias vanishes if either the unit weights, or the time weights, or the flexible low-rank regression adjustment successfully remove the underlying biases. The bias of TROP depends on the product of unit imbalance, time imbalance, and the bias arising from a potentially misspecified regression adjustment.

Figure 2: Key difference between TROP and SDID – adding the low-rank factor component in the outcome model
Double Robustness (SDID): SDID (Arkhangelsky et al. 2021) achieves only double robustness, where bias vanishes if either unit imbalance or time imbalance is negligible. It lacks the third robust channel provided by the flexible regression adjustment.

Figure 3: In SDID, there is no low-rank component $L_{it}$

How to come up with this TROP estimator? My Thought

In my view, the idea behind TROP is to marry the SDID framework with the matrix completion estimator

Recall that:

Matrix completion methods (Athey et al. 2021) focus solely on modeling the outcomes as a function of the latent factors. These methods aim to accurately model the control outcome process using a low-rank factor structure that generalizes the two-way fixed effects (TWFE) model
SDID is a “local” TWFE estimator, which essentially corresponds to modeling the assignment mechanism, via local weighting (i.e. unit and time weights)

The estimated average counterfactual from the TROP estimator has the following form:

Figure 4: The estimated average counterfactual from the TROP estimator. I think the key intuition here is the 'bias correction'

It contains both unit/time weights (the key idea in SDID) and outcome model (low-rank factor model) in matrix completion methods. The TROP estimator, $\hat{\tau}^{\text{TROP}}$, applies unit and time weights, $\omega_i$ and $\theta_t$ respectively, to correct the bias from the outcome model $(\mathbf{L})$.

This form also reminds me the idea in AIPW estimator or the generic debiased framework (Chernozhukov, Newey, and Singh 2022) – leveraging the Riesz representer to correct for bias.

Figure 5: screen shot from my previous post about AutoDML - using Riesz Representer to correct for bias

Check more details from my previous post: “Intuition for Doubly Robust Estimator”.

Reference

Athey, Susan, Guido Imbens, Zhaonan Qu, and Davide Viviano (2025), “Triply robust panel estimators.” https://doi.org/10.48550/arXiv.2508.21536

Athey, Susan, Mohsen Bayati, Nikolay Doudchenko, Guido Imbens, and Khashayar Khosravi (2021), “Matrix Completion Methods for Causal Panel Data Models,” Journal of the American Statistical Association, 116 (536), 1716–30.

Arkhangelsky, Dmitry, Susan Athey, David A. Hirshberg, Guido W. Imbens, and Stefan Wager (2021), “Synthetic Difference-in-Differences,” American Economic Review, 111 (12), 4088–4118.

Chernozhukov, Victor, Whitney K. Newey, and Rahul Singh (2022), “Automatic Debiased Machine Learning of Causal and Structural Effects,” Econometrica, 90 (3), 967–1027.

AIPW vs. Residual-on-Residual regression: Non-Parametric Flexibility or Efficiency?

Wed, 11 Jun 2025 00:00:00 +0000

Introduction

Two powerful tools in causal inference are the Augmented Inverse Propensity Weighting (AIPW) estimator and the Residual-on-Residual regression estimator for partially linear models. Drawing from Wager’s notes (2024), this post breaks down how these estimators work, compares their strengths and weaknesses, and offers tips for when to use each.

Residual-on-Residual regression

The semiparamteric partially linear model (PLM) assumes the outcome $ Y $ can be written as:

$$ Y = \theta D + g(X) + \epsilon \tag{1}$$

Here, $ D $ is a binary treatment, $ \theta $ is the causal parameter of interest, $ g(X) $ is some unknown function of covariates $ X $, and $ \epsilon $ is random noise with $\E[\epsilon \mid D, X] = 0$. The residual-on-residual estimator (or Robinson (1988) estimator) isolates $ \theta $ in two steps:

Partialing out covariates: “Partialing out” $ X $ from $D$ and $Y$ using nonparametric regression and cross-fitting: $ \tilde{D} = D - \hat{\E}[D \mid X] $ and $ \tilde{Y} = Y - \hat{\E}[Y \mid X] $.
Estimate the effect: Use linear regression $ \tilde{Y} \sim \tilde{D} $ to get $ \hat{\theta}$.

However, the model is not fully general, because it imposes a parametric specification on the key component of interest. It imposes additivity in $g(X)$ and $D$.

Augmented Inverse Propensity Weighting (AIPW) Estimator

AIPW takes a fully non-parametric approach, aiming to estimate ATE, $ \tau = E[Y(1) - Y(0)] $, where $ Y(1) $ and $ Y(0) $ are potential outcomes under treatment and control.

$$ \hat{\tau}_{\text{AIPW}} = \frac{1}{n} \sum_{i=1}^n \left[ \hat{\mu}_1(X_i) - \hat{\mu}_0(X_i) + D_i \frac{Y_i - \hat{\mu}_1(X_i)}{\hat{e}(X_i)} - (1 - D_i) \frac{Y_i - \hat{\mu}_0(X_i)}{1 - \hat{e}(X_i)} \right] $$

Here, $ \hat{\mu}_1(X) $ and $ \hat{\mu}_0(X) $ are outcome regression estimators for treated and untreated units, and $ \hat{e}(X) $ is the estimator of propensity score. AIPW combines outcome modeling with inverse propensity weighting, making it doubly robust: it’s consistent if either the outcome model or propensity score is consistent.

The Key Difference: Non-Parametric vs. Partially Linear

AIPW is fully non-parametric, imposing no specific parametric form on the treatment effect, while residual-on-residual regression estimator assumes a partially linear structure

AIPW’s flexibility:
- AIPW does not assume a specific parametric form!
- AIPW is efficient in the generic non-parametric setting.
Residual-on-residual regression’s structure: What partially linear assumption buys us is that residual-on-residual estimators that exploit this constraint can have smaller variance than AIPW. In other words, adding this additional structure makes the residual-on-residual estimator more efficient than AIPW.

A risk of using the residual-on-residual estimator is that constant treatment effect model (1) may be misspecified.

Why This Matters

The choice between AIPW and residual-on-residual regression reflects a deeper trade-off in causal inference: flexibility versus efficiency. AIPW’s non-parametric nature makes it a Swiss Army knife for complex data, while PLM structure is like a precision tool—effective when conditions are right.

As Wager’s notes highlight:

Both AIPW and residual-on-residual regression are Neyman-orthogonal, making them robust to first‐stage errors.
However, their assumptions shape their performance. AIPW attains the lowest possible asymptotic variance for ATE under unconfoundedness. The residual-on-residual estimator, by imposing extra structure, can go beyond that bound when its structure is correct but at the cost of vulnerability to misspecification.

Conclusion

AIPW’s fully non-parametric approach offers robustness and flexibility, while residual-on-residual regression’s partially linear structure prioritizes efficiency when assumptions hold.

Reference

Wager, S. (2024). Causal inference: A statistical learning approach. https://web.stanford.edu/~swager/causal_inf_book.pdf

Notes on Propensity Score Methods

Tue, 03 Jun 2025 00:00:00 +0000

Introduction

Here are my notes on propensity scores, mainly from Prof. Ding’s textbook (2024).

The traditional propensity score analysis workflow is shown in the image below, which I will not cover in detail. Instead, I will summarize the key theorems and results from Ding’s textbook.

Figure 1: Traditional propensity score analysis workflow

I will also provide some connections with Riesz Representer (RR).

Why connect with the Riesz Representer (RR)? The connection provides a powerful generalization of the foundational Rosenbaum-Rubin (1983) result.
Rosenbaum and Rubin showed that conditioning on the propensity score is sufficient for removing confounding bias when estimating causal effects. The Riesz representer extends this principle: it suffices to regress on the Riesz representer to obtain unbiased estimates of the average treatment effect.
The key insight is that the Riesz representer, like the propensity score, serves as a sufficient statistic – it captures all the confounding information necessary for unbiased estimation of your target causal parameter.

Setting & Notation

Binary treatment $Z$
Potential outcomes $\{Y(0), Y(1)\}$
Propensity score: $\P(Z = 1 \mid X)$, where $X$ represents covariates

Two approaches learning causal relationships:

Outcome process (via outcome regression)
Treatment assignment mechanism (via propensity score)

The following summarizes the key theorems and results related to propensity scores from Prof. Ding’s textbook.

1. The propensity score as a dimension reduction tool

Theorem 1 (pscore as dimension reduction tool).

$$ \text { If } Z \indep \{Y(1), Y(0)\} \mid X, \text { then } Z \indep \{Y(1), Y(0)\} \mid e(X) . $$

Covariates $X$ can be high dimensional, but the propensity score, $e(X) \in \R$, is a 1-dimensional scalar
We can view the propensity score as a dimensional reduction tool

2. Propensity score stratification

Idea: Discretize the estimated propensity score by its $K$ quantiles:
$$ Z \indep \{Y(1), Y(0)\} \mid \hat{e}^{\prime}(X)=e_k \quad(k=1, \ldots, K) . $$
Estimate ATE within each subclass and then average by the block size
Advantage: The propensity score stratification estimator only requires the correct ordering of the estimated propensity scores rather than their exact values, which makes it relatively robust compared with other methods

3. Propensity score weighting

Theorem 2 (Invese propensity score weighting (IPW)).

If $Z \indep \{Y(1), Y(0)\} \mid X$ and $ 0 < e(X) < 1$, then $$E\{Y(1)\}=E\left\{\frac{Z Y}{e(X)}\right\}, \quad E\{Y(0)\}=E\left\{\frac{(1-Z) Y}{1-e(X)}\right\}$$ and $$ \begin{aligned} \tau &=E\{Y(1)-Y(0)\}\\ &=E\left\{\frac{Z Y}{e(X)}-\frac{(1-Z) Y}{1-e(X)}\right\} \\ &=E\left\{HY \right\} \end{aligned} $$ where $H := \left[\frac{Z}{e(X)}-\frac{(1-Z) }{1-e(X)}\right]$ is called the Horvitz-Thompson transform.

Connection the Riesz Representer (RR)

Remark 1 (RR in the case of ATE).

In the case of ATE, the Riesz Representer, $\alpha(Z, X)$, has the same form as above Horvitz-Thompson transform, $$ \alpha(Z, X) = \left[\frac{Z}{e(X)}-\frac{(1-Z) }{1-e(X)}\right] $$

3.1 Estimation

The sample version of IPW is called the Horvitz–Thompson (HT) estimator,
$$ \hat{\tau}^{\mathrm{ht}}=\frac{1}{n} \sum_{i=1}^n \frac{Z_i Y_i}{\hat{e}\left(X_i\right)}-\frac{1}{n} \sum_{i=1}^n \frac{\left(1-Z_i\right) Y_i}{1-\hat{e}\left(X_i\right)} $$
HT estimator $\hat{\tau}^{\mathrm{ht}}$ has many problems
Problem: lack of invariance, i.e. if we replace $Y_i$ by $Y_i + c$, $\hat{\tau}^{\mathrm{ht}}$ changed because it depends on $c$. This is not reasonable.
Solution: normalizing the weights
$$ \hat{\tau}^{\text {hajek }}=\frac{\sum_{i=1}^n \frac{Z_i Y_i}{\hat{e}\left(X_i\right)}}{\sum_{i=1}^n \frac{Z_i}{\hat{e}\left(X_i\right)}}-\frac{\sum_{i=1}^n \frac{\left(1-Z_i\right) Y_i}{1-\hat{e}\left(X_i\right)}}{\sum_{i=1}^n \frac{1-Z_i}{1-\hat{e}\left(X_i\right)}} . $$
Hajek estimator is invariant to the location transformation

3.2 Strong overlap condition

Many asymptotic analyses require a strong overlap condition,

$$ 0<\alpha_{\mathrm{L}} \leq e(X) \leq \alpha_{\mathrm{U}}<1 $$ In practice,

Crump et al. (2009) suggested $α_L = 0.1$ and $α_U = 0.9$
Kurth et al. (2005) suggested $α_L = 0.05$ and $α_U = 0.95$

4. Balancing property

Theorem 3 (balancing property).

Conditional on $e(X)$, the treatment and the covariates are independent
Within the same level of the propensity score, the covariate distributions are balanced across the treatment and control groups
Useful implication: we can check whether the propensity score model is specified well enough to ensure the covariate balance in the data

4.1 Propensity score is a balancing score

Definition 1.

Theorem 4 (Propensity score is a balancing score).

This is relevant in subgroup analysis
The conditional independence in (11.5) ensures unconfoundedness holds given the propensity score, within each level of $X_1$. Therefore, we can perform the same analysis based on the propensity score, within each level of $X_1$, yielding estimates for two subgroup effects

5. Doubly Robust or AIPW

The following Theorem is summarized from Prof. Wager’s lecture notes (2024).

Theorem 5 (strong double robustness of AIPW estimator).

Define the outcome regression as $$ \mu_{(z)}(x)=\mathbb{E}\left[Y_i(z) \mid X_i=x\right], $$ Define AIPW estimator as $$ \begin{aligned} \hat{\tau}_{A I P W} & =\underbrace{\frac{1}{n} \sum_{i=1}\left(\hat{\mu}_{(1)}\left(X_i\right)-\hat{\mu}_{(0)}\left(X_i\right)\right)}_{\text { outcome regression estimator }} \\ & +\underbrace{\frac{1}{n} \sum_{i=1}^n\left(\frac{Z_i}{\hat{e}\left(X_i\right)}\left(Y_i-\hat{\mu}_{(1)}\left(X_i\right)\right)-\frac{1-Z_i}{1-\hat{e}\left(X_i\right)}\left(Y_i-\hat{\mu}_{(0)}\left(X_i\right)\right)\right)}_{\text { applying IPW to the regression residuals }} \end{aligned} $$ If we use estimators $\hat{\mu}_{(z)}(x)$ and $\hat{e}(x)$ that are both consistent with root-mean squared error (RMSE) decaying faster than $n^{-\alpha_\mu}$ and $n^{-\alpha_e}$ respectively, and if furthermore $\alpha_\mu+\alpha_e \geq 1 / 2$, then $$ \begin{aligned} & \sqrt{n}\left(\hat{\tau}_{A I P W}-\tau\right) \Rightarrow \mathcal{N}\left(0, V_{A I P W}\right) \\ & V_{A I P W}=\operatorname{Var}\left[\tau\left(X_i\right)\right]+\mathbb{E}\left[\frac{\sigma_0^2\left(X_i\right)}{1-e\left(X_i\right)}\right]+\mathbb{E}\left[\frac{\sigma_1^2\left(X_i\right)}{e\left(X_i\right)}\right] \end{aligned} $$ where $$ \sigma_{(z)}^2(x)=\operatorname{Var}\left[Y_i(z) \mid X_i=x\right] $$

Check my previous post: Intuition for Doubly Robust Estimator
AIPW provides a natural starting point for understanding Double Machine Learning
Key insight of RR in DML framework: Leverage the Riesz Representer, a “generalized version of propensity score” to “correct the bias”

More general, Li et al. (2018a) gave a unified discussion of the causal estimands in observational studies.

Theorem 6 (Ding (2024), Section 13.4).

Summary Table of common estimands:

This table provides us a good way to understand and remember IPW estimator for ATT
How to remember $\tau^h$? Apply IPW on “pseudo outcome” $Yh(X)$ then divide by $E(h(X))$
When the parameter of interest is ATT, then $$E(h(X)) = E(e(X)) = E(E(Z \mid X)) = E(Z) = \P(Z = 1) = e$$
Use it to better understand IPW for ATT

7. Propensity Score in Regression

PS as a covariate

Theorem 7 (regression with pscore as a covariate).

Under unconfoundedness, the coefficient of $Z$ in the population OLS fit of $$ Y \sim 1+Z + e(X) $$ equals $\tau_{\mathrm{O}}$, $$ \tau_{\mathrm{O}}=\frac{E[e(X)\{1-e(X)\} \tau(X)]}{E[e(X)\{1-e(X)\}]} ,$$ which is the overlap-weighted average treatment effect

Based on above Theorem, we also have:

Corollary 1.

Under unconfoundedness,
1. the coefficient of $Z$ in the population OLS fit of $$ Y \sim 1+Z + e(X) + X $$ also equals $\tau_{\mathrm{O}}$,

2. the coefficient of $Z-e(X)$ in the population OLS fit of $$ Y \sim [Z - e(X)] \quad \text{or} \quad Y \sim 1 + [Z - e(X)] $$ also equals $\tau_{\mathrm{O}}$

PS as a weight

There is a convenient way to obtain $\hat{\tau}^{\text{hajek}}$ based on WLS.

Proposition 1 (convenient to obtain $\hat{\tau}^{\text{hajek}}$ based on WLS).

Need to use bootstrap for standard error
Why does the WLS give a consistent estimator for $\tau$ ?
In RCT with a constant propensity score, we can simply use the coefficient of $Z_i$ in the OLS fit of $Y_i$ on ( $1, Z_i$ ) to estimate $\tau$
In observational studies, we need to deal with the selection bias. The key idea is:
- If we weight the treated units by $\frac{1}{e(X_i)}$ and the control units by $\frac{1}{1-e(X_i)}$, then both treated and control groups can represent the whole population
- Thus, by weighting, we effectively have a pseudo-randomized experiment
- Remark 2 (IPCW).
  
  Inverse Probability of Censoring Weighting (IPCW) follows the same idea — it adjusts for censoring bias by reweighting observations based on their probability of being uncensored.
Consequently, the difference between the weighted means is consistent for $\tau$. The numerical equivalence of $\hat{\tau}^{\text {hajek }}$ and WLS is not only a fun numerical fact itself but also useful for motivating more complex estimators with covariate adjustment

Reference

Ding, Peng (2024), A First Course in Causal Inference, CRC Press.

Wager, S. (2024). Causal inference: A statistical learning approach. https://web.stanford.edu/~swager/causal_inf_book.pdf

Notes on DML for DiD: A Unified Approach

Mon, 02 Jun 2025 00:00:00 +0000

Introduction

This blog post explores how Double Machine Learning (DML) extends to conditional Difference-in-Differences (DiD), focusing on doubly robust estimators. The key insight is that conditional DiD can be understood through the lens of cross-sectional ATT estimation.

Foundation: Cross-Sectional ATT Estimation

To build intuition, we start with the familiar cross-sectional setting. Standard identification requires three assumptions: SUTVA, unconfoundedness, and overlap.

Step 1: Propensity Score Approach for ATT

Unlike ATE, ATT estimation requires only “one-sided” unconfoundedness and overlap conditions.

Assumption 1 (Identification Assumptions).

$$Z \indep Y(0) \mid X \text{ and } e(X) < 1$$
Estimate ATT using IPW

Theorem 1 (Ding (2024), Section 13.2).
More general, Li et al. (2018a) gave a unified discussion of the causal estimands in observational studies.

Theorem 2 (Ding (2024), Section 13.4).

Summary Table of common estimands:
- This table provides us a good way to understand and remember IPW estimator for ATT
- How to remember $\tau^h$? Apply IPW on “pseudo outcome” $Yh(X)$ then divide by $E(h(X))$
- When the parameter of interest is ATT, then $$E(h(X)) = E(e(X)) = E(E(Z \mid X)) = E(Z) = \P(Z = 1) = e$$
- Use it to better understand IPW for ATT

Step 2: Doubly Robust ATT Estimator

Combines outcome regression and IPW methods
For DR estimator of ATT, check my previous post
More generally, we have

Theorem 3 (DR for general estimand, see Ding (2024), page 191).

Extension to Conditional DiD

Identification Assumptions

Conditional DiD relies on two core assumptions: conditional parallel trends and no anticipation, plus an overlap condition.

Assumption 2 (CausalML Book, page 457).

How to understand the overlap condition (16.3.3)? It essentially imposes that there are control observations available for every value of $X$.

The Key Insight: Transformation to Cross-Sectional Problem

By taking the difference,

$$ \Delta Y = Y_{\text{after}} - Y_{\text{before}} $$

we transform panel data into a cross-sectional problem. This allows us to apply the same doubly robust framework used for cross-sectional ATT.

The Unified Result

The Neyman orthogonal score for conditional DiD is identical to the cross-sectional ATT score, where the outcome variable is simply the difference $\Delta Y$.

Neyman orthogonal score for ATT in conditional DiD

Proposition 1 (see CausalML Book).
Neyman orthogonal score for ATT in cross-sectional setting

Proposition 2 (see CausalML Book).
Comparing to the score for the ATT in cross-sectional setting, we see that DiD score is identical to that for learning the ATT under unconfoundedness where the outcome variable is simply defined as $\Delta Y$

This elegant connection demonstrates that the doubly robust estimator for conditional DiD is equivalent to the doubly robust ATT estimator applied to the differenced outcome $\Delta Y$.

References

Chernozhukov, Victor, Christian Hansen, Nathan Kallus, Martin Spindler, and Vasilis Syrgkanis (2024), “Applied causal inference powered by ML and AI.”

Ding, P. (2024). A First Course in Causal Inference. CRC Press.

Callaway, Brantly and Pedro H. C. Sant’Anna (2021), “Difference-in-Differences with multiple time periods,” Journal of Econometrics, Themed Issue: Treatment Effect 1, 225 (2), 200–230.

Chernozhukov, Victor, Whitney K Newey, and Rahul Singh (2022), “Debiased machine learning of global and local parameters using regularized Riesz representers,” The Econometrics Journal, 25 (3), 576–601.

Chernozhukov, Victor, Whitney K. Newey, Victor Quintas-Martinez, and Vasilis Syrgkanis (2024), “Automatic debiased machine learning via riesz regression.”

Notes on Callaway & Sant’Anna (2021) – Staggered Adoption DiD

Tue, 27 May 2025 00:00:00 +0000

0. Motivation

Staggered‐adoption policies break the canonical two-period / two-group DiD model. It has been shown that the traditional two-way fixed-effects (TWFE) regression can assign negative weights to treatment effects, thereby obscuring their dynamic and heterogeneous patterns.

Callaway & Sant’Anna (2021) propose a divide-and-conquer strategy:

Divide the messy staggered panel into many honest $2\times2$ DiDs
Conquer by estimating each “little” DiD under familiar assumptions, then combine them with user-chosen weights to answer specific questions

The key takeaway:

Use $\operatorname{ATT}(g, t)$ as a building block so we can transparently see how things are constructed
Many different aggregation schemes are possible: they deliver different parameters
Can allow for covariates via regressions adjustments, IPW, and DR.

1. Setup

Data structure: Panel of units $i$ over time $t = 1,\dots,T$
Let $D_{it}$ be a binary variable. $D_{it} = 1$ if unit $i$ is treated in period $t$; $D_{it} = 0$ otherwise
Cohorts
- Define $G$ as the time period when a unit first becomes treated. For all units that eventually get treated, $G$ defines which “group” they belong to
- Define $G_{g}$ as a binary variable. $G_{g} = 1$ if a unit is first treated in period $g$ (i.e. $G_{i,g} = \1\{G_i = g\}$ )
- Define $G=\infty$ as “never treated” group.
Potential outcomes
- $Y_{it}(g)$: outcome at time $t$ if first treated in period $g$.
- $Y_{it}(\infty)$: outcome if never treated.
Cohort-time ATT

Assume $\mathrm{iid}$, drop unit index $i$. A parameter of interest that has clear interpretation is the $\operatorname{ATT}(g, t)$: $$ \operatorname{ATT}(g, t)=\mathbb{E}\left[Y_t(g)-Y_t(\infty) \mid G_g=1\right], \text { for } t \geq g . $$ This defines one “clean” $2\times2$ DiD for each pair $(g,t)$.

Remark 1.

Why focus on time after treatment starting period, $t \ge g$? Because we need the no anticipation assumption. Before treatment taking place, $t < g$, there is no treatment effect.

2. Key Assumptions

Given that we never observe $Y(\infty)$ in post-treatment periods among units that have been treated, we need to make assumptions to identify $\operatorname{ATT}(g, t)$.

Assumption 1 (No anticipation).

$$ Y_{i t}(g)=Y_{i t}\left(g^{\prime}\right) \quad \forall \ i \text{, } t<\min \left\{g, g^{\prime}\right\} $$ Treatment cannot affect pre-treatment outcomes.

The no anticipation assumption has exactly the same content as in the $2\times2$ case.

Assumption 2 (Parallel trends based on 'never-treated' control).

For each $t \in\{2, \ldots, T\}, g \in \mathcal{G}$ such that $\textcolor{red}{t \geq g}$, $$ \mathbb{E}\left[Y_t(\infty)-Y_{t-1}(\infty) \mid G_g=1\right]=\mathbb{E}\left[Y_t(\infty)-Y_{t-1}(\infty) \mid G=\infty\right], \tag{1} $$

Note that, (1) is equivalent to the following: for $t \in\{2, \ldots, T\}, g \in \mathcal{G}, t \geq g$ ,

$$ \mathbb{E}\left[Y_t(\infty)-Y_{g-1}(\infty) \mid G_g=1\right]=\mathbb{E}\left[Y_t(\infty)-Y_{g-1}(\infty) \mid G=\infty\right], \tag{1’} $$

Assumption 3 (Parallel trends based on 'Not-Yet-Treated' groups).

For each $(s, t) \in\{2, \ldots, T\} \times\{2, \ldots, T\}, g \in \mathcal{G}$ such that $\textcolor{red}{t \geq g, s \geq t}$ $$ \mathbb{E}\left[Y_t(\infty)-Y_{t-1}(\infty) \mid G_g=1\right]=\mathbb{E}\left[Y_t(\infty)-Y_{t-1}(\infty) \mid D_s=0, G_g=0\right], \tag{2} $$

Similarly, (2) is equivalent to (2’), that is, changing $Y_{t-1}(\infty)$ to $Y_{g-1(\infty)}$ in (2).

3. Identification: Long-Difference Estimands

Under no anticipation and one of the parallel-trends assumptions, each $\mathrm{ATT}(g,t)$ equals a simple long-difference DID:

Using never-treated

$$ \operatorname{ATT}^{\text {never}}(g, t)=\mathbb{E}\left[Y_{t}-Y_{g-1} \mid G_g=1 \right]-\mathbb{E}\left[Y_{ t}-Y_{g-1} \mid G=\infty\right] $$

Using not-yet-treated

$$ \operatorname{ATT}^{\text {not-yet}}(g, t)=\mathbb{E}\left[Y_{t}-Y_{g-1} \mid G_g=1 \right]-\mathbb{E}\left[Y_{ t}-Y_{g-1} \mid D_t = 0, G=\infty\right] $$

Why it’s called "long difference"? Longer Time Span!

Remark 2.

Q: Why it’s called "long difference"?

A: Longer Time Span! Check my plot above. The "long difference" refers to the fact that the comparison often spans from a pre-treatment period (i.e. $g-1$) to a later period (post-treatment, $t \ge g$), potentially covering multiple time periods. This contrasts with "short differences," which might involve comparing outcomes in consecutive periods or shorter time windows.

For each treated cohort, the method computes the difference in outcomes between the pre-treatment period and a specific post-treatment period, potentially far apart in time. This extended gap emphasizes the "long" aspect, as it captures the cumulative effect of the treatment over time.

Moreover, with covariates, one can form doubly-robust estimators that combine generalized propensity scores $p_g(X_i)$ and outcome models $m_{g,t}(X_i)$. At high level, the form of this estimator is identical to AIPW–ATT in cross–section, but we need to replace “levels” (e.g. $Y$) with “changes” (e.g. $\Delta Y$).

For more details, check Theorem 1 in the paper.

4. Aggregation of $\mathrm{ATT}(g,t)$

Any overall summary $\theta$ is a weighted average of the cell-specific ATTs:

$$ \theta=\sum_{g=2}^T \sum_{t=g}^T w_{g, t} \operatorname{ATT}(g, t), \quad \sum_{g, t} w_{g, t}=1 . $$

Common choices:

Cohort-heterogeneity: Average effect of participating in the treatment that units in group $g$ experienced,

$$ \theta_S(g)=\frac{1}{T-g+1} \sum_{t=2}^T 1\{g \leq t\} \mathrm{ATT}(g,t) $$

Calendar time heterogeneity
Event-study / dynamic treatment effects

5. Limitation & Extension

Lee & Wooldridge (2023) argue that Callaway & Sant’Anna (2021) method is less efficient but more resilient to functional form of covariates. The following is from their working paper:

… CS (2021) method uses only the period just prior to the intervention in defining the control group, thereby discarding potentially useful information in earlier time periods.

In fact, Wooldridge (2021) shows that, under the standard “error components” structure on the error, with a homoskedastic time-constant component and homoskedastic and serially uncorrelated idiosyncratic errors, the POLS estimator is both best linear unbiased (BLUE) and asymptotically efficient. These theoretical results imply that the CS (2021) estimators are inefficient under a standard set of assumptions. The simulations in Wooldridge (2021) bear this out, showing the CS approach can be very inefficient. Balanced against the loss in precision is that the CS approach can be less biased when parallel trends are violated.

To improve the efficiency, instead of using long differences, Lee and Wooldridge (2023) use all suitable control observations in transforming the outcome variable. Specifically, rather than using the single period just prior to the treatment, $Y_{g-1}$, they use the pre-treatment average, $\frac{1}{g-1}\sum_{s = 1}^{g-1}Y_s$.

The rolling method (Lee and Wooldridge, 2023) transforms panel data with staggered interventions by subtracting each unit’s average outcome across all pre-treatment periods from their outcome in the current period of interest. This transformation, combined with no anticipation and parallel trends assumptions, makes the treatment assignment unconfounded for the transformed outcome in each cohort/time cross-section. With unconfoundedness holding, we can then apply standard treatment effects estimators, including doubly robust methods and matching, utilizing all not-yet-treated units as the valid control group for that specific cross-section.

6. R code Example

The following R code example is provided by Professor Scott Cunningham.

library(readstata13)
library(ggplot2)
library(did) # Callaway & Sant'Anna

castle <- data.frame(read.dta13('https://github.com/scunning1975/mixtape/raw/master/castle.dta'))
castle$effyear[is.na(castle$effyear)] <- 0 # untreated units have effective year of 0

# Estimating the effect on log(homicide)
atts <- att_gt(yname = "l_homicide", # LHS variable
               tname = "year", # time variable
               idname = "sid", # id variable
               gname = "effyear", # first treatment period variable
               data = castle, # data
               xformla = NULL, # no covariates
               #xformla = ~ l_police, # with covariates
               est_method = "dr", # "dr" is doubly robust. "ipw" is inverse probability weighting. "reg" is regression
               control_group = "nevertreated", # set the comparison group which is either "nevertreated" or "notyettreated" 
               bstrap = TRUE, # if TRUE compute bootstrapped SE
               biters = 1000, # number of bootstrap iterations
               print_details = FALSE, # if TRUE, print detailed results
               clustervars = "sid", # cluster level
               panel = TRUE) # whether the data is panel or repeated cross-sectional

# Aggregate ATT
agg_effects <- aggte(atts, type = "group")
summary(agg_effects)

# Group-time ATTs
summary(atts)

# Plot group-time ATTs
ggdid(atts)

# Event-study
agg_effects_es <- aggte(atts, type = "dynamic")
summary(agg_effects_es)

# Plot event-study coefficients
ggdid(agg_effects_es)

Reference

Callaway, Brantly and Pedro H. C. Sant’Anna (2021), “Difference-in-Differences with multiple time periods,” Journal of Econometrics, Themed Issue: Treatment Effect 1, 225 (2), 200–230.

Lee, S. J., & Wooldridge, J. M. (2023). A Simple Transformation Approach to Difference-in-Differences Estimation for Panel Data (SSRN Scholarly Paper No. 4516518). Social Science Research Network. https://doi.org/10.2139/ssrn.4516518

Sant’Anna, Pedro H. C. and Jun Zhao (2020), “Doubly robust difference-in-differences estimators,” Journal of Econometrics, 219 (1), 101–22.

How does doubly robust DiD estimator works? Check this: Lecture 5: How Covariates can make your DiD More Plausible

did R package 📦

Intuition for Doubly Robust Estimator

Sun, 18 May 2025 00:00:00 +0000

Introduction

To estimate ATE, we can either use outcome regression or inverse propensity weighting (IPW). While each approach has merits, combining them offers significant advantage – double robustness. In this post, I summarize the intuition for doubly robust estimator from Professor Ding’s textbook (Ding 2024), and connects this framework to debiased machine learning (DML) through Riesz representation theory. By understanding these connections, we can gain some insights into how modern causal inference methods effectively correct for bias in treatment effect estimation.

Two characterizations of the ATE

Assumption 1 (Basic setting).

SUTVA, unconfoundedness and overlap

Let $Y(0), Y(1)$ be potential outcomes and $D$ be a binary treatment variable. Consider ATE.

First, we can use the outcome regression,

$$ \tau = \E\{\mu_1(X) - \mu_2(X) \}, $$ where $$ \mu_1 = \E\{Y(1) \mid X \} = \E\{Y \mid D = 1, X \}, $$ $$ \mu_0 = \E\{Y(0) \mid X \} = \E\{Y \mid D = 0, X \} $$

Second, we can use the inverse propensity score weighting (IPW) approach,

$$\tau = \E\left\{\frac{DY}{e(X)} \right\} - \E\left\{\frac{(1-D)Y}{1-e(X)} \right\},$$

where $e(X) = \P(D = 1 \mid X)$ is the propensity score.

It completely ignores the outcome model. However, if there exist covariates $X$ that are predictive of $Y$, then even when the outcome model is misspecified, including it can reduce the variance compared to using IPW alone.

Key Insights

This motivates combining the two approaches to:

1. Reduce the variance of the IPW estimator

2. Reduce the bias of the outcome regression

Reducing the Variance

$$ \mu_1=E\{Y(1)\}=E\left\{Y(1)-\mu_1\left(X, \beta_1\right)\right\}+E\left\{\mu_1\left(X, \beta_1\right)\right\} . $$

Idea:

View $Y-\mu_1\left(X, \beta_1\right)$ as a “pseudo potential outcome”
Then apply IPW to it:

$$ \begin{aligned} \mu_1 & = E\left\{\frac{D\left\{Y-\mu_1\left(X, \beta_1\right)\right\}}{e(X)}\right\}+E\left\{\mu_1\left(X, \beta_1\right)\right\} \\ & =E\left\{\frac{D\left\{Y-\mu_1\left(X, \beta_1\right)\right\}}{e(X)}+\mu_1\left(X, \beta_1\right)\right\}, \end{aligned} $$

Similarly,

$$ \begin{aligned} \mu_0 & =E\left\{\frac{(1-D)\left\{Y-\mu_0\left(X, \beta_0\right)\right\}}{1-e(X)}\right\}+E\left\{\mu_0\left(X, \beta_0\right)\right\} \\ & =E\left\{\frac{(1-D)\left\{Y-\mu_0\left(X, \beta_0\right)\right\}}{1-e(X)}+\mu_0\left(X, \beta_0\right)\right\}, \end{aligned} $$

Notice that,

$$ \begin{aligned} \mu_1 - \mu_0 & = E\left\{\frac{D\left\{Y-\mu_1\left(X, \beta_1\right)\right\}}{e(X)}+\mu_1\left(X, \beta_1\right)\right\} \\ &\quad - E\left\{\frac{(1-D)\left\{Y-\mu_0\left(X, \beta_0\right)\right\}}{1-e(X)}+\mu_0\left(X, \beta_0\right)\right\} \end{aligned} $$

$\mu_1 - \mu_2$ gives us the AIPW estimator.

Reducing the Bias

$$ \mu_1 =E\left\{\frac{Z\left\{Y-\mu_1\left(X, \beta_1\right)\right\}}{e(X)}\right\}+E\left\{\mu_1\left(X, \beta_1\right)\right\} $$

Idea: We can view $Y-\mu_1\left(X, \beta_1\right)$ as the regression residuals, from which we apply IPW to extract useful signals. Alternatively, we can view $\mu_1\left(X, \beta_1\right) - Y$ as the bias, which we then use IPW to correct the bias.

Connecting to DDML: Riesz Representation for Bias Correction

In the generic debiased framework (Chernozhukov, Newey, and Singh 2022), we leverage the Riesz representer to correct for bias, similar to the approach described above. This method parallels our use of IPW to extract signals from residuals, but specifically employs the Riesz representer for bias correction.

Data is $Z = \{Y, D, X\}$
Let $g$ be outcome regression, $g(D, X)=E[Y \mid D, X]$
Suppose moment is of the form: for some moment $m()$ that is linear in $g$

$$\theta-E[m(Z ; g)]=0,$$ $$\theta = E[m(Z ; g)] :=E[g(1, X)-g(0, X)]$$

Then debiased version of the moment is of the form:

$$ \theta-E[m(Z ; g)+a(X) \cdot(Y-g(X))]=0, $$

where

$a(X)$ is the Riesz Representer of the linear functional $L(g) := E[m(Z ; g)]$. The existence of $a(X)$ is guaranteed by the Riesz representation theorem, for all square integrable $g(\cdot)$, we have

$$ E[m(Z ; g)]=E[a(X) \cdot g(X)]. $$

As we consider the ATE, the Riesz Representer are just some “inverse propensity score terms”,

$$ a(D, X):=\frac{D}{e(X)}-\frac{(1-D)}{1-e(X)}, $$

Consider the expression

$$ E[m(Z ; g)+a(X) \cdot(Y-g(X))] $$

The key intuition here is that $Y−g(X)$ represents the residual part, and we use the Riesz Representer $a(X)$ to correct this bias/residual term.

Double robust estimation of ATT

How about the double robust estimator for ATT? Can we also use the same idea to derive it? Yes!

Assumption 2 ( "one-sided" unconfoundedness and overlap).

$$ D \indep Y(0) \mid X \text { and } e(X)<1 . $$

Theorem 1.

Under the "one-sided" unconfoundedness and overlap assumption, $$ \E\{Y(0) \mid D=1\}=\frac{1}{e}\E\left\{\frac{e(X)}{1-e(X)} (1-D) Y\right\} \tag{1} $$ and $$ \tau_{\mathrm{T}}=\E(Y \mid D=1)-\E\left\{\frac{e(X)}{e} \frac{1-D}{1-e(X)} Y\right\}, $$ where $e=\P(D=1)$ is the marginal probability of the treatment.

We also have a doubly robust estimator for $\E\{Y(0) \mid D=1\}$ which combines the propensity score and the outcome models.

Theorem 2.

Define $$ \tilde{\mu}_{0 \mathrm{T}}^{\mathrm{dr}} := \frac{1}{e} \E\left[ \frac{e(X, \alpha)}{1 - e(X, \alpha)}(1-D)\left\{Y-\mu_0\left(X, \beta_0\right)\right\}+D \mu_0\left(X, \beta_0\right)\right] , $$ Under above Assumption, if either $e(X, \alpha)=e(X)$ or $\mu_0\left(X, \beta_0\right)=\mu_0(X)$, then $\tilde{\mu}_{0 \mathrm{T}}^{\mathrm{dr}}= \E\{Y(0) \mid D=1\}$.

How to come up with $\tilde{\mu}_{0 \mathrm{T}}^{\mathrm{dr}}$ ? Exactly the same idea as before!

$$ \E\{Y(0) \mid D=1\} = \E\{Y(0) - \mu_0\left(X, \beta_0\right) \mid D=1\} + \E\{\mu_0\left(X, \beta_0\right) \mid D=1\} $$

Now, we can view $Y(0) - \mu_0\left(X, \beta_0\right)$ as a “pseudo potential outcome” under the control and apply eqn (1) to weight it, then we can get the form of $\tilde{\mu}_{0 \mathrm{T}}^{\mathrm{dr}}$ .

Reference

Ding, Peng (2024) A first course in causal inference. Chapman & Hall.

Chernozhukov, Victor, Whitney K. Newey, and Rahul Singh (2022), “Automatic Debiased Machine Learning of Causal and Structural Effects,” Econometrica, 90 (3), 967–1027.

Big Picture of Debiased Machine Learning

Tue, 25 Mar 2025 00:00:00 +0000

Debiased machine learning (DML) is a generic recipe. The idea behind it is adding a correction term to the plug-in estimator of the functional, which leads to properties such as semi-parametric efﬁciency, double robustness, and Neyman orthogonality.

(Auto)-DML is a Method-of-Moments estimator

debiased/orthogonal moment scores
- Why it matters?
  - try to solve: model selection and/or regularization bias from ML learners (e.g. Lasso)
  - Neyman orthogonality: ensure the parameter of interest insentitive to first order perturbation of nuisance estimation
  - double robustness
  - asymptotic normality
- Key Idea: Debiasing is achieved by adding a correction term to the plug-in estimator of the functional
  - Three representations: $\theta = \mathbb{E}[m(W,g)] = \mathbb{E}[Y\alpha(W)] = \mathbb{E}[g(W)\alpha(W)]$, where
    - $g()$ is outcome regression;
    - $\alpha()$ is Rieze Representer (RR);
    - $m()$ is a continuous linear functional;
    - $W = (D, X)$ is data containing treatment $D$ and covariates $X$;
    - $Y(d)$ is potential outcome
  - Correct the residual using RR
    - $\mathbb{E}\{m(W,g) - \theta + \alpha(W)[Y-g(W)]\} = 0$
- How to construct orthogonal moment function?
  - orthogonal moment function = identifying moment function + first step influence function (FSIF)
  - identifying moment function: $m(W,g) - \theta$
    - involving outcome regression
  - FSIF: $\alpha(W)[Y-g(W)]$
    - correct the residual using Rieze Representer (RR)
    - Rieze Representer (RR)
      - In the case of ATE with binary treatment, RR are inverse propensity score terms
      - RR can be automatically characterized; NO NEED to know its analytical form
      - Can use random forests and NNet learners of RR
- Double Robustness
  - $\mathbb{E}[m(W ; g) -\theta_0 \left.+\alpha(W)(Y-g(W))\right] =-\mathbb{E}\left[\left(\alpha-\alpha_0\right)\left(g-g_0\right)\right]$
  - The score will be zero in expectation when either $\alpha(W) = \alpha_0(W)$ or $g(W) = g_0(W)$
Cross-fitting
- Why it matters?
  - Reduce overfitting bias

Notes on Causal Survival Forest 🌲⏳

Thu, 05 Sep 2024 00:00:00 +0000

In this post, I provide summary notes on the paper “Estimating Heterogeneous Treatment Effects with Right-Censored Data via Causal Survival Forests” by Cui et al. (2023).

Motivation

How to estimate heterogeneous treatment effects with right-censored data?

Heterogeneous treatment effect (HTE) estimation plays a central role in data-driven personalization
Existing methods often can’t handle censored survival outcomes, common in medical/business applications

Causal Survival Forests (CSF)

To address this challenge, the paper proposes causal survival forests (CSF)

An adaptation of the causal forest algorithm of Athey et al. (2019)
It adjusts for censoring using doubly robust estimating equations developed in the survival analysis literature

Advantages

Doubly robust, computationally tractable, and outperforms available baselines in our experiments
Good statistical properties – UCAN

Setting and notation

Assume i.i.d tuples $\{X_i, T_i, C_i, W_i\}$, where

$X_i \in \mathcal{X}$ denote covariates
$T_i \in \mathbb{R}_{+}$ is the survival time for $i$th unit
$C_i \in \mathbb{R}_{+}$ is the censoring time (the time at which $i$th unit gets censored)
$W_i \in \{0,1\}$ denotes a binary treatment

Using potential outcome framework, posit potential outcomes $\{T_i(1), T_i(0)\}$ s.t. $T_i = T_i(W_i)$, we need to estimate the conditional average treatment effect (CATE)

$$ \tau(x)=\mathbb{E}\left[y(T_i(1)) - y(T_i(0)) \mid X_i=x\right], $$

where $y()$ is the outcome transformation. For example,

$y(T) = T$
$y(T) = T \wedge h$ for the restricted mean survival time (RMST); here $h$ is some chosen maximum considered time
$y(T) = \1\{T \ge h\}$ for the survival probability

To estimate $\tau(x)$, the main challenge is that $T_i$ is not always observable. We can only observe:

censored survival time: $U_i = T_i \wedge C_i$
non-censoring indicator: $\Delta_i = 1\{T_i \le C_i\}$

Based on Assumption 1 (see later), we define the effective non-censoring indicator as follows:

\begin{align} \Delta_i^h &= 1\{T_i \wedge h \le C_i\} \\ &\overset{(2)}{=} \Delta_i \vee 1\{U_i \ge h\} \end{align}

Note that, for the eqn (2), everything is observed. We can regard an observation with $\Delta_i^h = 1$ as a complete observation.

Assumptions

In order to identify treatment effects, we need to rely on two sets of assumptions.

Assumption 2-4 enable us to identify the causal effect of $W_i$ on $T_i$ without censoring
Assumption 5-6 is to guarantee that censoring due to $C_i$ does not break identification results

Assumption 1 (Finite Horizon)

$$ y(t) = y(h), \quad \forall \ t \ge h, \ 0<h<\infty $$

Assumption 2 (Potential Outcomes)

$$ \{T_i(1), T_i(0)\} \quad s.t. \quad T_i = T_i(W_i) \quad a.s. $$ Assumption 3 (Ignorability)

$$ \{T_i(1), T_i(0)\} \perp W_i \mid X_i $$

Assumption 4 (Overlap)

Propensity score $e(x) = \mathbb{P}(W_i = 1\mid X_i = x)$ is uniformly bounded away from 0 and 1,

$$ \eta_e \le e(x)\le 1- \eta_e, \quad 0< \eta_e \le \frac{1}{2} $$

Assumption 5 (Ignorable censoring)

Censoring is independent of survival time conditionally on treatment and covariates,

$$ T_i \perp C_i \mid W_i, X_i $$

Assumption 6 (Positivity)

$$ \mathbb{P}(C_i < h | W_i, X_i) \le 1- \eta_c, \quad 0<\eta_c\le1 $$

Causal Forests Without Censoring

How does causal forest work?

Essentially, we are running a “forest”-localized version of Robinson’s regression

$$ \tau(x):=\operatorname{lm}\left(Y_i-\hat{m}^{(-i)}\left(X_i\right) \sim W_i-\hat{e}^{(-i)}\left(X_i\right), \text { weights }=\textcolor{blue}{\alpha_i(x)}\right), $$

where

$\textcolor{blue}{\alpha_i(x)}$ capture how “similar” a target sample $x$ is to each of the training samples $X_i$
$\hat{m}$ and $\hat{e}$ are machine learning estimates for the outcome and propensity score models with cross-fitting

Using notations in previous section, we estimate $\tau(x)$ by solving the following equation,

$$ \sum_{i=1}^n \alpha_i(x) \psi_{\hat{\tau}(x)}^{(c)}\left(X_i, y\left(T_i\right), W_i ; \hat{e}, \hat{m}\right)=0 \tag{cf} $$

where,

\begin{aligned} \psi_\tau^{(c)}\left(X_i, y\left(T_i\right), W_i ; \hat{e}, \hat{m}\right) &= \left[ W_i - \hat{e}\left(X_i\right) \right] \quad \times\\ & \left[ y\left(T_i\right) - \hat{m}\left(X_i\right) - \tau \left( W_i - \hat{e}\left(X_i\right) \right) \right] \end{aligned}

is the orthogonal complete score function (shown as up-script $^{(c)}$),

$e(x) = \mathbb{P}(W_i = 1\mid X_i = x)$
$m(x) = \mathbb{E}(y(T_i) \mid X_i)$
$\hat{e}(X_i)$ and $\hat{m}(X_i)$ are estimates derived via cross-fitting

Adjusting for Censoring via Weighting

In the presence of censoring, the $T_i$ in equation (cf) is no longer observable.

Simply ignoring censoring and building models on with complete observations (i.e. $\Delta_i^h = 1$) would lead to bias.

Simple Censoring Adjustment via IPCW

Define the conditional survival function for censoring process as $$ S_w^C(s \mid x)=\mathbb{P}\left[C_i \geq s \mid W_i=w, X_i=x\right] $$ We have, $$ \mathbb{P}\left[\Delta_i^h=1 \mid X_i, W_i, T_i\right]=S_{W_i}^C\left(T_i \wedge h \mid X_i\right) $$

the LHS is the conditional probability of observing a complete observations (i.e. $\Delta_i^h = 1$)
the RHS is the conditional probability that censoring time is greater than survival time
Does the above $\mathbb{P}(\Delta_i^h = 1 \mid \cdots)$ look like propensity score function?

The main idea of IPCW estimation is to only consider complete cases, but up-weight all complete observations by $1/S_{W_i}^C\left(T_i \wedge h \mid X_i\right)$ to compensate for censoring.

As a result, IPCW estimators succeed in eliminating censoring bias.

With IPCW, we estimate $\tau(x)$ by solving the following equation,

$$ \sum_{\left\{i: \Delta_i^h=1\right\}} \frac{\alpha_i(x)}{\hat{S}_{W_i}^C\left(T_i \wedge h \mid X_i\right)} \psi_{\hat{\tau}(x)}^{(c)}\left(X_i, y\left(T_i\right), W_i ; \hat{e}, \hat{m}\right)=0 \tag{IPCW} $$ Let’s compare the equation (cf) v.s (IPCW),

For eqn (cf), we sum over all observations; for eqn (IPCW), we only sum over complete observations
For eqn (IPCW), we add $\frac{1}{1/S_{W_i}^C\left(T_i \wedge h \mid X_i\right)}$ as a part of weight

For more details on IPCW, please check:

Chapter 8 and 12 in the textbook Causal Inference: What If" (Hernán and Robins, 2020). In particular, “Ch 12.6 Censoring and missing data” is very helpful.
Chapter 21 “Treatment Heterogeneity with Survival Outcomes” in the textbook Handbook of Matching and Weighting Adjustments for Causal Inference (Zubizarreta et al., 2023)

A Doubly Robust Correction

Two limitations of IPCW approach:

Only use complete observations; throw away all observations with $\Delta_i^h = 0$, and this may hurt efficiency
IPCW-type methods are generally not robust to estimation errors; Neyman orthogonality condition (Chernozhukov et al. 2018) does not hold

CSF Method

CSF method does not rely on IPCW. Instead, it relies on a more robust approach to making estimating equations robust to censoring.

Recall the simplest case (without censoring), we have,

$$ \sum_{i=1}^n \alpha_i(x) \psi_{\hat{\tau}(x)}^{(c)}\left(X_i, y\left(T_i\right), W_i ; \hat{e}, \hat{m}\right)=0 \tag{cf} $$

where

$\psi_{\hat{\tau}(x)}^{(c)} (\cdot)$ is the score function with the complete data

The idea of CSF is to convert the above (cf) equation into a censoring robust estimating equation by using estimates of the survival and censoring processes.

Now, we estimate the $\tau(x)$ by solving the following equation,

$$ \sum_{i=1}^n \alpha_i(x) \psi_{\hat{\tau}(x)}\left(X_i, y\left(U_i\right), U_i \wedge h, W_i, \Delta_i^h ; \hat{e}, \hat{m}, \hat{\lambda}_w^C, \hat{S}_w^C, \hat{Q}_w\right)=0, $$

where the score function is,

the conditional expectation of the transformed survival time is defined as:

$$ Q_w(s \mid x)=\mathbb{E}\left[y\left(T_i\right) \mid X_i=x, W_i=w, T_i \wedge h>s\right] $$

and the associated conditional hazard function is defined as:

$$ \lambda_w^{\mathrm{C}}(s \mid x)=-\frac{d}{d s} \log S_w^{\mathrm{C}}(s \mid x) $$

$\hat{Q}_w(s \mid x), \hat{S}_w^C(s \mid x)$ and $\hat{\lambda}_w^C(s \mid x)$ are cross-fit nuisance parameter estimates.

How to understand the above score function?

The short answer is that that functional form emerges for the math (i.e., the desire for a doubly robust adjustment); and, unlike with the basic AIPW formula, it’s not as immediately intuitive.¹

💡 KEY Points

We should think about the Neyman-orthogonal property. In summary, CSF alleviates the drawbacks of IPCW so by taking the (complete-data) causal forest estimating equation $\psi_{\tau(x)}^{(c)}(T, W, \ldots)$ (the “R-learner”) and turn it into a censoring robust estimating equation $\psi_{\tau(x)}(Y, W, \ldots)$ by using estimates of the survival and censoring processes². ("…" refers to additional nuisance parameters):

censoring process: $P\left[C_i>t \mid X_i=x, W_i=w\right]$
survival process: $P\left[T_i>t \mid X_i=x, W_i=w\right]$

The upshot of this “orthogonal” estimating equation is that it will be consistent if either the survival or censoring process is correctly specified, which is very beneficial when we want to estimate these by modern ML tools, such as random survival forests.

CSF approach is doubly robust in the sense that we can obtain the consistent estimator either the survival or censoring process is correctly specified.

For more details, Rubin & van der Laan (2007) and the chapter on RCTs with time-to-event data in Targeted Learning (2011) gives some more digestible details on doubly robust estimation with survival data.³

I also found the code in grf GitHub repo helpful to understand the implement procedure. Specifically, check the following lines:

# The conditional survival function for the survival process.
sf.survival <- do.call(grf::survival_forest, c(list(X = cbind(X, W), Y = Y, D = D), args.nuisance))

# The conditional survival function for the censoring process.
sf.censor <- do.call(grf::survival_forest, c(list(X = cbind(X, W), Y = Y, D = 1 - D), args.nuisance))

References

Athey, Susan, Julie Tibshirani, and Stefan Wager. 2019. “Generalized Random Forests.” The Annals of Statistics 47 (2): 1148–78. https://doi.org/10.1214/18-AOS1709.

Chernozhukov, Victor, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. 2018. “Double/Debiased Machine Learning for Treatment and Structural Parameters.” The Econometrics Journal 21 (1): C1–68. https://doi.org/10.1111/ectj.12097.

Cui, Yifan, Michael R Kosorok, Erik Sverdrup, Stefan Wager, and Ruoqing Zhu. 2023. “Estimating Heterogeneous Treatment Effects with Right-Censored Data via Causal Survival Forests.” Journal of the Royal Statistical Society Series B: Statistical Methodology 85 (2): 179–211. https://doi.org/10.1093/jrsssb/qkac001.

Hernán MA, Robins JM (2020). Causal Inference: What If. Boca Raton: Chapman & Hall/CRC.

Zubizarreta, J. R., Stuart, E. A., Small, D. S., & Rosenbaum, P. R. (2023). Handbook of Matching and Weighting Adjustments for Causal Inference. CRC Press.

This was suggested by Professor Wager in an email conversation. ↩︎
Check more on grf tutorial: Causal forest with time-to-event data ↩︎
Suggested by Erik Sverdrup. Many thanks! ↩︎

A walkthrough of how Causal Forest 🌲 works

Tue, 20 Aug 2024 00:00:00 +0000

Introduction

In this post, I will go over how causal forest works based on the tutorial in grf R package. Causal Forests offer a flexible, data-driven approach to estimating varied treatment effects, bridging machine learning and causal inference techniques.

Common Setting

If we are working on an observational study, we have the following data:

Outcome variable: $Y_i$
Binary treatment indicator: $W_i = \{0, 1\}$
A set of covariates: $X_i$

Let’s assume that the following conditions hold:

(Assumption 1) $W_i$ is unconfounded given $X_i$ (i.e. treatment is as good as random given covariates).

$$\{Y_i(0), Y_i(1)\} \perp W_i | X_i$$

(Assumption 2) The confounders $X_i$ have a linear effect on $Y_i$.
(Assumption 3) The treatment effect $\tau$ is constant.

Then we could run a regression of the type

$$ Y_i = \tau W_i + \beta X_i + \epsilon_i $$

and interpret the estimate of $\hat{\tau}$ as the average treatment effect (ATE) $\tau = \mathbb{E}(Y_i(1) - Y_i(0))$.

Relaxing Assumptions

Assumption 1 is an “identifying” assumption we have to live with
Assumption 2 and Assumption 3 are modeling assumptions that we can question.

Relaxing Assumption 2: Partially Linear Model (PLR)

Assumption 2 is a strong parametric modeling assumption that requires the confounders to have a linear effect on the outcome, and that we should be able to relax by relying on semi-parametric statistics.

We can instead posit the partially linear model:

$$ Y_i = \tau W_i + f(X_i) + \epsilon_i, \ \ \ \mathbb{E}(\epsilon_i | X_i, W_i) = 0 $$

How do we get around estimating $\tau$ when we do not know $f(X_i)$?

Define the propensity score as $$e(x) = \mathbb{E}(W_i | X_i = x),$$

and the conditional mean of $Y$ as

$$m(x) = \mathbb{E}(Y_i | X_i = x) = f(x) + \tau e(x).$$

By Robinson (1988), we can rewrite the above equation in “centered” form:

$$Y_i - m(x) = \tau \cdot [W_i - e(x)] + \epsilon_i$$

This formulation has great practical appeal, as it means $\tau$ can be estimated by residual-on-residual regression.

Good properties 😀: Robinson (1988) shows that this approach yields root-n consistent estimates of $\tau$, even if estimates of $m(x)$ and $e(x)$ converge at a slower rate (“4-th root” in particular). This property is often referred to as orthogonality and is a desirable property that essentially tells you that given noisy “nuisance” estimates ($m(x)$ and $e(x)$) you can still recover “good estimates of your target parameter ($\tau$). For more details, please refer to Wager, Stefan “STATS 361: Causal Inference” Lecture 3.

But how to estimate $m(x)$ and $e(x)$?

Use modern machine learning models! One could use boosting, random forest, and etc to estimate $m(x)$ and $e(x)$ because what we need is just “reasonable accurate” predictions, i.e.

$$\mathbb{E}\left[(\hat{m}(X)-m(X))^2\right]^{\frac{1}{2}}, \mathbb{E}\left[(\hat{e}(X)-e(X))^2\right]^{\frac{1}{2}}=o_P\left(\frac{1}{n^{1 / 4}}\right)$$

Issue with Direct Plug-in of Estimates: Directly plugging in $\hat{m}(x)$ and $\hat{e}(x)$ into the residual-on-residual regression typically leads to bias because modern ML methods regularize to trade off bias and variance.
Solution via Cross-Fitting: cross-fitting, where the prediction for observation $i$ is obtained without using unit $i$ for estimation, can help overcome this bias (Chernozhukov et al. 2018).

Recap: We have a way to adopt the modern ML toolkit to non-parametrically control for confounding when estimating an ATE, and still retain desirable statistical properties such as unbiased-ness and consistency.

Relaxing Assumption 3: Non-constant treatment effects

Non-constant treatment effects occur when the impact of a treatment varies across different subgroups or individuals. This concept relaxes the assumption of homogeneous treatment effects, where the treatment is assumed to have the same impact on all units.

We could specify certain subgroups and run separate regressions for each subgroup and obtain different estimates of $\tau$. To avoid false discoveries, this approach would require us to specify potential subgroups without looking at the data. How can we use the data to inform us of potential subgroups?

Let’s define,

$$ Y_i=\textcolor{red}{ \tau\left(X_i\right) } W_i+f\left(X_i\right)+\varepsilon_i, \quad E\left[\varepsilon_i \mid X_i, W_i\right]=0, $$

where $\textcolor{red}{ \tau\left(X_i\right) }$ is the conditional ATE, i.e., $\textcolor{red}{ \tau\left(X_i\right) } :=\mathbb{E}\left[Y_i(1)-Y_i(0) \mid X_i\right]$. How do we estimate this?

Idea: If we imagine we had access to some neighborhood $\mathcal{N}(x)$ where $\tau$ was constant, we could proceed exactly as before, by doing a residual-on-residual regression on the samples belonging to $\mathcal{N}(x)$, i.e.:

$$ \tau(x) := \operatorname{lm}\left(Y_i - \hat{m}^{(-i)}(X_i) \sim W_i - \hat{e}^{(-i)}(X_i), \text{ weights } = \mathbf{1}\{X_i \in \mathcal{N}(x)\} \right) $$

This is conceptually what Causal Forest does, it estimates the treatment effect $\tau(x)$ for a target sample $X_i = x$ by running a weighted residual-on-residual regression on samples that have similar treatment effects.

Recap: Causal Forest is running a a “forest”-localized version of Robinson’s regression.

These weights play a crucial role, so how does grf 📦 find them?

Random forest as an adaptive neighborhood finder

Breiman’s random forest for predicting the conditional mean $m(x) = \mathbb{E}(Y_i | X_i = x)$ can be briefly summarized in two steps:

Building phase: Build $B$ trees which greedily place covariate splits that maximize the squared difference in subgroups means

$$ n_L \cdot n_R \cdot ( \bar{y}_L - \bar{y}_R )^{2} $$

Prediction phase: Aggregate each tree’s prediction to form the final point estimate by averaging the outcomes $Y_i$ that fall into the same terminal leaf $L_b(X_i)$ as the targets sample $x$:

$$\begin{align} \hat{m}(x) &= \frac{1}{B} \sum_{b=1}^B \sum_{i=1}^n \frac{Y_i \mathbf{1}\left\{X_i \in L_b(x)\right\}}{\left|L_b(x)\right|} \\ &= \sum_{i=1}^n \frac{1}{B} \sum_{b=1}^B Y_i \frac{\mathbf{1}\left\{X_i \in L_b(x)\right\}}{\left|L_b(x)\right|} \tag{1} \\ &= \sum_{i=1}^n Y_i \textcolor{blue}{\frac{1}{B} \sum_{b=1}^B \frac{\mathbf{1}\left\{X_i \in L_b(x)\right\}}{\left|L_b(x)\right|}} \\ & =\sum_{i=1}^n Y_i \textcolor{blue}{\alpha_i(x)} \tag{2}, \end{align}$$

Note that, this procedure is a double summation, first over trees, then over training samples (see equation (1)). We can swap the order of summation and obtain $\textcolor{blue}{\alpha_i(x)}$ in the equation (2). We have defined $\textcolor{blue}{\alpha_i(x)}$ as the frequency with which the $i$-th training sample falls into the same leaf as $x$.

Does above $\textcolor{blue}{\alpha_i(x)}$ remind you the traditional deterministic kernel function and bandwidth story?

The following image illustrates how the $\textcolor{blue}{\alpha_i(x)}$ are calculated: Some dots are larger because they are used by all trees, while some are smaller because they are only used by a few trees.

Figure 1: Weights from Causal Forests

Causal Forest

Causal Forest essentially combines Breiman (2001) and Robinson (1988) by modifying the steps above to:

Building phase: Greedily places covariate splits that maximize the squared difference in subgroup treatment effects $$n_L \cdot n_R \cdot ( \hat{\tau}_L - \hat{\tau}_R )^{2},$$ where $\hat{\tau}$ is obtained by running Robinson’s residual-on-residual regression for each possible split point.
Use the resulting forest weights $\textcolor{blue}{\alpha_i(x)}$ to estimate $$\tau(x):=\operatorname{lm}\left(Y_i-\hat{m}^{(-i)}\left(X_i\right) \sim W_i-\hat{e}^{(-i)}\left(X_i\right), \text { weights }=\textcolor{blue}{\alpha_i(x)}\right),$$ where $\textcolor{blue}{\alpha_i(x)}$ capture how “similar” a target sample $x$ is to each of the training samples $X_i$

That is, Causal Forest is running a “forest”-localized version of Robinson’s regression. This adaptive weighting (instead of leaf-averaging) coupled with some other forest construction details known as “honesty” and “subsampling” can be used to give asymptotic guarantees for estimation and inference with random forests (Wager & Athey, 2018)

Efficiently estimating summaries of the CATEs

What about estimating summaries of $\tau(x)$, in terms of estimands like the average treatment effect (ATE), or the best linear projection (BLP), that have guaranteed $\sqrt{n}$ rate of convergence along with exact confidence intervals?

For estimating ATE, the most intuitive approach is to average the CATE, i.e. $\frac{1}{n}\sum_{i=1}^n \tau(X_i)$, right? However, there are more efficient methods than simply averaging individual CATE estimates.

Robins, Rotnitzky & Zhao (1994) showed that the so-called Augmented Inverse Probability Weighted (AIPW) estimator is asymptotically optimal for $\tau$ (meaning that among all non-parametric estimators, it has the lowest variance).

$$\begin{gathered} \hat{\tau}_{AIPW}=\frac{1}{n} \sum_{i=1}^n\left(\hat{\mu}_{(1)}\left(X_i\right)-\hat{\mu}_{(0)}\left(X_i\right)+\frac{W_i}{\hat{e}\left(X_i\right)}\left(Y_i-\hat{\mu}_{(1)}\left(X_i\right)\right)\right. \\ \left.-\frac{1-W_i}{1-\hat{e}\left(X_i\right)}\left(Y_i-\hat{\mu}_{(0)}\left(X_i\right)\right)\right), \end{gathered}$$

where $$\mu_{(w)}(x):=\mu(x, w) := \mathbb{E}\left[Y_i \mid X_i=x, W_i=w\right]$$ and $$e(x)=\mathbb{P}\left[W_i=1 \mid X_i=x\right]$$

To interpret the AIPW estimator $\hat{\tau}_{AIPW}$ , it is helpful to decompose it into two components: Let $\hat{\tau}_{AIPW} = A + B$ , where

$$ A=\frac{1}{n} \sum_{i=1}^n\left(\hat{\mu}_{(1)}\left(X_i\right)-\hat{\mu}_{(0)}\left(X_i\right)\right) $$ $$ B=\frac{1}{n} \sum_{i=1}^n\left(\frac{W_i}{\hat{e}\left(X_i\right)}\left(Y_i-\hat{\mu}_{(1)}\left(X_i\right)\right)-\frac{1-W_i}{1-\hat{e}\left(X_i\right)}\left(Y_i-\hat{\mu}_{(0)}\left(X_i\right)\right)\right) $$

$A$ represents the outcome regression adjustment estimator using $\hat{\mu}_{(w)}$
$B$ is an inverse propensity score weighting (IPW) estimator applied to the residuals $Y_i-\hat{\mu}_{\left(W_i\right)}\left(X_i\right)$
The AIPW estimator utilizes propensity score weighting on the residuals to debias the direct estimate
One key property of the AIPW estimator is its “double robustness”, which means that the estimator remains consistent and asymptotically normal even if either the outcome model or the propensity score model is misspecified. For proof, please refer to Stefan Wager’s Lecture 3 notes in “STATS 361: Causal Inference”.

The expression for above AIPW estimator can be rearranged and expressed as

$$ \begin{aligned} &\frac{1}{n} \sum_{i=1}^n\left(\tau\left(X_i\right) + \textcolor{red}{\left[\frac{W_i}{e(X_i)} - \frac{1-W_i}{1-e(X_i)}\right]} \textcolor{violet}{\left[Y_i-\mu\left(X_i, W_i\right)\right]} \right) \\ & = \frac{1}{n} \sum_{i=1}^n\left(\tau\left(X_i\right)+\frac{W_i-e\left(X_i\right)}{e\left(X_i\right)\left[1-e\left(X_i\right)\right]}\left(Y_i-\mu\left(X_i, W_i\right)\right)\right) \\ &\triangleq \frac{1}{n} \sum_{i=1}^n \Gamma_i \end{aligned} $$

We can understand above terms as:

$\tau(X_i)$ is an initial treatment effect estimate
$\textcolor{red}{\left[\frac{W_i}{e(X_i)} - \frac{1-W_i}{1-e(X_i)}\right]}$ is the Riesz Representer (Chernozhukov et al., 2022), which is used to “correct” the bias
$\textcolor{violet}{\left[Y_i-\mu\left(X_i, W_i\right)\right]}$ is the residual part
$\Gamma_i$ is called double robust score

References

Chernozhukov, Victor, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. 2018. “Double/Debiased Machine Learning for Treatment and Structural Parameters.” The Econometrics Journal 21 (1): C1–68. https://doi.org/10.1111/ectj.12097.
Athey, Susan, Julie Tibshirani, and Stefan Wager. 2019. “Generalized Random Forests.” The Annals of Statistics 47 (2): 1148–78. https://doi.org/10.1214/18-AOS1709.
Robinson, Peter M. “Root-N-consistent semiparametric regression.” Econometrica: Journal of the Econometric Society (1988): 931-954.
Wager, Stefan, and Susan Athey. “Estimation and inference of heterogeneous treatment effects using random forests.” Journal of the American Statistical Association 113.523 (2018): 1228-1242.
Wager, S. (2022). STATS 361: Causal Inference Lecture notes. Stanford University. https://web.stanford.edu/~swager/stats361.pdf
Chernozhukov, V., Newey, W. K., & Singh, R. (2022). Automatic Debiased Machine Learning of Causal and Structural Effects. Econometrica, 90(3), 967–1027. https://doi.org/10.3982/ECTA18515

doubly robust | Chen Xing

Notes on Doubly Robust Censoring Unbiased Transformation

1. Setup

2. The Challenge: Unidentifiability

The Surrogate-Response Strategy

3. Two Classical Transformations

A. The Buckley–James Transformation (depends on $\bar{F}$)

B. The IPCW Transformation (depends on $\bar{G}$)

4. The Doubly Robust Censoring Unbiased Transformation

5. Intuition: An AIPW in Disguise

Takeaway

Reference

Calibrate Your Nuisances: A Simple Fix for Doubly Robust Inference

Doubly Robust Inference via Calibration

TL;DR

What is this paper about?

Motivation

What do the authors do?

Why is this important?

Who should care?

Do we have code?

Reference

Balancing Weights for Causal Inference

TL;DR

What is this paper about?

What do the authors do?

Why is this important?

Who should care?

Do we have code?

Reference

Triply Robust Panel Estimators 🛡: When You Don't Know Which Assumptions Hold

TL;DR

What is this paper about?

What do the authors do?

Why is this important?

Who should care?

Do we have code?

Points I enjoyed

TROP vs. SDID

How to come up with this TROP estimator? My Thought

Reference

AIPW vs. Residual-on-Residual regression: Non-Parametric Flexibility or Efficiency?

Introduction

Residual-on-Residual regression

Augmented Inverse Propensity Weighting (AIPW) Estimator

The Key Difference: Non-Parametric vs. Partially Linear

Why This Matters

Conclusion

Reference

Notes on Propensity Score Methods

Introduction

Setting & Notation

1. The propensity score as a dimension reduction tool

2. Propensity score stratification

3. Propensity score weighting

3.1 Estimation

3.2 Strong overlap condition

4. Balancing property

4.1 Propensity score is a balancing score

5. Doubly Robust or AIPW

6. Other Estimands related to IPW

7. Propensity Score in Regression

PS as a covariate

PS as a weight

Reference

Notes on DML for DiD: A Unified Approach

Introduction

Foundation: Cross-Sectional ATT Estimation

Step 1: Propensity Score Approach for ATT

Step 2: Doubly Robust ATT Estimator

Extension to Conditional DiD

Identification Assumptions

The Key Insight: Transformation to Cross-Sectional Problem

The Unified Result

References

Notes on Callaway & Sant’Anna (2021) – Staggered Adoption DiD

0. Motivation

1. Setup

2. Key Assumptions

3. Identification: Long-Difference Estimands