How does Generalized Random Forest calculate the gradient of the score function?

Question

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?

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.