We have now introduced the inverse probability weighting (IPW) estimator as
\[ \hat{\tau}_{\text{ipw}} = \frac{1}{n} \sum_{i=1}^n \left( \frac{Y_i A_i}{\hat{e}(X_i)} - \frac{Y_i (1 - A_i)}{1 - \hat{e}(X_i)} \right) \tag{1} \]
On the other hand, if we are able to come up with some good regression estimator of the expectation of the outcome \(Y\), say, \(\hat{\mu}_{(1)}(X)\), for those who received treatment, and \(\hat{\mu}_{(0)}(X)\), for those who did not receive treatment, then we can use the regression estimator of the ATE as
\[ \hat{\tau}_{\text{reg}} = \frac{1}{n} \sum_{i=1}^n \left( \hat{\mu}_{(1)}(X_i) - \hat{\mu}_{(0)}(X_i) \right) \tag{2} \]
An interesting idea of combining these two estimators is called an augmented inverse probability weighted (AIPW) estimator (Robins, Rotnitzky, and Zhao 1994), which is defined as
\[ \hat{\tau}_{\text{aipw}} = \frac{1}{n} \sum_{i=1}^n \left( \hat{\mu}_{(1)}(X_i) - \hat{\mu}_{(0)}(X_i) + \frac{A_i}{\hat{e}(X_i)} \left(Y_i - \hat{\mu}_{(1)}(X_i) \right) - \frac{(1 - A_i)}{1 - \hat{e}(X_i)} \left(Y_i - \hat{\mu}_{(0)}(X_i) \right) \right) \tag{3} \]
By just looking at this estimator, it is attempting to do a correction of the regression estimator. For an observation, \(Y_i - \hat{\mu}_{(1)}(X_i)\) is the error of the regression estimator for those who received treatment, and using a weighted version of this quantity would adjust the regression estimator. The same is true for those who did not receive treatment. Interestingly, this estimator displays a doubly robust property, which means that the AIPW estimator is consistent if either the outcome regression functions or the propensity score function is consistently estimated, under the causal identification and overlap assumptions introduced previously.
To understand the double robustness property, let’s consider the expectation version. This is reasonably close to the sample average as long as we have a large sample size.
\[ \mathbb{E} \left[ \hat{\mu}_{(1)}(X) + \frac{A}{\hat{e}(X)} \left(Y - \hat{\mu}_{(1)}(X) \right) \right] - \mathbb{E} \left[ \hat{\mu}_{(0)}(X) + \frac{(1 - A)}{1 - \hat{e}(X)} \left(Y - \hat{\mu}_{(0)}(X) \right) \right] \tag{4} \]
The first term should target \(\E [ Y(1) ]\). Let’s rearrange the first term and look at the difference between the two:
\[ \begin{aligned} &\E \left[ \hat{\mu}_{(1)}(X) + \frac{A}{\hat{e}(X)} \left(Y - \hat{\mu}_{(1)}(X) \right) - Y(1) \right] \\ =& \E \left[ \frac{A}{\hat{e}(X)} \left(Y - \hat{\mu}_{(1)}(X) \right) - \left( Y(1) - \hat{\mu}_{(1)}(X) \right) \right] \\ =& \E \left[ \frac{A - \hat{e}(X)}{\hat{e}(X)} \left( Y(1) - \hat{\mu}_{(1)}(X) \right) \right] \quad \text{by SUTVA} \\ =& \E \left[ \E \left[ \frac{A - \hat{e}(X)}{\hat{e}(X)} \left( Y(1) - \hat{\mu}_{(1)}(X) \right) \Biggm| X \right] \right] \\ =& \E \left[ \E \left[ \frac{A - \hat{e}(X)}{\hat{e}(X)} \Biggm| X \right] \times \E \left[ Y(1) - \hat{\mu}_{(1)}(X) \Biggm| X \right] \right] \quad \text{by Unconfoundedness}\\ =& \E \left[ \frac{e(X) - \hat{e}(X)}{\hat{e}(X)} \times \left[ \mu_{(1)}(X) - \hat{\mu}_{(1)}(X) \right] \right] \end{aligned} \]
Here, we treat the two plug-in estimates as fixed, as would be justified by evaluating them on an independent sample. We will return to this point in the cross-fitting section. It is easy to see that the remainder is zero whenever \(\hat{\mu}_{(1)}(X) = \mu_{(1)}(X)\) or \(\hat{e}(X) = e(X)\). The derivation of the second term is almost identical. It may also be interesting to establish an upper bound for this remainder. If we apply the Cauchy-Schwarz inequality,
\[ |\E(X Y)| \leq \sqrt{\E(X^2) \E(Y^2)} \]
Under the overlap condition, we have
\[ \begin{aligned} & \biggm| \mathbb{E} \left[ \frac{e(X) - \hat{e}(X)}{\hat{e}(X)} \times \left[ \mu_{(1)}(X) - \hat{\mu}_{(1)}(X) \right] \right] \biggm| \\ \leq & \sqrt{\mathbb{E} \left[ \left( \frac{e(X) - \hat{e}(X)}{\hat{e}(X)} \right)^2 \right] \times \mathbb{E} \left[ \left( \mu_{(1)}(X) - \hat{\mu}_{(1)}(X) \right)^2 \right]} \end{aligned} \]
This suggests that the remainder is upper bounded by the product of the \(L_2\) estimation errors of \(\hat{e}(X)\) and \(\hat{\mu}_{(1)}(X)\), after accounting for the propensity-score denominator. Hence, if either nuisance function is consistently estimated, this product remainder converges to zero under suitable regularity conditions.
An alternative view of this bias reduction can be seen by reorganizing the terms in (4):
\[ \mathbb{E} \left[ \hat{\mu}_{(1)}(X) - \hat{\mu}_{(0)}(X) \right] + \mathbb{E} \left[ \frac{A}{\hat{e}(X)} \left(Y - \hat{\mu}_{(1)}(X) \right) - \frac{(1 - A)}{1 - \hat{e}(X)} \left(Y - \hat{\mu}_{(0)}(X) \right) \right] \]
In this case, the first term \(\mathbb{E} \left[ \hat{\mu}_{(1)}(X) - \hat{\mu}_{(0)}(X) \right]\) is the regression estimand for the ATE, while the second term corrects it using propensity-weighted outcome residuals.
The variance of the AIPW estimator is a bit more complicated. Let’s re-organize the terms in (4) a bit before the analysis
\[ \mathbb{E} \left[ \hat{\mu}_{(1)}(X) - \hat{\mu}_{(0)}(X) \right] + \mathbb{E} \left[ \frac{A}{\hat{e}(X)} \left(Y - \hat{\mu}_{(1)}(X) \right) - \frac{(1 - A)}{1 - \hat{e}(X)} \left(Y - \hat{\mu}_{(0)}(X) \right) \right] \]
At the sample level, both the empirical average of the fitted conditional effects and the residual correction contribute to the variance. The first component need not have negligible variance, especially when treatment effects vary across \(X\). To see why accurate outcome regressions can still help, we can analyze the residual correction associated with \(A = 1\):
\[ \begin{aligned} & \Var \left[ \frac{A}{\hat{e}(X)} \left( Y - \hat{\mu}_{(1)}(X) \right) \right] \\ =& \Var \left[ \frac{A}{\hat{e}(X)} \left( Y(1) - \hat{\mu}_{(1)}(X) \right) \right] \quad \text{by SUTVA} \\ \leq & \E \left[ \frac{A}{\hat{e}^2(X)} \left( Y(1) - \hat{\mu}_{(1)}(X) \right)^2 \right] \quad \text{by}\,\, \Var(X) = \E(X^2) - \E^2(X)\\ = & \E \left[ \frac{e(X)}{\hat{e}^2(X)} \times \E \left[ \left( Y(1) - \hat{\mu}_{(1)}(X) \right)^2 \biggm| X \right] \right] \quad \text{by Unconfoundedness} \end{aligned} \]
When \(\hat{\mu}_{(1)}(X)\) is accurate, this residual can be less variable than the uncentered outcome used by IPW. This provides an efficiency motivation for augmentation, but it does not guarantee that AIPW has smaller variance than IPW in every finite sample.
Let’s compare several estimators we have covered so far:
The following example is modified from Ding (2023).
ATE_Est <- function(a, y, x,
truncps = c(0, 1)) # truncate propensity score if needed (positivity assumption)
{
ate_data = data.frame("a" = a, "y" = y, x)
# fitted propensity score
pscore <- glm(a ~ . - y, data = ate_data, family = "binomial")$fitted.values
pscore <- pmax(truncps[1], pmin(truncps[2], pscore))
# fitted potential outcomes
# weights simply restrict to part of the data based on the label by still predict all subjects
outcome1 <- lm(y ~ . - a, data = ate_data, weights = a)$fitted.values
outcome0 <- lm(y ~ . - a, data = ate_data, weights = (1 - a))$fitted.values
# outcome regression estimator
ate_reg <- mean(outcome1 - outcome0)
# IPW estimators
y_treat <- mean(a * y / pscore)
y_control <- mean((1 - a) * y / (1 - pscore))
ate_ipw <- y_treat - y_control
# Hajek estimator (see HW1)
one_treat <- mean(a / pscore)
one_control <- mean((1 - a) / (1 - pscore))
ate_hajek <- y_treat / one_treat - y_control / one_control
# doubly robust estimator
res1 <- y - outcome1
res0 <- y - outcome0
r_treat <- mean(a * res1 / pscore) # residual correction term for treated
r_control <- mean((1 - a) * res0 / (1 - pscore)) # residual correction term for control
ate_dr <- ate_reg + r_treat - r_control
return(c(ate_reg, ate_ipw, ate_hajek, ate_dr))
}We will consider four different cases to compare these estimators:
n = 500
nsim = 1000
set.seed(1)
data_gen <- function(n,
ps = TRUE, # TRUE/FALSE for correct propensity score
reg = TRUE) # TRUE/FALSE for correct regression model
{
# two dimensional covariates plus intercept
x <- matrix(rnorm(n * 2), n, 2)
if (ps) # TRUE/FALSE for correct propensity score
{
# generate propensity score and treatment label
beta_ps <- c(1, 1)
pscore <- 1 / (1 + exp(-as.vector(x %*% beta_ps)))
a <- rbinom(n, 1, pscore)
}else{
# nonlinear propensity score
x1 = cbind(x, exp(x))
beta_ps <- c (0, 0, 1, -1)
pscore <- 1 / (1 + exp(1 - as.vector(x1 %*% beta_ps)))
a <- rbinom(n, 1, pscore)
}
if (reg) # TRUE/FALSE for correct regression model
{
# generate potential outcomes
beta_y1 <- c(2, 1)
beta_y0 <- c(2, 1)
y1 <- rnorm(n, 1 + x %*% beta_y1)
y0 <- rnorm(n, 1 + x %*% beta_y0)
}else{
# generate potential outcomes
x1 = cbind(x, exp(x))
beta_y1 <- c(1, -1, 0.5, -0.2)
beta_y0 <- c(-1, 1, -0.5, 0.2)
y1 <- rnorm(n, 1 + x1 %*% beta_y1)
y0 <- rnorm(n, 1 + x1 %*% beta_y0)
}
# observed outcome
y <- a * y1 + (1 - a) * y0
return(list(a = a, y = y, x = x,
tau = mean(y1 - y0))) # tau is the true SATE for this simulated dataset
}In this case, the data generator is exactly what our model is. Hence, the doubly robust estimator should perform well. It should also be noticed that since the regression estimator is also correct and pretty much optimal (root-\(n\) sense), there is really not much to improve from it.
result_pool = matrix(NA, nsim, 4)
for (k in 1:nsim)
{
mydata = data_gen(n, ps = TRUE, reg = TRUE)
# estimated values
result_pool[k, ] = ATE_Est(mydata$a, mydata$y, mydata$x) - mydata$tau
}
# summarize results from all simulations
summary_mat = rbind(apply(result_pool, 2, mean),
apply(result_pool, 2, sd))
rownames(summary_mat) <- c("estimate", "se")
colnames(summary_mat) <- c("reg", "ipw", "Hajek", "DR")
round(summary_mat, 5)
## reg ipw Hajek DR
## estimate 0.00150 -0.00357 0.01170 0.00563
## se 0.08714 0.35163 0.29993 0.10571For this case, we specify a different data generator for the propensity score. Hence our working model would be specified wrong. The IPW estimator can be very biased in this case.
result_pool = matrix(NA, nsim, 4)
for (k in 1:nsim)
{
mydata = data_gen(n, ps = FALSE, reg = TRUE)
# estimated values
result_pool[k, ] = ATE_Est(mydata$a, mydata$y, mydata$x) - mydata$tau
}
# summarize results from all simulations
summary_mat = rbind(apply(result_pool, 2, mean),
apply(result_pool, 2, sd))
rownames(summary_mat) <- c("estimate", "se")
colnames(summary_mat) <- c("reg", "ipw", "Hajek", "DR")
round(summary_mat, 5)
## reg ipw Hajek DR
## estimate -0.00470 -0.68136 -0.64346 -0.01031
## se 0.10487 0.71156 0.50183 0.18670The outcome model is modified to include nonlinear terms. In this case, the regression estimator is biased, while the DR estimator may still show some finite-sample error. However, the IPW and Hajek estimators remain unbiased because they are not affected by the outcome regression model. The doubly robust estimator remains consistent because the propensity score model is correctly specified.
result_pool = matrix(NA, nsim, 4)
for (k in 1:nsim)
{
mydata = data_gen(n, ps = TRUE, reg = FALSE)
# estimated values
result_pool[k, ] = ATE_Est(mydata$a, mydata$y, mydata$x) - mydata$tau
}
# summarize results from all simulations
summary_mat = rbind(apply(result_pool, 2, mean),
apply(result_pool, 2, sd))
rownames(summary_mat) <- c("estimate", "se")
colnames(summary_mat) <- c("reg", "ipw", "Hajek", "DR")
round(summary_mat, 5)
## reg ipw Hajek DR
## estimate -0.13398 -0.00158 -0.00435 -0.00228
## se 0.11941 0.31770 0.29824 0.20697In this case, neither condition for double robustness holds, so the AIPW estimator is no longer protected against model misspecification. In other words, doubly robust does not mean robust when both nuisance models are misspecified.
result_pool = matrix(NA, nsim, 4)
for (k in 1:nsim)
{
mydata = data_gen(n, ps = FALSE, reg = FALSE)
# estimated values
result_pool[k, ] = ATE_Est(mydata$a, mydata$y, mydata$x) - mydata$tau
}
# summarize results from all simulations
summary_mat = rbind(apply(result_pool, 2, mean),
apply(result_pool, 2, sd))
rownames(summary_mat) <- c("estimate", "se")
colnames(summary_mat) <- c("reg", "ipw", "Hajek", "DR")
round(summary_mat, 5)
## reg ipw Hajek DR
## estimate -0.20039 -0.26808 -0.38877 0.32407
## se 0.16309 0.48331 0.46917 0.51623There is actually a little catch in our previous derivation: we treated the fitted propensity score and outcome models as fixed when evaluating the AIPW score. Cross-fitting provides this separation when flexible machine learning methods are used, but it is not required for the basic double-robustness property under standard parametric modeling conditions. To illustrate the idea, let’s consider splitting the data into two halves. In equation (3), we estimate all \(\hat{e}(X)\) and \(\hat{\mu}(X)\)’s using the first half of the data, and then use them as plug-in estimators to calculate the ATE terms using the second half. We can then reverse the roles of the two halves and combine all terms. Conditional on each training sample, the nuisance estimates used for an observation are independent of that observation. This idea is connected to semiparametric efficiency theory, which is an advanced statistical topic Bickel et al. (1993). More details can be found in the textbook Kosorok (2008).
The following updated code implements this idea. However, for this case, it does not really improve the standard error much, mainly because \(\hat{e}(X)\) and \(\hat{\mu}(X)\) are simple parametric estimators with root-\(n\) rates. Cross-fitting becomes especially useful when more general machine learning models are used (Chernozhukov et al. 2018). A key result from semiparametric efficiency theory is that, under suitable regularity and overlap conditions, a sufficient requirement is that the product of the two \(L_2\) nuisance estimation errors is \(o_p(n^{-1/2})\). For example, this condition holds if both nuisance estimators converge faster than \(n^{-1/4}\), allowing the doubly robust estimator to remain root-\(n\) consistent. 1
dr_cross <- function(a, y, x, idx)
{
alldata = data.frame(a, y, x)
data_plugin = alldata[-idx, ]
data_est = alldata[idx, ]
# fitted propensity score
ps_model <- glm(a ~ . - y, data = data_plugin, family = "binomial")
pscore_est <- predict(ps_model, newdata = data_est, type = "response")
# fitted potential outcomes
outcome_model1 <- glm(y ~ . - a, data = data_plugin, weights = a, family = "gaussian")
outcome_est1 <- predict(outcome_model1, newdata = data_est)
outcome_model0 <- glm(y ~ . - a, data = data_plugin, weights = (1 - a), family = "gaussian")
outcome_est0 <- predict(outcome_model0, newdata = data_est)
# doubly robust estimator
res1_est <- data_est$y - outcome_est1
res0_est <- data_est$y - outcome_est0
r_treat_est <- mean(data_est$a * res1_est / pscore_est) # residual correction term for treated
r_control_est <- mean((1 - data_est$a) * res0_est / (1 - pscore_est)) # residual correction term for control
dr_crossfit <- mean(outcome_est1 - outcome_est0) + r_treat_est - r_control_est
# estimated values
return(dr_crossfit)
} result_drcross = rep(NA, nsim)
for (k in 1:nsim)
{
mydata = data_gen(n, ps = TRUE, reg = TRUE)
# fit the cross-fitted doubly robust estimator
M = 5 # use a 5 fold version
allidx = sample(1:n, n)
index <- seq_along(allidx)
factor_levels <- cut(index, breaks=M, labels=FALSE)
chunks <- split(allidx, factor_levels)
cross_est_fold = rep(NA, M)
for (j in 1:M)
{
cross_est_fold[j] = dr_cross(mydata$a, mydata$y, mydata$x, chunks[[j]])
}
# estimated values
result_drcross[k] = mean(cross_est_fold) - mydata$tau
}
# summarize results from all simulations
mean(result_drcross)
## [1] 0.01099088
sd(result_drcross)
## [1] 0.117282Sometimes it could be difficult to derive the variance of an ATE estimator analytically. However, as we introduced previously, the bootstrap method can be used to estimate the variance of an estimator by re-sampling the data with replacement and re-calculating the estimator on each bootstrap sample. Using this idea, we can easily implement a bootstrap estimation of the variance for any ATE estimator. The following code illustrates this for all four estimators above.
n = 500
nbs = 500
mydata = data_gen(n, ps = TRUE, reg = TRUE)
point.est = ATE_Est(mydata$a, mydata$y, mydata$x)
boot.est = replicate(nbs, {
id.boot = sample(1:n, n, replace = TRUE )
ATE_Est(mydata$a[id.boot], mydata$y[id.boot], mydata$x[id.boot, ]) - mydata$tau
})
rbind(point.est,
"sd" = apply(boot.est, 1, sd))
## [,1] [,2] [,3] [,4]
## point.est -0.008017671 -1.050069 -0.8889289 -0.3034686
## sd 0.098965547 1.259396 0.7804645 0.2444319grf PackageRecall our discussion of random forest estimation with influence functions. The grf package treats the doubly robust estimator as an efficient influence function estimator. Here, we update our previous formula explicitly into the doubly robust version
\[ \psi(Y_i, A_i, X_i; \mu_0, \mu_1, e) = \left[ \mu_1(X_i) - \mu_0(X_i) \right] + \frac{A_i}{e(X_i)} (Y_i - \mu_1(X_i)) - \frac{1 - A_i}{1 - e(X_i)} (Y_i - \mu_0(X_i)) - \tau(x) \]
And we proceed with the forest weighted moment condition to get the local estimator of \(\tau(x)\):
\[ \sum_{i=1}^n w_i(x) \psi(Y_i, A_i, X_i; \mu_0, \mu_1, e) = 0 \]
with \(\mu_0\), \(\mu_1\) and \(e\) being estimated from the data. The following code illustrates how to use the grf package to estimate the ATE with doubly robust estimation, on the lalonde dataset:
library(grf)
library(Matching)
data("lalonde")
# prepare data
y = lalonde$re78
a = lalonde$treat
x = lalonde[, !(names(lalonde) %in% c("re78", "treat"))]
# fit the causal forest with doubly robust estimation
cf_dr = causal_forest(x, y, a,
Y.hat = NULL, W.hat = NULL, # let grf estimate the nuisance functions
num.trees = 2000)
# estimate ATE
ate_dr_grf = average_treatment_effect(cf_dr, target.sample = "all")
ate_dr_grf
## estimate std.err
## 1549.7443 669.5565The estimation here would include all samples and estimate the ATE when applying the treatment on all of them. If one wants to consider just the ones being treated, then target.sample = "treated" can be used.
# estimate ATE for treated only
ate_dr_grf_treated = average_treatment_effect(cf_dr, target.sample = "treated")
ate_dr_grf_treated
## estimate std.err
## 1725.7991 687.4033Out-of-bag prediction in random forests uses trees that do not contain the target observation. It is closely related to sample splitting and can reduce own-observation overfitting, although it is not identical to fixed-fold cross-fitting.↩︎