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

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?

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.

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.