# Formula for E(MSE) from Recursive Partitioning for Heterogeneous Causal Effects (Athey and Imben 2015)

I'm reading Recursive Partitioning for Heterogeneous Causal Effects (Athey and Imben 2015) and I'm confused about the formula for $$\hat E[\mu(x;\Pi)]$$ on page 8. The formula offered in the paper is

but I don't understand how they got this formula. Given that the formula for variance is

$$Var(\mu) = E(\mu^2)-E(\mu)^2$$ or $$E(\mu^2) = E(\mu)^2 + Var(\mu)$$

shouldn't the formula from the paper be $$\hat E[\mu^2(x;\Pi)] = \hat\mu^2(x; S^{tr},\Pi) + \frac{S^2_{S^{tr}}(l(x;\Pi))}{N^{tr}(l(x;\Pi))} \text{ ?}$$ I.e. shouldn't the minus sign be a plus? Can someone explain how they derived that formula?

Their formula has the correct sign, and the crux of the reason why is that they're doing the variance decomposition with respect to a distribution different than the one you are. Considering this different distribution will make it so that the square is on the outside of the expectation, rather than on the inside as you have it. We'll take notation from their paper without explanation, and, as they do, we will treat the partition $$\Pi$$ as fixed.
First, notice that $$\mathbb{E}_{\mathcal{S}^\textrm{tr}}[\hat{\mu}(x; \mathcal{S}^\textrm{tr}, \Pi)] = \mu(x; \Pi),$$ so that $$\hat{\mu}(x; \mathcal{S}^\textrm{tr}, \Pi)$$ is an unbiased estimator for $$\mu(x; \Pi)$$. We will use these to make an estimator for the square $$\mu^2(x; \Pi)$$, via $$\underbrace{\mathbb{E}_{\mathcal{S}^\textrm{tr}}[\hat{\mu}^2(x; \mathcal{S}^\textrm{tr}, \Pi)]}_{(1)} - \underbrace{\mathrm{Var}_{\mathcal{S}^\textrm{tr}} \left( \hat{\mu}(x; \mathcal{S}^\textrm{tr}, \Pi) \right)}_{(2)} = \left( \mathbb{E}_{\mathcal{S}^\textrm{tr}}[\hat{\mu}(x; \mathcal{S}^\textrm{tr}, \Pi)] \right)^2 = \mu^2(x; \Pi).$$ The first term $$(1)$$ can be estimated simply by $$\hat{\mu}^2(x; \mathcal{S}^\textrm{tr}, \Pi)$$, and the second term $$(2)$$ can be estimated by the quantity in the paper.
Ultimately, the goal is to estimate $$\mathbb{E}_{X_i} \left[ \mu^2 (X_i; \Pi) \right],$$ where $$X_i$$ is a draw from the test set. This is done by averaging the estimator (described above) over the test points.