The reference is GENERALIZED RANDOM FORESTS by ATHEY, TIBSHIRANI and WAGER (2019). They construct a general algorithm to grow trees and forest for estimation of target parameters that are conditional on a covariate vector $X$.

Notation and Setup

Observable data $O_{i}=(Y_{i},W_{i})$, where $Y$ defines the outcome, $W\in\{0,1\}$ defines the binary treatment, and $X\in\mathbb{R}^{p}$ represents a covariates matrix. In matrix form

$$W=\left(\begin{array}{c} W_{1}\\ \vdots\\ W_{N} \end{array}\right)_{N\times1}Y=\left(\begin{array}{c} Y_{1}\\ \vdots\\ Y_{N} \end{array}\right)_{N\times1}X=\left(\begin{array}{ccc} X_{11} & \cdots & X_{1p}\\ \vdots & \ddots\\ X_{N1} & & X_{Np} \end{array}\right)_{N\times p} $$

The paper defines the target parameter $\theta(x)$ (a possible nuisance parameters $\nu(x)$) as solutions to local moment conditions $$\mathbb{E}\left[\psi_{\theta,\nu}\left(O\right)\mid X=x\right]=0.$$

In the example of conditional treatment effect (CATE) estimation we assume $Y=W_ib_i+e_i$ and conditional independence, and the CATE is defined as $\beta(x)=\mathbb{E}[b_i\mid X_i=x]$. Then a for a given contrast $\xi$, the target parameter $\theta(x)=\xi\beta(x)$ is identified via the score function \begin{align*} \psi_{\beta,c}(Y_{i},W_{i})_{\mid X_{i}=x} & =\left[\begin{array}{c} Y_{i}-\beta(x)W_{i}-c(x)\\ \left(Y_{i}-\beta(x)W_{i}-c(x)\right)W_{i} \end{array}\right]_{2\times1} \end{align*} which defines the two moment conditions in experiments, \begin{align*} \mathbb{E}\left[\psi_{\beta,c}\left(Y,W\right)\mid X_{i}=x\right]=0\leftrightarrow\left[\begin{array}{c} \mathbb{E}\left[Y-\beta(x)W-c(x)\mid X_{i}=x\right]\\ \mathbb{E}\left[\left(Y-\beta(x)W-c(x)\right)W\mid X_{i}=x\right] \end{array}\right]_{2\times1}=\left[\begin{array}{c} 0\\ 0 \end{array}\right]_{2\times1} \end{align*} for all $x\in\mathcal{X}$.

A gradient of $\mathbb{E}[\psi\mid X_i=x]$

The generalized random forest performs CART splits on psuedo-outcomes, calculated as follows. Given a parent node $P$, we define the pseudo-outcome $\rho$ by

$$\rho_{i}=-\xi^{T}A_{p}^{-1}\psi_{\hat{\theta}_{P},\hat{\nu}_{P}}(O_{i}),$$ where $A_{P}$ is a consistent estimate for the gradient of $\mathbb{E}\left[\psi_{b,c}\left(Y,W\right)\mid X_{i}\in P\right]$.

In the CATE estimation this amounts to \begin{align} \rho_i &= A_P^{-1}\left(W_i-\bar{W}_P\right)\left(Y_i-\bar{Y}_P-\left(W_i-\bar{W}_P\right)\hat{\beta}_P\right) \\ A_P &= \frac{1}{\left| \{i:X_i\in P\} \right|} \left[ \sum_{\{i:X_i\in P\}} \left(W_i-\bar{W}_P\right)\left(W_i-\bar{W}_P\right)^T \right] \end{align}

My question is: why? Why is $A_p$ the gradient of the local expectation score function $\mathbb{E}[\psi\mid X_i=x]$? How can $A_p$ be defined as a gradient, when the $\mathbb{E}[\psi\mid X_i=x]$ is in fact a matrix of vector?


1 Answer 1


Equation (20) is the locally centered version of the gradient

\begin{align} A_P &= \frac{1}{\left| \{i:X_i\in P\} \right|} \left[ \sum_{\{i:X_i\in P\}} \left(\tilde{W_i}\right)\left(\tilde{W_i}\right)^T \right] \end{align}

where $\tilde{W_{i}}$ is a $(q + 1) \times 1$ vector that has a 1 in the first position and the vector $W_{i}$ in the next $q$ positions. I explain in detail below.

Gradients and Jacobians

I'll follow the gradient conventions laid out in the Mathematical Appendix of Patterns, Predictions and Actions by Hardt and Recht (2022).

  • Let $\Phi: \mathbb{R}^{d} \to \mathbb{R}$. The gradient of $\Phi(x)$ with respect to the variable $x$ evaluated at a value $w \in \mathbb{R}^{d}$ is the vector of partial derivatives

$$ \nabla_{x} \Phi(w) =\left[\begin{array}{c} \frac{\partial \Phi(w)}{\partial x_1}\\ \vdots\\ \frac{\partial \Phi(w}{\partial x_d} \end{array}\right]_{d\times1}.$$

  • Let $\Phi: \mathbb{R}^{n} \to \mathbb{R}^{m}$ a multivariate mapping. The Jacobian of $\Phi(x)$ with respect to the variable x evaluated at a value $w \in \mathbb{R}^{d}$ is the $m \times n$ matrix

$$ D_{x}\Phi(w) = \left[ \frac{\partial\Phi_{i}(w)}{\partial x_{j}} \right]_{(i=1,...,m),(j=1,...,n)} $$

Estimating conditional average partial effects

We observe samples $(X_{i},Y_{i},W_{i}) \in \mathcal{X} \times \mathbb{R} \times \mathbb{R}^{q}$ and our conditional parameters of interest are $\beta:\mathcal{X} \to \mathbb{R}^{q}$ and $c:\mathcal{X} \to \mathbb{R}$. The scoring function $\Psi: \mathbb{R}^{q+1} \mapsto \mathbb{R}^{q+1}$ is

\begin{align} \Psi_{\beta(x),c(x)}(Y_{i},W_{i})^{\color{green}{[(q+1) \times 1]}} &= (Y_{i}^{\color{green}{[1 \times 1]}} - {W_{i}^{T}}^{\color{green}{[1 \times q]}}\beta(x)^{\color{green}{[q \times 1]}}-c(x)^{\color{green}{[1 \times 1]}})^{\color{green}{[1 \times 1]}}{[\quad 1^{\color{green}{[1 \times 1]}} \quad {W_{i}^{T}}^{\color{green}{[1 \times q]}} \quad ]^{T}} ^{\color{green}{[(q+1) \times 1]}} \\ &= (Y_{i} - {W_{i}^{T}}\beta(x)-c(x))^{\color{green}{[1 \times 1]}} \left[\begin{array}{c} 1 \\ W_{i1} \\ W_{i2} \\ \vdots \\ W_{iq} \end{array}\right]^{\color{green}{[(q+1) \times 1]}} \end{align}

In the strictest sense, given this setup, differentiation of $\Psi$ is only well defined for its arguments $(Y_{i},W_{i})$. However, we are interested in relating changes in the conditional parameters $(\beta(x),c(x))$ to changes in the scoring function value, so we will be taking partial derivatives with respect to the parameters rather than the arguments. It might help to think about a modified mapping $\tilde{\Psi}:\mathbb{R}^{q+1} \mapsto \mathbb{R}^{q+1}$ of the form $\tilde{\Psi}_{Y_{i},W_{i}}(\beta(x),c(x))$ where $(Y_{i},W_{i})$ are parameters and $(\beta(x),c(x))$ are arguments.

Note that just like in the case of $\Psi$, $\tilde{\Psi}$ is a multivariate mapping and so speaking about its gradient makes no sense. I think that a good part of the OP confusion has to do with the fact that the object $\nabla_{(\beta(x),c(x))} \Psi_{\beta(x),c(x)}(Y_{i},W_{i})$ is described as a "gradient" in the paper. According to the convention described above, it is more accurately described as the Jacobian matrix $D_{(\beta(x),c(x))} \tilde{\Psi}_{Y_{i},W_{i}}(\beta(x),c(x))$ of dimensions $(q+1) \times (q+1)$.

With that out of the way, we can focus on calculating the Jacobian:

\begin{align} D_{(\beta(x),c(x))} \tilde{\Psi}_{Y_{i},W_{i}}(\beta(x),c(x)) &= \begin{bmatrix} \frac{\partial(Y_{i} - {W_{i}^{T}}\beta(x)-c(x))}{\partial c(x)} & \frac{\partial(Y_{i} - {W_{i}^{T}}\beta(x)-c(x))}{\partial\beta(x)_{1}} & \ldots & \frac{\partial(Y_{i} - {W_{i}^{T}}\beta(x)-c(x))}{\partial \beta(x)_{q}} \\ \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{i1})}{\partial c(x)} & \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{i1})}{\partial\beta(x)_{1}} & \ldots & \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{i1})}{\partial \beta(x)_{q}} \\ \vdots & \vdots & \ddots & \vdots \\ \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{iq})}{\partial c(x)} & \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{iq})}{\partial\beta(x)_{1}} & \ldots & \frac{\partial((Y_{i} - {W_{i}^{T}}\beta(x)-c(x))W_{iq})}{\partial \beta(x)_{q}} \end{bmatrix} \\ &= \begin{bmatrix} 1 & W_{i1} & \ldots & W_{iq} \\ W_{i1} & W_{i1}^{2} & \ldots & W_{iq}W_{i1} \\ \vdots & \vdots & \ddots & \vdots \\ W_{iq} & W_{i1}W_{iq} & \ldots & W_{iq}^{2} \end{bmatrix}\\ &= [1 \quad W_{i}] [ 1 \quad W_{i}]^{T}\\ &= \tilde{W_{i}} \tilde{W_{i}}^{T} \end{align}

Plugging this into equation (7) in the paper gives the result stated at the beginning of this answer.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.