Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
Rendered from LaTeX for readability, not typeset faithfully. Citation keys are highlighted; maths is left as source; figures, tables and equation environments are summarised rather than reproduced; unrecognised commands are greyed out so nothing is silently dropped. Email addresses are removed.
On the adaptation of causal forests to manifold data
{5pt}
{5pt}
{5pt}
{5pt}
\hypersetup{colorlinks,breaklinks,urlcolor=blue,linkcolor=blue}
abstractResearchers often hold the belief that random forests are “the cure to the world's ills” bickel2010leo. But how exactly do they achieve this?
Focused on the recently introduced causal forests athey2016recursive,wager2018estimation, this manuscript aims to contribute to an ongoing research trend towards answering this question, proving that causal forests can adapt to the unknown covariate manifold structure. In particular, our analysis shows, for the first time, that a causal forest estimator can achieve the optimal rate of convergence for estimating the conditional average treatment effect, with the covariate dimension automatically replaced by the manifold dimension. These findings align with analogous observations in the realm of deep learning and resonate with the insights presented in Peter Bickel's 2004 Rietz lecture.
{\bf Keywords}: random forests, causal forests, manifold data, adaptation.
Introduction
Suppose that we observe a binary treatment data. It contains $n$ observations of $(X, D, Y(D))$ from an independent and identically distributed sample $\{(X_i, D_i,Y_i(0), Y_i(1))\}_{i=1}^n$ of $\{X, D, Y(0), Y(1)\}$, with
$$
X_i\in\mathcal{X}\subset\mathbb{R}^d, ~~D_i\in\{0,1\},~~ {\rm and}~~ (Y_i(0), Y_i(1))\in\mathbb{R}^2
$$
as the potential outcomes neyman1923applications,rubin1974estimating. The goal of interest is to, based on such a data, infer the conditional average treatment effect (CATE, athey2016recursive,athey2019machine):
align[align omitted — 122 chars of source]
Causal inference towards estimating $\tau(x)$ often invokes some nonparametric/semiparametric regression procedures van2006targeted,wager2018estimation,kunzel2019metalearners,kennedy2020towards,knaus2021machine,nie2021quasi. To justify these methods, besides the commonly imposed causality assumptions of unconfoundedness and overlap imbens2015causal, it is often to assume that:
itemize• {\bf distribution assumption}: the Lebesgue density and conditional density of $X$ given $D$ are (locally) “sufficiently regular”;
• {\bf regression assumption}: the regression functions $\mu_\omega(x):={\mathrm E}[Y~|~X=x,D=\omega]$, for $\omega=0$ and $1$, are (locally) “sufficiently smooth”.
This paper aims to challenge the distribution assumption {\bf (D)} and proposes to replace it with the following covariate manifold assumption:
itemize• {\bf manifold assumption}: the distribution of $X$ lives on a low-dimensional manifold and its Hausdorff (conditional) density is (locally) “sufficiently regular”.
There are three main motivations to consider the assumption {\bf (M)} in contrast to {\bf (D)}. Foremost and of paramount importance, in practical terms, instances of data exhibiting intrinsic low-dimensional structures seem to be ubiquitous. If we embrace this perspective, it becomes judicious to postulate a manifold assumption when scrutinizing statistical methodologies related to such data.
Second, from a methodological point of view, although the mathematical statistics literature has long been aware of the phenomenon that “local methods” can adapt to the covariate manifold structure bickel2007local,kpotufe2010curse,
the causal inference literature—except for a notable work of khosravi2019non—is largely silent about such phenomena. Discussions that address the mitigation of the curse of dimensionality through local, in contrast to global, causal methods could prove beneficial in this regard.
Lastly, from a theoretical standpoint, since Peter Bickel's seminal 2004 Rietz lecture, there has been a surge of theoretical interest in understanding adaption to manifold structures. Noteworthy contributions in this realm include bickel2007local and cheng2013local on local linear regression, kpotufe2011k on nearest neighbor regression, kpotufe2013adaptivity on kernel regression, yang2015minimax on Gaussian process regression, kpotufe2012tree on tree methods, liao2021multiscale on multiscale regression, and schmidt2019deep, chen2022nonparametric, and jiao2023deep on deep learning methods.
Interestingly, amidst these advancements, discussions regarding the adaptation of random forests to manifold structures—albeit alluded to in various contexts (cf. biau2016random and references therein)—still remain notably absent. Our work thus also complements the studies of biau2008consistency, biau2012analysis, louppe2013understanding, scornet2015consistency, zhu2015reinforcement, oprescu2019orthogonal, syrgkanis2020estimation, mourtada2020minimax, klusowski2021sharp, cattaneo2022pointwise, chi2022asymptotic, among many other, which are focused on adaptation of random forests methods to different types of regression models.
Specifically, we provide a statistical analysis of the causal forest method of lin2022regression—a slight modification to the causal forests introduced in athey2016recursive and wager2018estimation—based on the assumptions {\bf (R)} and {\bf (M)}. We explore pointwise rates of convergence and central limit theorems for the causal forest estimator $\hat\tau(x)$ of $\tau(x)$, showing that
itemize• the pointwise mean squared error (MSE) could achieve the minimax risk $n^{-2/(m+2)}$, with $m$ representing the manifold dimension in contrast to the ambient space dimension;
• a rescaled version of $\hat\tau(x)$ could admit a central limit theorem towards inferring $\tau(x)$, with the asymptotic variance consistently estimable;
• all the theoretical outcomes are established within a broader framework that encompasses a variety of nonparametric methods beyond the scope of the causal forest approach.
The focus of this paper is on the theoretical side—about the CATE estimation based on causal forests under a manifold setting—rather than developing new algorithms expressly crafted for manifold data. Our intent is to expand the already extensively explored theory in mathematical statistics to encompass random forests-based methods, and to connect them to causal inference. Through a particular version of the CATE estimator (i.e. the causal forest), our analysis suggests how precise a theoretical formulation can be made and used for justifying other local methods.
{\bf Paper organization.} The rest of this manuscript is organized as follows. Section (ref) introduces the setup of the problem and the studied causal forest estimator. Section (ref) presents the main theory of this manuscript, highlighting the adaption of causal forests to manifold structure based on the criteria of consistency, rate of convergence, and central limit theorem. Section (ref) lifts everything in Section (ref) to a general framework, covering potentially many other local methods. All proofs are relegated to the appendix.
Setup
We consider the setting where there are $n$ studied units, indexed by $i\in[n]:=\{1,\ldots,n\}$. Each unit receives a binary treatment, indicated by $D_i\in\{0,1\}$, with $D_i=1$ or $D_i=0$ demonstrating that the $i$-th unit is in the treatment or control group, respectively. Following the Neyman-Rubin causal model, each unit then has a pair of potential outcomes, $(Y_i(0), Y_i(1))$. The realized outcome for the $i$-th unit is denoted by
$$
Y_i = Y_i(D_i)=
casesY_i(0), & if D_i=0,\\
Y_i(1), & if D_i=1.
$$
Let $X_i\in\cX\subset\mathbb{R}^d$ represent a vector of pretreatment variables, not affected by the treatment. The observed data then is $\{(X_i,D_i,Y_i)\}_{i=1}^n$, which we believe to be exchangeable and, especially, contain no interference.
In order to estimate the CATE $\tau(x)$ in (ref), we are interested in a causal forest estimator first proposed in lin2022regression.
To begin with, consider a tree estimator as athey2016recursive. Let two {\it generic} trees, denoted by $T^1$ and $T^0$, be constructed for the treated and control groups, $\{(X_i,Y_i)\}_{i=1,D_i=1}^n$ and $\{(X_i,Y_i)\}_{i=1,D_i=0}^n$, respectively. Let $L^1$ and $L^0$ represent the sets of leaves generated by the two trees; they each represent a partition of the covariate space $\cX$.
The corresponding tree regression estimator of the {\it unobserved} potential outcome at the unit-$i$ is then defined as follows: it is the sample average of outcomes $Y_j(D_j)$ of individuals whose covariates $X_j$ are in the same leaf as $X_i$ but receiving the opposite treatment (i.e., $D_j=1-D_i$).
Next, we grow a forest from such trees through {\it subsampling} breiman2001random. Let's first fix two parameters. Set $s$ to represent the subsample size and set $B$ to represent the number of replications. In each round of subsampling, indexed by $b\in[B]$, let ${\mathcal I}_b^\omega$ ($\omega=0,1$) stand for a size-$s$ set being randomly sampled without replacement from $\{i:D_i=\omega\}$. For each round $b\in[B]$, one could then grow two causal trees, $T_b^1$ and $T_b^0$, from $\{(X_i,Y_i)\}_{i\in{\mathcal I}_b^1}$ and $\{(X_i,Y_i)\}_{i\in{\mathcal I}_b^0}$, respectively, leading to $2B$ many trees.
The {\it forest} then aggregates all $2B$ trees by averaging. For each $i\in[n]$, the final forest estimator of $(Y_i(0),Y_i(1))$, combined with a regression adjustment step due to rubin1973use and abadie2011bias, is defined to be
align*[align* omitted — 372 chars of source]
and
align*[align* omitted — 372 chars of source]
where for any set $A$, $|A|$ represents its cardinality; for any $\omega\in\{0,1\}$, $b\in[B]$, and $x\in\cX$, $L_b^\omega(x)$ stands for the leaf of $T_b^\omega$ that contains $x$; and $\hat\mu_\omega(\cdot)$ is a {\it regression adjustment function} allowed to be set 0 in estimating the CATE.
Lastly, the estimator for $\tau(x)$ is articulated as a “kernel smoothed average” of $(\hat Y_i^{\rm RF}(1)-\hat Y^{\rm RF}_i(0))$ values corresponding to covariates proximate to $x$. Delving into the specifics, let $K_h(\cdot):\mathbb{R}^d\to\mathbb{R}$ represent a {\it multivariate kernel function} incorporating a user-defined tuning parameter $h\in \mathbb R^{> 0}$. The ultimate causal forest estimator is
align[align omitted — 150 chars of source]
remarkA notable distinction between the investigated estimator in (ref) and the original causal tree and forest estimators introduced in athey2016recursive and wager2018estimation lies in the construction process. The former entails the creation of two distinct trees, in line with the methodology proposed in lin2022regression. In contrast, the latter approach involves the development of a single tree that mixes both treatment and control data. While we conjecture that this divergence may not be pivotal, the separation approach significantly streamlines subsequent analysis and facilitates intricate computations to be developed in Section (ref). Consequently, we adhere to this separation methodology throughout this manuscript, while acknowledging that, in practical applications, adopting a unified tree-building strategy as presented in athey2016recursive and wager2018estimation could potentially suffice.
remarkThe estimator in (ref) further diverges from the causal forest estimator introduced in wager2018estimation due to the incorporation of kernel smoothing. In its initial form, aside from the distinction of constructing trees using combined data versus two distinct groups, the causal forest employs a straightforward difference, $\hat Y_i^{\rm RF}(1)-\hat Y_i^{\rm RF}(0)$, to approximate $\tau(X_i)$. Our approach, on the other hand, is underpinned by lin2022regression, which posits that (ref) can also yield a natural estimator for the average treatment effect (ATE). As a theoretical exposition, we refrain from engaging in a direct comparison between these two variants, instead concentrating our efforts on the adaptation of such methodologies to manifold structures.
remarkIn their seminal paper, athey2019generalized introduced the concept of “local-centering”, a notion advocating the utilization of residuals from outcome and treatment variables—rather than original values—as data inputs when implementing a causal forest algorithm. This approach aligns with the “residual-on-residual” strategy robinson1988root and the paradigm of double machine learning chernozhukov2018double, while also exhibiting a fundamental connection to the concept of “double robustness” robins1994estimation, bang2005doubly. In a series of recent studies, lin2021estimation and lin2022regression demonstrated that, in the context of ATE estimation, the advantage of double robustness can also be harnessed by incorporating a regression adjustment approach devised in rubin1973use and abadie2011bias, with $\hat\mu_w(\cdot)$ designed to approximate the conditional expectation $\mu_w(x):={\mathrm E}[Y|X=x,D=w]$. In Section (ref), we will delve deeply into an enhanced iteration of (ref), which discusses this supplementary regression adjustment step in more details.
Theory
To present the theory, we first introduce the assumptions that regulate the data generating process and permit causal identification of $\tau(x)$ for {\it any} $x\in\cX$.
assumption\phantomsection
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• $[(X_i,D_i,Y_i)]_{i=1}^n$ are independent and identically distributed (i.i.d.) following the joint distribution of $(X,D,Y(D))$.
• For almost all $x \in \cX$, $D$ is independent of $(Y(0),Y(1))$ conditional on $X=x$, and there exists a fixed constant $\eta > 0$ such that $\eta < {\mathrm P}(D=1 \,|\, X=x) < 1-\eta$.
• Letting $\mu_\omega(x):={\mathrm E}(Y(\omega)|X=x, D=\omega)$ and $U_\omega(x):=Y(\omega)-\mu_\omega(x)$, it is assumed that $U_\omega(\cdot)$ is uniformly bounded and $\mu_\omega(\cdot)$ is bounded continuous for all $x \in \cX$ and $\omega \in \{0,1\}$.
\end{enumerate}
remarkAssumption (ref)(ref) is the standard i.i.d.-ness condition commonly adopted in observational studies. Assumption (ref)(ref) is the unconfoundedness and overlap conditions for identifying the causal effect. Assumption (ref)(ref) regulates the regression function and its residual, constituting Assumption {\bf (R)} presented in Section (ref).
The next set of assumptions regulates the distribution of $X$, constituting Assumption {\bf (M)} presented in Section (ref).
assumption\phantomsection
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• It is assumed that $X \in \cM$, where $\cM$ is an $m$-dimensional $C^{\infty}$ manifold in $\mathbb{R}^d$ with the manifold dimension $m \leq d$.
\itemThe marginal law of $X$, denoted by $\zeta$, is assumed to be absolutely continuous with respect to the restricted m-dimensional Hausdorff measure in $\mathbb{R}^d$ on $\cM$ gariepy2015measure, denoted by $\mathcal{H}^m$.
\itemThe support of $\zeta$, denoted by $\operatorname{supp}(\zeta)$, is assumed to be compact.
• Write $f:={\mathrm d} \zeta / {\mathrm d} \cH^m$ to be the Radon-Nikodym derivative of $\zeta$ with respect to $\cH^m$.
We assume that $f(\cdot)$ is bounded and bounded away from zero for any $x\in\operatorname{supp}(\zeta)$.
• Let the law of $X|D=\omega$ for $\omega\in\{0,1\}$ be denoted by $\zeta_\omega$, and the corresponding restriction to a set $U$ be denoted by $\zeta_{\omega, U}$. For any point $x \in \operatorname{supp}(\zeta)$ and any chart $(U, \psi)$ such that $U=U_x$ is a coordinate neighborhood of $x$ and $\psi: U \rightarrow$ $V \subset \mathbb{R}^m$, let $g_{\omega,x}=g_{\omega,U_x}:=\mathrm{d}\left(\psi_* \zeta_{\omega,U}\right) / \mathrm{d} \lambda$ be the Radon-Nikodym derivative of the restricted pushforward measure $\psi_* \zeta_{\omega,U}$ with respect to the Lebesgue measure $\lambda$. It is assumed that $g_{\omega,x}(\cdot)$ is locally Lipschitz.
\end{enumerate}
remarkAssumption (ref)(ref) regulates the random variable $X$ to be resided on an $m$-dimensional smooth manifold. Assumptions (ref)(ref) and (ref)(ref) require the Hausdorff density to exist and being compactly supported. Assumptions (ref)(ref) and (ref)(ref) further require the density to be “sufficiently regular”, namely, bounded and bounded away from 0 and (locally) Lipschitz smooth.
remarkTwo additional remarks are worth noting. First, when $\cM$ is set to be the $d$-dimensional space without any more structure, Assumption (ref) simplifies to the commonly employed density assumptions as seen in, e.g., wager2018estimation. Second, although in Assumption (ref) and at various other places ahead we formulate “global” assumptions that need to apply to {\it any} $x\in\cM$, it is straightforward to adapt these assumptions to a {\it local} context by considering those within a small neighborhood of a particular $x$. This is by confining our investigation to a localized subset of the manifold.
The next set of assumptions regulates the kernel smoothing function used in (ref).
assumption\phantomsection
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• We define $K_h(x)=h^{-d/2}K(\|h^{-1/2}x\|)$, with $\|\cdot\|$ representing the Euclidean norm and $K(\cdot):\mathbb{R}^{\geq 0}\to \mathbb{R}^{\geq 0}$ assumed to be of a bounded support, $\operatorname{supp}(K)$, i.e, the diameter of $\operatorname{supp}(K)$, denoted by $d_K$, is finite.
• We assume $K$ is Lipchitz over $\operatorname{supp}(K)$.
• We assume $h \to 0$ and $n h^{m/2} \to \infty$.
\end{enumerate}
remarkAssumptions (ref)(ref) and (ref)(ref) are added for simplifying the subsequent theoretical analysis. They are satisfied by using, e.g., the box kernel or a truncated Gaussian kernel. Lastly, Assumption (ref)(ref) is easily satisfied by choosing not too small an $h$.
In the end, we regulate the causal forest algorithm introduced in Section (ref). In the following, for any real sequences $\{a_n\}, \{b_n\}$, we adopt the notation “$a_n\lesssim b_n$” to mean $|a_n|\leq C|b_n|$ for all sufficiently large $n$. We write $a_n\asymp b_n$ if both $a_n\lesssim b_n$ and $b_n\lesssim a_n$ hold.
assumption\phantomsection
We assume that
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• $s=s_n \rightarrow \infty$ and $n/B=O(1)$;
• there exists a constant $ \epsilon \in(0, 1/(m+2))$ such that for any tree $T$ considered in Section (ref) and any leaf $L_t$ in $T$,
$$
\Big\lVert{\rm diam}(L_t\cap\cM)\Big\rVert_\infty \lesssim h^{1/2+\epsilon},~~{\rm uniformly},
$$
where ${\rm diam}(\cdot)$ outputs the diameter of the input set and $\|\cdot\|_{\infty}$ stands for the $L^{\infty}$ norm.
\end{enumerate}
assumption\phantomsection
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• We assume that for the same $\epsilon$ in Assumption (ref), there exists a constant $\beta>0$ such that the subsample size $s$ satisfies $\lim _{n \rightarrow \infty} s^3 h^{\beta}=\infty$, and for $\omega\in\{0,1\}$,
$$
\Big\lVert {\mathrm E}\Big[\rvert L^{\omega}(X_1) \rvert^{-\beta}\,\Big{|}\, D_1=\omega, D_2,\ldots,D_n,X_1,1\in{\mathcal I}^\omega\Big] \Big\lVert_\infty \lesssim s^{-2-\beta/2} (\log s)^{-\beta/2}h^{(2+\beta)\epsilon},
$$
where $\lvert L \rvert$ represents the number of samples in a leaf $L$ and the subscript $b$ in both $L^\omega$ and ${\mathcal I}^\omega$ is hidden.
• The tree is honest, that is, the partition is independent of $Y_i$'s.
\end{enumerate}
assumption\phantomsection
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• We assume that the leaf size satisfies that for $\omega\in\{0,1\}$,
$$
\Big\lVert {\mathrm E}\Big[\Big(\zeta_\omega(L^\omega(X_1)\cap\cM)\Big)^{-1}\,\Big{|}\, D_1=\omega, D_2,\ldots,D_n,X_1, 1 \in {\mathcal I}^{\omega}\Big]\Big\rVert_\infty
\lesssim sh^{2\epsilon}.
$$
• The tree is extremely honest, that is, the partition is independent of the data.
\end{enumerate}
remarkAssumption (ref)(ref) weakens the first part of Assumption 5.1(ii) in lin2022regression about the subsample size $s$. Assumption (ref)(ref) as well as Assumptions (ref) and (ref) enforce different decaying rates for the diameter of $L_t$, the number of samples in each leaf, and the subsample size $s$. They are stronger than the corresponding requirements made in Assumptions 5.1 and 5.3 of lin2022regression, but are necessary for studying CATE, instead of ATE, estimation. These assumptions can be satisfied by properly trimming the trees while growing the forest; see also Lemma 5.1 and Proposition 5.1 in lin2022regression. Assumption (ref)(ref) and Assumption (ref)(ref) are the famous “honesty” wager2018estimation and “purely random forests” conditions breiman2004consistency,biau2008consistency,biau2012analysis, both of which can be satisfied by, e.g., implementing sample splitting or employing a pre-trained random forest model.
remarkBy imposing a stronger independence condition between the partitions and the data, as compared to Assumption (ref), Assumption (ref) effectively eases the constraints on the leaf size, $|L_t|$. This relaxation plays a pivotal role in the subsequent derivation of a minimax optimal result for $\hat\tau_{\rm RF}(x)$. It is noteworthy that we employ the “extremely honest” condition primarily for its technical convenience, as elaborated in, for instance, the discussion presented in biau2016random.
remarkAs a matter of fact, under extremely honest condition, Assumptions (ref)(ref) and (ref)(ref) can be replaced by
\begin{align*}
\Big\lVert \max_t{\rm diam}(L_t\cap\cM) \Big\rVert_{4m} \lesssim h^{1/2+4 \epsilon} and \Big\lVert \Big(\max_t \zeta_{\omega}(L_t^{\omega}\cap\cM)\Big)^{-1} \Big\rVert_{4} \lesssim sh^{2\epsilon},
\end{align*}
for $\omega \in \{0,1\}$. Here $\lVert\cdot\rVert_{p}$ stands for the $L^{p}$ norm. This is by modifying the proof of Lemmas (ref)-(ref) in the appendix and conditioning on $\{L^{\omega}_{bt}\}_{t \ge 1}$ instead of just on $\{D_j\}_{j=1}^n$ and $X_i$. However, in practice, this difference appears to be not huge, and accordingly we choose to present them at their current forms.
With the above assumptions held, we are then ready to introduce the first theorem about the causal forest estimator $\hat\tau_{\rm RF}$.
theorem[Consistency]
Assume Assumptions (ref)-(ref) hold, with either Assumption (ref) or Assumption (ref), and in (ref) set $\hat\mu_0=\hat\mu_1=0$. We then have, for any interior point $x\in\operatorname{supp}(\zeta)$,
$$
\hat\tau_{\rm RF}(x)-\tau(x)\stackrel{\sf p}{\longrightarrow} 0.
$$
Next, we establish the rate of convergence for $\hat\tau_{\rm RF}(x)$ approximating $\tau(x)$ based on the criterion of the MSE. To this end, some additional assumptions on $\mu_\omega$ and the kernel function are needed.
assumption\phantomsection For $\omega \in \{0,1\}$, assume $\mu_\omega$ to be Lipchitz over $\operatorname{supp}(\zeta)$.
assumption\phantomsection
We assume $K$ is bounded away form zero over $\operatorname{supp}(K)$.
remarkAssumption (ref) is a stronger version of the continuity condition in Assumption (ref)(ref) and requires Lipschitz continuity of the regression functions. This type of smoothness is routinely posed; cf. wager2018estimation.
theorem[MSE]
Assume Assumptions (ref)-(ref), (ref)-(ref) hold, with either Assumption (ref) or Assumption (ref) held, and in (ref) set $\hat\mu_0=\hat\mu_1=0$. We then have, for any interior point $x\in\operatorname{supp}(\zeta)$,
$$
{\rm MSE}(\hat\tau_{\rm RF}(x)):={\mathrm E}\Big[\Big(\hat\tau_{\rm RF}(x)-\tau(x)\Big)^2\Big]\lesssim h+\frac{1}{nh^{m/2}},
$$
In particular, when Assumption (ref) holds, and as we choose $h\asymp n^{-2/(m+2)}$,
it holds that the MSE admits the following minimax risk for estimating a Lipschitz function,
$$
{\rm MSE}(\hat\tau_{\rm RF}(x)) \lesssim n^{-2/(m+2)},
$$
with the ambient dimension $d$ replaced by the intrinsic manifold dimension $m$.
Of note, when Assumption (ref) holds, the rate in Theorem (ref) is {\it minimax optimal} for pointwisely estimating a Lipschitz smooth function yang1999information when the covariate resides in an $m$-dimensional space. The message conveyed here then is clear: causal forests are able to rate-optimally adapt to the unknown manifold structure.
If only an honest condition as made in Assumption (ref) holds, the situation is a little bit more complex. In detail,
if in Assumption (ref)(ref) we indeed have ${\rm diam}(L_t) \asymp h^{1/2+\epsilon}$, then by Assumption (ref)(ref), the displayed constraint in Assumption (ref)(ref) restricts that
\[
\frac{n^{1-4/\beta}}{\log(n)} h^{m+2(m+1)\epsilon+ 4\epsilon/\beta}\gtrsim 1,
\]
with $s \asymp n$. If we we set $\epsilon$ in Assumption (ref)(ref) to be sufficiently close to 0 and $\beta$ in Assumption (ref) to be sufficiently large, then when $m>1$, the above displayed constraint yields $h\gtrsim (nh^{m/2})^{-1}$ and thus an ${\rm MSE}(\hat\tau_{\rm RF}(x)) \lesssim h$, yielding an approximately “best” rate
$$
(\log n)^{1/(m+2(m+1)\epsilon+ 4\epsilon/\beta)}/n^{(1-4/\beta)/(m+2(m+1)\epsilon+ 4\epsilon/\beta)},
$$ which is apparently suboptimal. The phenomenon of obtaining a sub-optimal rate in analyzing even a simple random forest model (e.g., honest/pure ones) is again well known; see, for example, biau2012analysis and klusowski2021sharp. The fact that we can obtain an improved, minimax optimal, rate via adopting Assumption (ref), on the other hand, is surprising and suggests the importance of leveraging kernel smoothing in (ref).
Lastly, we consider the problem of statistically inferring $\tau(x)$ using $\hat\tau_{\rm RF}(x)$. This is done by {\it under-smoothing} $\hat\tau_{\rm RF}(x)$, as well as enforcing the following additional moment assumptions on the data generating distribution. Below, Assumption (ref) corresponds to Assumption 3.5 in lin2022regression, which is put in order to leverage the Lindberg-Feller central limit theorem for analyzing $\hat\tau_{\rm RF}(x)$. Assumption (ref) is added for facilitating a closed form of the asymptotic variance of $\hat\tau_{\rm RF}(x)$; cf. the term $\Sigma(x)$ in Equation (ref) ahead.
assumption\phantomsection
For any interior point $x \in \operatorname{supp}(\zeta)$ and $\omega \in \{0,1\}$, ${\mathrm E} [U^2_\omega(x) \,|\, X=x]$ is uniformly bounded away from zero.
assumption\phantomsection
We assume
$\sigma^2_\omega(x) := {\mathrm E} [U^2_\omega(x) \,|\, X=x] $ is Lipchitz over $\operatorname{supp}(\zeta)$ for $\omega \in \{0,1\}$.
theorem[Central limit theorem]
Assume Assumptions (ref)-(ref), (ref)-(ref) hold and in (ref) set $\hat\mu_0=\hat\mu_1=0$. Further, assume that $nh^{m/2+1}\to 0$. We then have, for any interior point $x \in \operatorname{supp}(\zeta)$,
\begin{align*}
\sqrt{nh^{m/2}} \left(\hat\tau_{\rm RF}\left(x\right) - \tau\left(x\right) \right) \stackrel{\sf d}{\longrightarrow} N(0,\Sigma(x)),
\end{align*}
where
\begin{align}
\Sigma(x):=\frac{1}{c_K^2 f(x)}\left(\frac{\sigma^2_1(x)}{e(x)} + \frac{\sigma^2_0(x)}{1-e(x)}\right)\int_{\mathbb{R}^m} K^2(\lVert t\rVert) {\mathrm d} t
\end{align}
with $c_K:=\int K(\|t\|){\mathrm d} t$ and $e(x):={\mathrm P}(D=1|X=x)$ representing the propensity score rosenbaum1983central.
remarkThe asymptotic variance $\Sigma(x)$ can be consistently estimated combining consistent estimators of the manifold density $f(x)$, the propensity score $e(x)$, and the residual variances $\sigma_\omega^2(x)$ for $\omega=0,1$. Here, consistent estimators of manifold densities can be found in, e.g., pelletier2005kernel and le2019approximation. Propensity score estimation is a classic regression fitting problem, and we refer the readers to imbens2015causal for a complimentary review. Lastly, estimation of the residual variance in a (nonparametric) regression model has been well studied in the literature; cf. wang2008effect and shen2020optimal and the references therein for details.
remarkIn Theorems (ref) and (ref) and various other places in Section (ref) ahead, for tuning the parameters in order to obtain the optimal rate or check the assumption validity, one needs to have some knowledge about the manifold dimension, $m$. While we agree with levina2004maximum that cross-validation should be able to automatically select the “optimal” tuning parameters, in theory one could use some of the existing methods for consistently estimating the manifold dimension; see, e.g., levina2004maximum, farahmand2007manifold, and block2022intrinsic.
remarkIn this section, all theoretical results are demonstrated by setting $\hat\mu_0=\hat\mu_1=0$, representing misspecified mean function models. Despite this, our causal forest estimator (ref) still enjoys consistency due to its nature as a doubly robust estimator, as will be formally demonstrated in Section (ref). Using random forests to impute missing potential outcomes automatically provides us with correctly specified odds ratios for the propensity score, a similar phenomenon as formally characterized in lin2021estimation for the nearest neighbor matching algorithm and in lin2022regression for more general cases.
To conclude this section, we make a brief observation regarding the manifold assumption itself. Although our paper primarily concentrates on scenarios where $X$ resides precisely within a smooth manifold, we posit that all the findings should, in principle, extend to cases where $X$ is distributed over an {\it approximate} manifold. This perspective aligns with existing work such as kpotufe2011k and jiao2023deep. Nevertheless, the elegant exact manifold assumption proves sufficient to fulfill our research objectives, and we avoid confusing our readers more by delving to more technical assumptions.
A general framework
This section aims to establish a general framework for analyzing the CATE estimator introduced in (ref), while bearing the potential to cover more. To this end, we adopt the framework introduced in lin2022regression, and consider such imputation estimators that belong to the family of {\it linear smoothers} buja1989linear.
Using the same notation as above, instead of imputing the missing potential outcomes as the forest-based ones, we consider a more general approach, defining
align*[align* omitted — 220 chars of source]
and
align*[align* omitted — 215 chars of source]
where $w_{i\leftarrow j}$'s are the {\it smoothing parameters} usually learnt from the data. The corresponding CATE estimator is then defined to be
$$
\hat\tau_{w}(x):=\frac{\sum_{i=1}^nK_h(X_i-x)(\hat Y_i(1)-\hat Y_i(0))}{\sum_{i=1}^nK_h(X_i-x)},
$$
with the subindex $w$ in $\hat\tau_{w}(x)$ highlighting the dependence of the general CATE on the smoothing parameters. In particular, when setting
$$
w_{i\leftarrow j}:= \frac1B \sum_{b=1}^B \frac{\ind\Big(j \in {\mathcal I}^{1-D_i}_b:X_j \in L^{1-D_i}_b(X_i)\Big)}{\Big\lvert \Big\{k \in {\mathcal I}^{1-D_i}_b:X_k \in L^{1-D_i}_b(X_i)\Big\} \Big\rvert},
$$
we recover the forest-based estimator $\hat\tau_{\rm RF}(x)$.
We now provide some general regulations on the smoothing parameters $w_{i\leftarrow j}$'s and the regression adjustment terms $\hat\mu_0,\hat\mu_1$. They are in parallel to the corresponding assumptions in lin2022regression (suggested in the parenthesis of each assumption) and are intent to be general.
assumption[Assumption 3.2 in lin2022regression] \phantomsection
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• Consider any permutation $\pi: \llbracket n\rrbracket \to \llbracket n\rrbracket$. Let $[w_{i\leftarrow j}]_{D_i+D_j=1}$ and $[w^\pi_{i\leftarrow j}]_{D_i+D_j=1}$ be the weights constructed by $[(X_i,D_i,Y_i)]_{i=1}^n$ and $[(X_{\pi(i)},D_{\pi(i)},Y_{\pi(i)})]_{i=1}^n$, respectively. It is assumed that for any $i,j \in \llbracket n\rrbracket$ such that $D_i+D_j=1$ and any permutation $\pi$, $w_{i\leftarrow j} = w^\pi_{\pi(i)\leftarrow \pi(j)}$ holds true.
• The weights satisfy
\begin{align*}
\sum_{j:D_j=1-D_1} w_{1\leftarrow j}=1.
\end{align*}
\end{enumerate}
assumption\phantomsection
Assume that, $\sum_{j:D_j=1-D_i} \lvert w_{i\leftarrow j}\rvert$ is bounded.
assumption[Assumption 3.3 in lin2022regression] \phantomsection
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• We assume that for $\omega \in \{0,1\}$ a deterministic function $\bar{\mu}_\omega(\cdot):\bR^d \to \bR$ exists such that (a) $\bar{\mu}_\omega(x)$ is uniformly bounded for almost all $x \in \operatorname{supp}(\zeta)$, and (b) the regression adjustment term $\hat{\mu}_\omega(x)$ satisfies
$$
\lVert \hat{\mu}_\omega - \bar{\mu}_\omega \rVert_\infty = o_{\mathrm P}(1),
$$
where $\|\cdot\|_{\infty}$ represents the infinity norm (conditional on the data).
• For all $x \in \operatorname{supp}(\zeta)$, the weights satisfy that there exists $\epsilon>0$, such that
\begin{equation}
h^{d-m}{\mathrm E} \Big[\Big( \frac{1}{n} \sum_{i=1}^n \Big(\sum_{j:D_j=1-D_i} K_{h,j}w_{j\leftarrow i} - K_{h,i}\Big(D_i \frac{1-e(X_i)}{e(X_i)} + (1-D_i) \frac{e(X_i)}{1-e(X_i)} \Big)\Big) \Big)^2\Big]
= o(1),
\end{equation}
and
\begin{equation}
h^{d-m}{\mathrm E} \Big[ \Big(\sum_{j:D_j=1-D_1} K_{h,j}w_{j\leftarrow 1} - K_{h,1}\Big(D_1 \frac{1-e(X_1)}{e(X_1)} + (1-D_1) \frac{e(X_1)}{1-e(X_1)} \Big) \Big)^2\Big] \lesssim h^{-m/2} \Big(h^{\epsilon}+\frac{1}{n}\Big),
\end{equation}
where $K_{h,i}:= h^{-d/2}K(\|h^{-1/2}(X_i-x)\|)$.
\end{enumerate}
assumption[Assumption 3.4 in lin2022regression] \phantomsection
\begin{enumerate}[itemsep=-.5ex,label=(\roman*)]
• It is assumed that, for $\omega \in \{0,1\}$, the regression adjustment term $\hat{\mu}_\omega(x)$ satisfies
$$
\lVert \hat{\mu}_\omega - \mu_\omega \rVert_\infty = o_{\mathrm P}(1).
$$
• The weights $[w_{1\leftarrow j}]_{D_j=1-D_1}$ are constructed by $[(X_i,D_i)]_{i=1}^n$ only and not using any outcome information in $[Y_i]_{i=1}^n$.
• For all $x \in \operatorname{supp}(\zeta)$, the weights satisfy
\begin{align*}
h^{-m/2} {\mathrm E}\Big[\Big\lvert\sum_{D_j=1-D_1} K\left(\lVert h^{-1/2}\left(X_j-x\right)\rVert \right) w_{j \leftarrow 1}\Big\rvert\Big] \lesssim 1,
\end{align*}
and
\begin{align*}
{\mathrm E}\Big[\Big( h^{-m/2} \sum_{D_j=1-D_1} K\left(\lVert h^{-1/2}\left(X_j-x\right)\rVert \right) w_{j \leftarrow 1}\Big)^2\Big] \lesssim h^{-m/2}.
\end{align*}
\end{enumerate}
remarkAssumption (ref) corresponds to the case where the mean function models are misspecified but the propensity score model is correctly specified, and the parallel Assumption (ref) corresponds to the scenario where the mean function models are correctly specified.
With the above assumptions held, the first theorem of this section gives a {\it double robustness} result on $\hat\tau_w(\cdot)$. See, e.g., kennedy2020towards and diaz2018targeted about similar observations for other CATE estimators.
theorem[Double robustness of $\hat\tau_w$]
Suppose Assumptions (ref)-(ref) and (ref)-(ref) hold, and either Assumptions (ref) or Assumptions (ref) is true. We then have for any interior point $x \in \operatorname{supp}(\zeta)$,
\begin{align*}
\hat\tau_w(x) - \tau(x) \stackrel{\sf p}{\longrightarrow} 0.
\end{align*}
We then move on to study the rates of convergence.
assumption\phantomsection Assume that, for $\omega \in \{0,1\}$, $\hat{\mu}_\omega$ is Lipchitz over $\operatorname{supp}\left(\zeta\right)$.
assumption\phantomsection There exists some constant $\gamma>1$ such that the weights satisfy
\begin{equation*}
{\mathrm E} \Big[\Big(\frac{1}{n} \sum_{i=1}^n h^{-m/2} K\left(\lVert h^{-1/2}\left(X_i-x\right)\rVert\right) \sum_{j:D_j=1-D_i} \lvert w_{i\leftarrow j} \rvert \lVert X_i-X_j \rVert\Big)^{2\gamma}\Big] \lesssim h^{\gamma}.
\end{equation*}
remarkIn the context of the central limit theorem, it is important to note that we can ease the Assumption (ref) to $L^1$ convergence.
assumption\phantomsection
For $\omega \in \{0,1\}$, the estimator $\hat{\mu}_\omega(x)$ satisfies
$$
{\mathrm E}\Big[\lVert \mu_\omega - \hat{\mu}_\omega \rVert_\infty^2\Big] \lesssim h.
$$
theorem[MSE of $\hat\tau_w(x)$]
Suppose Assumptions (ref)-(ref), (ref)-(ref), (ref)-(ref), (ref)(ref) hold, with either Assumptions (ref)(ref), (ref)-(ref) or Assumptions (ref)(ref), (ref) held. We then have, for any interior point $x \in \operatorname{supp}(\zeta)$,
\begin{align*}
MSE\left(\hat\tau_w \left(x\right)\right)={\mathrm E} \Big[\Big(\hat\tau_w(x) - \tau(x)\Big)^2\Big] \lesssim h + \frac{1}{n h^{m/2}}.
\end{align*}
In particular, when $h \asymp n^{-{\frac{2}{m+2}}}$,
\begin{equation}
MSE\left(\hat\tau_w \left(x\right)\right) \lesssim n^{-{\frac{2}{m+2}}}.
\end{equation}
Lastly, we consider the central limit theorem.
assumption\phantomsection There exist some functions, $\phi_0(x)$ and $\phi_1(x)$, as well as $\epsilon>0$, such that for all points $x \in \operatorname{supp}(\zeta)$, the weights satisfy
\begin{equation*}
h^{d-m}{\mathrm E} \Big[\Big(\sum_{j:D_j=1-D_1} K_{h,j}w_{j\leftarrow 1} - K_{h,1}\Big(D_1 \phi_1(X_1) + (1-D_1) \phi_0(X_1)\Big) \Big)^2\Big] \lesssim h^{-m/2}\Big(h^{\epsilon} + \frac{1}{n}\Big).
\end{equation*}
Furthermore, assume that the following limit exists:
\begin{equation*}
\lim_{n \to \infty } h^{-m/2} {\mathrm E} \Big[K^2\left(\lVert h^{-1/2}\left(X_1-x\right)\rVert\right)\Big[D_1\left(1+\phi_1\left(X_1\right)\right)+\left(1-D_1\right)\left(1+\phi_0\left(X_1\right)\right)\Big]^2 \sigma_{D_1}^2(X_1)\Big]=: \tilde{\Sigma}(x).
\end{equation*}
remarkWe note that under Assumption (ref)(ref), $\phi_0(x)$ and $\phi_1(x)$ in Assumption (ref) take the following specific forms
\begin{equation*}
\quad \phi_0(x) = \frac{e(x)}{1-e(x)} \quad and\quad \phi_1(x) = \frac{1-e(x)}{e(x)}.
\end{equation*}
For the doubly robust estimator, the propensity score function can be misspecified, in which case $\phi_0(x)$ and $\phi_1(x)$ can be general functions without the above specific forms. From a technical standpoint, a central limit theorem without closed forms for $ \phi_0(x)$ and $ \phi_1(x)$ allows for more flexible assumption on the smoothing weight.
theorem[Central limit theorem of $\hat\tau_w(x)$]
Suppose Assumptions (ref)-(ref), (ref)-(ref), (ref)-(ref), (ref)(ref) hold with either Assumptions (ref)(ref), (ref)-(ref), or Assumptions (ref)-(ref) held. Furthermore, assume $nh^{m/2+1}\to 0$. We then have, for any interior point $x \in \operatorname{supp}(\zeta)$,
\begin{align*}
\sqrt{nh^{m/2}} \left(\hat\tau_w\left(x\right) - \tau\left(x\right) \right) \stackrel{\sf d}{\longrightarrow} N(0,\Sigma(x)).
\end{align*}
Here under Assumption (ref)(ref),
$$
\Sigma(x)=\frac{1}{c_K^2 f(x)}\left(\frac{\sigma^2_1(x)}{e(x)} + \frac{\sigma^2_0(x)}{1-e(x)}\right)\int_{\mathbb{R}^m} K^2(\lVert t\rVert) {\mathrm d} t;
$$
and under Assumption (ref),
$$
\Sigma(x)=\frac{\tilde{\Sigma}(x)}{c_K^2 f^2(x)},
$$
which, if $\phi_0(x)$ and $\phi_1(x)$ are Lipchitz and bounded and bounded away from zero, is
$$
\Sigma(x)=\frac{1}{c_K^2 f(x)}\Big( e(x)\left(1+\phi_1(x)\right)^2\sigma^2_1(x) + \left(1-e(x)\right)\left(1+\phi_0(x)\right)^2\sigma^2_0(x)\Big)\int_{\mathbb{R}^m} K^2(\lVert t\rVert) {\mathrm d} t.
$$
Discussions
In this study, we have demonstrated that causal forests can effectively adapt to unknown covariate manifold structures, under the (extremely) honesty assumption of trees in the forests. An interesting future direction is relaxing the honesty assumption and exploring data-dependent splits. We have considered fixed dimensions for the ambient space and manifold. Recent research scornet2015consistency,chi2022asymptotic has shown that random forests breiman2001random can adapt to model sparsity and high dimensionality, sparking interest in investigating a similar adaptability aspect in CATE estimation using random forests. This avenue of exploration holds potential for enhancing CATE estimation methodologies.
Ackowledgement
The authors thank Vasilis Syrgkanis for pointing out related research, and Jason M. Klusowski for helpful comments on the theory of random forests.