How is EMSE derived for causal trees in Athey and Imbens (PNAS 2016)?

Question

Athey and Imbens build a non-parametric matching procedure to identify and estimate causal effects. To this end, they minimize the expected mean squared error (EMSE) of their procedure, but I don't understand how they get their expression for the EMSE.

I would greatly appreciate any and all help.

They introduce their approach by illustrating it for regression trees as opposed to causal trees. For this, they split the sample in three parts: training sample ($S^{tr}$) in which the tree ($\pi(S^{tr})$ or, when fixed, $\Pi$) is grown, the estimation sample ($S^{est}$) on which the estimator for the conditional mean ($\hat\mu(x; S^{est}, \Pi)$) is estimated, and the test sample ($S^{te}$) on which the MSE is evaluated.

They define the MSE (page 5) and the EMSE (page 6) as

\begin{align} \tag{1} {\rm MSE}(S^{te}, S^{est}, \Pi) &\equiv \frac{1}{|S^{te}|} \sum_{i \in S^{te}} \left\{ (Y_i - \hat\mu(X_i; S^{est}, \Pi))^2 - Y_i^2 \right\}, \\ \tag{2} {\rm EMSE}(\Pi) &\equiv \mathbb{E}_{S^{te}, S^{est}}[ {\rm MSE}(S^{te}, S^{est}, \Pi)]. \end{align}

On page 7, they expand ${\rm EMSE}(\Pi)$ as

\begin{align} \tag{3} - {\rm EMSE}(\Pi) &= - \mathbb{E}_{(Y_i, X_i), S^{est}} [(Y_i - \mu(X_i; \Pi))^2 - Y_i^2] \end{align}

and subsequently present the formula

\begin{align} & - \mathbb{E}_{X_i, S^{est}} \left[ (\hat\mu(X_i; S^{est}, \Pi) - \mu(X_i; \Pi))^2 \right] = \\ &\quad = \mathbb{E}_{X_i} [\mu^2(X_i; \Pi)] - \mathbb{E}_{S^{est}, X_i} [\mathbb{V}(\hat\mu^2(X_i; S^{est}, \Pi))]. \tag{4} \end{align}

My page numbers follow the pre-print on arXiv: https://arxiv.org/pdf/1504.01132.pdf

My questions are:

1. What are the steps to get from (1) to (3)?
2. Is the left-hand side of (4) equal to (3)?
   • If yes, what are the steps to show this?
   • If no, why do they proceed to construct an estimator for the right-hand side of (4) instead of (3)?

For question 1, what I get when I try to work from (1) to (3):

\begin{align*} & {\rm EMSE}(\Pi) \equiv \\ &\quad\equiv \mathbb{E}_{S^{te}, S^{est}} \left[ \frac{1}{|S^{te}|} \sum_{i \in S^{te}} \left\{ (Y_i - \hat\mu(X_i; S^{est}, \Pi))^2 - Y_i^2 \right\} \right] = \\ &\quad= \mathbb{E}_{(Y_i, X_i), S^{est}} [ (Y_i - \hat\mu(X_i; S^{est}, \Pi))^2 - Y_i^2 ] = \\ &\quad= \mathbb{E}_{(Y_i, X_i), S^{est}} [ \hat\mu^2(X_i; S^{est}, \Pi) - 2 Y_i \hat\mu(X_i; S^{est}, \Pi) ] = \\ &\quad= \mathbb{E}_{X_i, S^{est}} [\hat\mu^2(X_i; S^{est}, \Pi)] - 2 \mathbb{E}_{(Y_i, X_i), S^{est}} \left[ \mathbb{E}_{S^{est}} \left\{ Y_i \hat\mu(X_i; S^{est}, \Pi) \right\} \right] = \\ &\quad= \mathbb{E}_{X_i, S^{est}} [\hat\mu^2(X_i; S^{est}, \Pi)] - 2 \mathbb{E}_{(Y_i, X_i), S^{est}} \left[ Y_i \mathbb{E}_{S^{est}} \left\{ \hat\mu(X_i; S^{est}, \Pi) \right\} \right] = \\ &\quad= \mathbb{E}_{X_i, S^{est}} [\hat\mu^2(X_i; S^{est}, \Pi)] - 2 \mathbb{E}_{(Y_i, X_i), S^{est}} \left[ Y_i \mu(X_i; \Pi) \right] \end{align*} where the last equality uses that $\mathbb{E}_S(\hat\mu(x; S, \Pi)) = \mu(x; \Pi)$, i.e., that $\hat\mu(x; S, \Pi)$ is an unbiased estimator. I don't see how to reach (3) from here.

For question 2, I also don't see how the expected value of the conditional variance of $\hat\mu$ (on the left-hand side of (4)) is equal to the expected value of the conditional variance of $Y_i$ minus $Y_i^2$ (the right-hand side of (3)).

Answer 1

1. I guess (3) is another typo in the paper, it should be $-\mathbb{E}_{(Y_i, X_i), S^{est}}[(Y_i - \hat{\mu}(X_i; S^{est}, \Pi))^2 - Y_i^2]$. Then OP's derivations show that this is indeed $\mathbb{E}_{S^{est}, S^{te}}[MSE(S^{te}, S^{est}, \Pi)]$. And my answer 2. shows how to get from this to (4)
2. As for question 2, the answer is "almost", first, notice the typo: should be $\mathbb{E}_{X_i, \mathcal{S}^{est}}[\mathbb{V}[\hat{\mu}(X_i, \mathcal{S}^{est}, \Pi)]]$ (not $\hat{\mu}^2$). Second, I think the notation $\mathbb{E}_{X_i, \mathcal{S}^{est}}[\mathbb{V}[\hat{\mu}(X_i, \mathcal{S}^{est}, \Pi)]]$ is a bit confusing since it seems like the inner variance (as usual variance) is constant, so the outter expectation does not make any difference. In fact, what they meant by this expression is $$\mathbb{E}_{X_i}[\mathbb{V}_{\mathcal{S}^{est}}[\hat{\mu}(X_i, \mathcal{S}^{est}, \Pi)]]$$ With this notation, it is clear that the inner variance is a random variable (as a function of $X_i$), so the outter expectation is taken with respect to $X_i$. To get from (3) to (4) add and subtract $\mu(X_i; \Pi)$: $$(Y_i - \mu(X_i; \Pi)+\mu(X_i; \Pi) - \hat{\mu}(X_i; \mathcal{S}, \Pi))^2 - Y_i^2$$ then open up the brackets, you'll get three terms:

• $(Y_i - \mu(X_i; \Pi))^2 - Y_i^2$
• $(\hat{\mu}(X_i; \mathcal{S}, \Pi) - \mu(X_i; \Pi))^2$
• $2(Y_i - \mu(X_i; \Pi))(\mu(X_i; \Pi) - \hat{\mu}(X_i; \mathcal{S}, \Pi))$

The first term then is simplified to $\mathbb{E}[\mu^2(X_i; \Pi) - 2Y_i \mu(X_i; \Pi)]$, the second to $\mathbb{E}_{X_i}[\mathbb{V}_{\mathcal{S}^{est}}[\hat{\mu}(X_i, \mathcal{S}^{est}, \Pi)]]$, the third (using the fact that $\mathbb{E}_{\mathcal{S}}[\hat{\mu}(x; \mathcal{S}, \Pi)] = \mu(x; \Pi)$ and that $(X_i, Y_i)\ \perp\ \mathcal{S}^{est}$) to $2 \mathbb{E}_{(Y_i, X_i)}[Y_i \mu(X_i; \Pi)] - 2\mathbb{E}_{X_i}[ \mu(X_i) \mu(X_i; \Pi)]$. Finally, using the law of total probability (partition = leaves) and the fact that $\mu(x; \Pi) = \mathbb{E}[\mu(X_i)|X_i \in \ell(x; \Pi)]$ it can be shown that $2\mathbb{E}_{X_i}[ \mu(X_i) \mu(X_i; \Pi)] = 2 \mathbb{E}_{X_i}[\mu^2(X_i; \Pi)]$ Then, combining all together you get the formula.