Variable Importance Measure for Heterogeneous Treatment Effect
Posted on
VIM for heterogeneous treatment effects
- introduce nonparameteric treatment effect variable importance measures (TE-VIMs), based on the MSE in predicting the individual treatment effect.
- TE-VIMs represent the increase in MSE when variables are removed from the CATE conditioning set.
CATE:
\[\tau(x) = E(Y^1 - Y^0\mid X=x)\]- one pioneering proposal for CATE variable attribution: extend random forest variable importance measures (VIMs) through the causal forest estimator of the CATE
- VIMs based on specific modeling strategies (decision tree, linear regression, etc.) are referred to as “algorithmic”
- While algorithmic-VIMs can provide useful insights, there remains a need for model-agnostic alternatives, there remains a need for model-agnostic alternatives, especially when the chosen CATE estimating strategy does not have a well-established algorithmic-VIM
- propose treatment effect variable importance measures (TE-VIMs) that are nonparameteric summary statistics, which measure the importance of variable subsets in predicting the individual treatment effect (ITE) $Y^1 - Y^0$
consider the mean-sqaured-error
\[L(f) = E[(Y^1-Y^0-f(X))^2]\]one key insight
\[L(f) = E[(\tau(X) - f(X))^2] + E[\var(Y^1-Y^0\mid X)]\]define the TE-VIM
\[\Theta_s = L(\tau_s) - L(\tau)\]where
\[\tau_s(x) = E(Y^1-Y^0\mid X_{-s} = x_{-s})\]is the CATE conditional on $X_{-s}$
We have \(\Theta_s = E[(\tau(X) - \tau_s(X))^2]\,.\)
Since \(\tau_s(X) = E[\tau(X)\mid X_{-s}]\)
Thus \(\Theta_s = E[\Var(\tau(X)\mid X_{-s})]\)
Then by total variance decomposition, \(\Var(\tau(X)) = \Var(\tau_s(X)) + E[\Var(\tau(X)\mid X_{-s})]\)
which implies \(\Theta_s = \Var(\tau(X)) - \Var(\tau_s(X))\)
the variance $\var(\tau(X))$ of the treatment effect (VTE), a global measure of treatment effect heterogeneity
- quantifies the treatment effect heterogeneity explained by $X_s$ beyond that already explained by $X_{-s}$
TE-VIMs connects to regression-VIMs, also called “leave-out covariates”
2 Methodology
2.1 Motivating the estimand
- $\tau_p$: the average treatment effect (ATE)
- $\Theta_p$: VTE
personally speaking, i felt we can write $\tau_{[p]}$ to avoid misleading
assuming
- consistency: $A=a \Rightarrow Y=Y^a$
- conditional exchangeability: $Y^a\ind A\mid X$ for $a=0, 1$
- positivity: $0 < \pi(x) < 1$ almost surely
the CATE is identified by
\[\tau(x) = \mu(1, x) - \mu(0, x)\]where
\[\mu(a, x) = E(Y\mid A=a, X=x)\]and $\pi(x) = E(A\mid X=x)$ is the propensity score.
Assumeing $\Vert \tau(X)\Vert <\infty$, then the TE-VIM
\[\Theta_s = \Vert \tau(X) - \tau_s(X)\Vert^2 < \infty\]is also finite.
the TE-VIM $\Theta_s$ is analogous to the regression-VIM $\Omega_s = \Vert Y-\mu_s(X)\Vert^2 - \Vert Y-\mu(X)\Vert^2$
propose scaled TE-VIMs as
\[\Psi_s = \Theta_s / \Theta_p = 1-\Var(\tau_s(X)) / \Theta_p\]modes:
- leave-one-out
- keep-one-in
- shapley values
- covariate grouping
2.2 CATE estimation
- metalearning strategies decompose CATE estimation into a sequence of subregression problems, which can be solved using off-the-shelf ML algorithms
- 2 metalearning algorithms: T-learner and DR-learner
- T-learner is based on the decomposition \(\tau(x) = \mu(1, x) - \mu(0, x)\) and estimate the CATE by \(\hat\tau^{(T)}(x) = \hat\mu(1, x) - \hat\mu(0, x)\)
- the T-learner is problematic for 2 main reasons:
-
- while regularization methods can be used to control the smoothness of $\hat\mu(a, x)$, the same is not true of $\hat\tau^{(T)}(x)$ which may be erratic. Slow convergence rate affecting $\hat\mu(a, x)$ may therefore propagate into $\hat\tau^{(T)}(x)$.
-
- $\hat\mu(1, x)$ is chosen to make an optimal bias-variance trade-off over the covariance distribution of the treated population. Likewise, $\hat\mu(0, x)$ is chosen to make an optimal bias-variance trade-off over the covariate distribution of the untreated population. When there is poor overlap between the treated and untreated subgroups, then $\hat\tau^{(T)}(x)$ may fail to deliver an optimal bias-variance trade-off over the population covariate distribution, making the T-learner potentially poorly targeted toward CATE estimation.
-
- the DR-learner is an alternative metalearning algorithm based on the decomposition \(\tau(x) = E[\phi(Z)\mid X=x]\) where, for $z = (y, a, x)$ \(\phi(z) = [y-\mu(a, x)]\frac{a-\pi(x)}{\pi(x)(1-\pi(x))}+\mu(1, x) - \mu(0, x)\) is called the pseudo outcome, or the augmented inverse propensity weighted (AIPW) score, and acts like the ITE in expectation.
- the DR-learner first estimates $\mu(a, x)$ and $\pi(x)$ to obtain the pseudo-outcome estiator $\hat\phi(Z)$, where $\mu(a, x)$ and $\pi(x)$ are replaced with estimates $\hat\mu(a, x)$ and $\hat\pi(x)$
- in a second step, the estimated pseudo-outcome, $\hat\phi(Z)$ is regressed on covariates $X$ to obtain $\hat\tau^{(DR)}(x)$.
- The DR-learner alleviates the issues related to the T-learner since the complexity of $\hat\tau^{(DR)}(x)$ can be controlled by regularizing the regression in the final stage of the procedure.
- one can trade-off accuracy in the outcome and propensity score estimators, a property known as double robustness, hence the name “DR-learner”
- estimation of the CATE $\tau_s(x)$ is complicated by the fact that one cannot assume that $Y\ind A\mid X_{-s}$ for an arbitrary subset of covariates $s$, a problem that is sometimes referred to as “runtime confounding”
- the DR-learner readily accommodates runtime confounding through the identity $\tau_s(x) = E[\phi(Z)\mid X_{-s}=x_{-s}]$. The identity implies that one may estimate $\tau_s(x)$ by regressing $\hat\phi(Z)$ on $X_{-s}$, that is, modifying the final regression step of the DR-learner.
2.3 TE-VIM estimation
2.3.1 Estimation of $\Theta_s$
consider estimators based on the efficient influence curve (IC) of $\Theta_s$ under the nonparameteric model.
ICs are mean 0 functions that characterize the sensitivity of so-called pathwise differentiable estimands to small changes in the data generating law.
ICs are useful for constructing efficient estimators and determining their asymptotic distribution
since the risk $L[f]$ is not identified. Consider that
\[\Theta_s = L_U[\tau_s] - L_U[\tau]\]where
\[L_U[f]= E[(U+\tau(X) - f(X))^2] = E[(\tau(X) - f(X))^2] + E[U^2]\]and $U$ is a random variable such that $E[U\mid X] = 0$. Note that $U = Y^1 - Y^0 - \tau(X)$ recovers the unidentified risk $L(f)$ and $U = \phi(Z)-\tau(X)$ recovers the DR-learner population risk
consider the risk $L^\star(f)$ obtained by setting $U=0$.
the IC of $L^\star(f)$ at a single observation $z$ is
\[[\phi(z) - f(x)]^2 - [\phi(z) - \tau(x)]^2 -L^\star[f]\]the IC of $\Theta_s$ is
\[\phi_s(z) = [\phi(z) - \tau(x)]^2 - [\phi(z) - \tau(x)]^2 - \Theta_s\]the interpretation of $\phi(Z)$ as a pseudo-outcome, which plays the role of the unobserved ITE $Y^1 - Y^0$
2.3.2 Importance testing
$H_0: \Theta_s = 0$
in which case $\phi_s(Z) = 0$. This IC degeneracy makes $H_0$ difficult to test, since $\hat\Theta_s$ is not necessarily asymptotically normal.
one solution to the IC degeneracy problem is to estimate $\var(\tau(X))$ and $\var(\tau_s(X))$ using efficient estimators in separate samples
2.3.3 Estimation of $\Psi_s$
the scaled TE-VIM
\[\Psi_s = \Theta_s / \Theta_p\]hac IC
\[\Phi_s(z) = (\phi_s(z) - \Psi_s\phi_p(z)) / \Theta_p\]The estimator $\hat\Theta_p$ and $\hat\Psi_s$ both rely on an ATE estimator $\hat\tau_p$.
The IC of the ATE, $\phi(z) - \tau_p$, implies an efficient estimator
\[\hat\tau_p^\star = n^{-1}\sum_{i=1}^n \hat\phi(z_i)\]known as the AIPW estimator
2.3.4 Plug-in estimation
a common framework for constructing debiased estimators is through targeted maximum likelihood estimators (TMLEs)