\section{Introduction} \vspace{-16pt}
	The understanding of treatment effects plays  an important role -- especially in shaping interventions and treatments -- in areas from  clinical trials \cite{rothwell2005subgroup,zhou2017residual} to recommender systems \cite{lada2019observational} to  public policy \cite{grimmer2017estimating}.  In many settings, the relevant population is diverse, and different parts of the population display different reactions to treatment.  
	In such settings, {\em heterogeneous treatment effect (HTE) analysis}   -- also called  {\em subgroup analysis}--  is used to find subgroups consisting of subjects who have similar covariates and display similar treatment responses \cite{imai2013estimating,foster2011subgroup}.  The identification of subgroups is informative of itself; it also  improves the interpretation of treatment effects across the entire population and makes it possible to develop more effective interventions and treatments and to improve the design of further experiments.  In a clinical trial, for example, HTE analysis can identify subgroups of the population for which the studied treatment is effective, even when it is found to be ineffective for the population in general  \cite{hobbs2011hierarchical}.
	
	To identify subjects who have similar covariates and display similar treatment responses, it is necessary to create reliable estimates of the treatment responses of individual subjects; i.e. of {\em individualized treatment effects (ITE)}.  The state-of-the-art work on HTE proceeds by simultaneously estimating ITE and recursively partitioning the subject population    \cite{athey2016recursive,su2009subgroup,tran2019learning,johansson2018interpretable}.
	In these HTE methods, the criterion for partitioning is maximizing  the heterogeneity of  treatment effects {\em across} subgroups, using a sample mean estimator, under the assumption that treatment effects are homogeneous {\em within}  subgroups.  In particular, the population (or any previously identified subgroup) would be  partitioned into two subgroups provided that the sample means of these subgroups are sufficiently different, ignoring the possibility that  treatment effects  might be very heterogeneous within the groups identified.  Put differently, these methods focus on inter-group heterogeneity but ignore intra-group heterogeneity.\looseness=-1
	
	
	\setlength{\intextsep}{7pt}%
	\begin{wrapfigure}{r}{0.4\textwidth}
		\centering
		\includegraphics[width=1.\linewidth]{figs/sim/toy_example}
		\caption{A toy example with two subgroups identified by HTE method in \cite{tran2019learning}. The solid red line shows the ITE estimation and their 95\% confidence interval is filled in red.}
		\label{fig:toy_example}
	\end{wrapfigure}
	
	An important problem with this approach is that, because it relies solely on inter-group heterogeneity based on sample means, it may lead to \textit{false discovery}.  To illustrate, consider the toy example depicted in Fig. \ref{fig:toy_example}.  
	In this example, the true treatment effect (shown on the vertical axis) was generated by iid random draws from the normal distribution having mean $0.0$ and standard deviation $0.1$.  In truth, the treatment under consideration is in fact {\em totally ineffective and innocuous};  on average, it has  {\em no effect at all} and the treatment effects are entirely uncorrelated with the single covariate (shown on the horizontal axis).  
	However if the observed data -- the realization of the random draws -- happens  to be the one shown in Fig. \ref{fig:toy_example}, standard methods will typically partition the population as shown in the figure, thereby ``discovering'' a segment of the population for whom the treatment is effective and a complementary segment where the treatment is dangerous.  Obviously, decisions based on such false discovery are useless -- or worse.
	Note that this false discovery occurs because, although the outcome variations between the two groups are indeed substantially different, the outcome variations within each group are just as different -- but the latter variation is entirely ignored in the creation of subgroups.\looseness=-1
	
	This paper proposes a robust recursive partitioning (R2P)\footnote{The code of R2P is available at: \url{https://bitbucket.org/mvdschaar/mlforhealthlabpub}.} method that avoids such false discovery.  R2P has several  distinctive characteristics. 
	\vspace{-\topsep}
	\begin{itemize}[leftmargin=*]\itemsep=0pt
		\item {\em R2P discovers interpretable subgroups in a way that is not tied to any {\em particular} ITE estimator.}  This is in sharp contrast with previous methods \cite{athey2016recursive,su2009subgroup,tran2019learning,johansson2018interpretable,seibold2016model}, each of which relies on a {\em specific} ITE estimator.  R2P can leverage {\em any} ITE estimator for subgroup analysis, e.g. an ITE estimator  based on Random Forest \cite{wager2018estimation}, or on multi-task Gaussian processes \cite{alaa2017bayesian} or on deep neural networks \cite{yoon2018ganite,shalit2017estimating,zhang2020learning}. This flexibility is important because no one ITE estimator is consistently the best in all different settings \cite{pmlr-v97-alaa19a}. 
			Furthermore, these ITE estimators are non-interpretable black-box models. R2P divides units into subgroups with respect to an interpretable tree-structure, and provides subgroup coverage guarantees for the ITE estimates in each subgroup. R2P enables these ITE estimators to produce trustworthy and interpretable ITE estimates in practice.
		
	
	
		
		\item  {\em R2P makes a conscious effort to guarantee homogeneity of treatment effects within each of the subgroups while maximizing heterogeneity across subgroups.} This is also different from previous methods, e.g., \cite{athey2016recursive,tran2019learning} where variation within the subgroup is largely ignored.
		
		\item {\em R2P produces  {\em confidence guarantees} with  narrower confidence intervals than previous methods.}  It accomplishes this by using methods of conformal prediction \cite{lei2018distribution} that produce valid confidence intervals.  Quantifying the uncertainty allows R2P to employ a novel  criterion we call {\em confident homogeneity} in order to create partitions that take account of both  heterogeneity across groups and homogeneity within groups. 
	\end{itemize}
	
		These characteristics make R2P both more reliable and more informative than the existing methods for subgroup analysis.
		Extensive experiments using synthetic and semi-synthetic datasets (based on real-world data) demonstrate that  R2P outperforms state-of-the-art methods by more robustly identifying subgroups while providing much narrower confidence intervals.

	

	
	\section{Robust Recursive Partitioning with Uncertainty Quantification}
	To highlight  the core design principles, we begin by introducing  robust recursive partitioning (R2P)  in the  regression setting;  we extend to the more complicated HTE setting in the next section.
	
	\subsection{Preliminaries}
	\vspace{-6pt}
	We consider a standard regression problem with a $d$-dimensional  covariate (input) space $\mathcal{X}\subseteq\mathbb{R}^d$ and a outcome space $\mathcal{Y}\subseteq\mathbb{R}$. We are given a dataset $\ensuremath{{\mathcal D}}=\{(\ensuremath{{\mathbf x}}_i,y_i)\}_{i=1}^n$, where, for the $i$-th sample, $\ensuremath{{\mathbf x}}_i\in\ensuremath{{\mathcal X}}$ is the vector of input covariates  and $y_i\in\ensuremath{{\mathcal Y}}$ is the outcome. We assume that  samples are independently drawn from an unknown distribution $\ensuremath{{\mathcal P}}_{X,Y}$ on 
	$\ensuremath{{\mathcal X}}\times\ensuremath{{\mathcal Y}}$. We are interested in estimating $\mu(\ensuremath{{\mathbf x}})=\mathbb{E}[Y|X=\ensuremath{{\mathbf x}}]$, which is the mean outcome conditional on $\ensuremath{{\mathbf x}}$. We denote the estimator by $\hat{\mu}:\ensuremath{{\mathcal X}}\rightarrow\ensuremath{{\mathcal Y}}$;  $\hat{\mu}$ predicts an outcome $\hat{y}= \hat{\mu}(\ensuremath{{\mathbf x}})$ on the basis of the covariate information $\ensuremath{{\mathbf x}}$. To quantify the uncertainty in the prediction, we apply the method of split conformal regression (SCR) \cite{lei2018distribution} to construct a confidence interval $\hat{C}$ that satisfies the rigorous frequentist guarantee in the finite-sample regime. (To the best of our knowledge, SCR is the simplest method that achieves this guarantee.)
	
	In SCR, we take as given a {\em miscoverage rate} $\alpha \in (0,1)$.  We split the samples in $\ensuremath{{\mathcal D}}$ into a training set $\ensuremath{{\mathcal I}}_1$ and a validation set $\ensuremath{{\mathcal I}}_2$ that are disjoint and have the same size. We train the estimator 
	$\hat{\mu}^{\ensuremath{{\mathcal I}}_1}$ on $\ensuremath{{\mathcal I}}_1$ and  compute the residual of $\hat{\mu}^{\ensuremath{{\mathcal I}}_1}$ on each sample in $\ensuremath{{\mathcal I}}_2$. For a testing sample $\ensuremath{{\mathbf x}}$ the confidence interval is given by
	\vspace{-6pt}
	\begin{equation}
		\label{eqn:conf_entire}
		\hat{C}^{\ensuremath{{\mathcal I}}_1,\ensuremath{{\mathcal I}}_2} (\ensuremath{{\mathbf x}}) = \left[\hat{\mu}^{\textrm{lo}}(\ensuremath{{\mathbf x}}),\hat{\mu}^{\textrm{up}}(\ensuremath{{\mathbf x}})\right]=\left[\hat{\mu}^{\ensuremath{{\mathcal I}}_1}(\ensuremath{{\mathbf x}})-\hat{Q}_{1-\alpha}^{\ensuremath{{\mathcal I}}_2},\hat{\mu}^{\ensuremath{{\mathcal I}}_1}(\ensuremath{{\mathbf x}})+\hat{Q}_{1-\alpha}^{\ensuremath{{\mathcal I}}_2}\right],
		\vspace{-6pt}
	\end{equation}
	where $\hat{Q}_{1-\alpha}^{\ensuremath{{\mathcal I}}_2}$ is defined to be the $(1-\alpha)(1+1/|\ensuremath{{\mathcal I}}_2|)$-th quantile of the set of residuals  $\{|y_i-\hat{\mu}^{\ensuremath{{\mathcal I}}_1}(\ensuremath{{\mathbf x}}_i)|\}_{i\in\ensuremath{{\mathcal I}}_2}$. Assuming that the  training and testing samples are drawn exchangeably from $\ensuremath{{\mathcal P}}_{X,Y}$,  the confidence intervals defined in (\ref{eqn:conf_entire})  satisfy the coverage guarantee $\mathbb{P}[y\in\hat{C}^{\ensuremath{{\mathcal I}}_1,\ensuremath{{\mathcal I}}_2}]\geq 1-\alpha$  \cite{lei2018distribution}.\footnote{Recall that assuming exchangeability is weaker than assuming iid.}  
	
	To illustrate, assume the miscoverage rate $\alpha$ is $0.05$ and we are given 1000 testing samples.  SCR prescribes a confidence interval for each sample in such a way that for at least 950 samples the prediction is within the prescribed confidence interval.  (We often say the sample is {\em covered}.)  However this coverage guarantee is marginal over the entire covariate space.  If we perform a  subgroup analysis that partitions the covariate space $\mathcal{X}$ into subgroups $\mathcal{X}_1, \mathcal{X}_2$ in such a way that $\mathcal{X}_1$ has 800  samples and $\mathcal{X}_2$ has 200 samples, it might be the case that 790 samples in $\mathcal{X}_1$ are covered but only $160$ samples in $\mathcal{X}_2$ are covered.  In this case,  80\% of the samples in $\mathcal{X}_2$ would be covered.  It seems obvious that such a situation is undesirable for subgroup analysis; we want to achieve the prescribed coverage rate for {\em each} subgroup, not just for the population as a whole.  R2P  overcomes this problem.   
	
	We begin by discussing how we use confidence intervals to quantify  outcome homogeneity within a subgroup. We then introduce our space partitioning algorithm and provide the required theoretical guarantee of subgroup coverage.
	
	\subsection{Partitioning for Robust Heterogeneity Analysis}
	\label{sec:criterion}
	Let $\Pi$ be a partition of the covariate space $\mathcal{X}$ into (disjoint) subgroups.  Write  $|\Pi|$ for the number of subgroups in the partition, $l_j$ for an element of $\Pi$ and $l(\ensuremath{{\mathbf x}};\Pi)$ for the subgroup that contains the sample $\ensuremath{{\mathbf x}}$. Write $\ensuremath{{\mathcal D}}_l=\{(\ensuremath{{\mathbf x}}_i,y_i)\in\ensuremath{{\mathcal D}}|\ensuremath{{\mathbf x}}_i\in l\}$ for the samples whose covariates belong to the subgroup $l$. Note that when we restrict to covariates in the subgroup $l$, the samples are drawn from the truncated distribution $\ensuremath{{\mathcal P}}_{X,Y}^l$ which is the distribution conditional on the requirement that the vector of covariates of samples lie in the subgroup $l$. 
	
	We evaluate  homogeneity within the subgroup $l$ by the concentration of outcome values for covariate vectors in the subgroup $l$. To do this, we apply SCR to the samples in  $\ensuremath{{\mathcal D}}_l$ by splitting it into two sets, $\ensuremath{{\mathcal I}}_1^l$ and $\ensuremath{{\mathcal I}}_2^l$. Write $\hat{\mu}_l(\ensuremath{{\mathbf x}})$ denote the mean outcome model trained on  $\ensuremath{{\mathcal I}}_1^l$. As in (\ref{eqn:conf_entire}), we  obtain the confidence interval $\hat{C}_l(\ensuremath{{\mathbf x}})$ for subgroup $l$ by setting the upper and lower endpoints to be
	$\hat{\mu}_l^{\textrm{up}}(\ensuremath{{\mathbf x}})=\hat{\mu}_l(\ensuremath{{\mathbf x}})+\hat{Q}_{1-\alpha}^{\ensuremath{{\mathcal I}}_2^l}$ and $\hat{\mu}_l^{\textrm{lo}}(\ensuremath{{\mathbf x}})=\hat{\mu}_l(\ensuremath{{\mathbf x}})-\hat{Q}_{1-\alpha}^{\ensuremath{{\mathcal I}}_2^l}$, respectively. (To avoid notational complications, omit reference to the subsets $\ensuremath{{\mathcal I}}_1^l$ and $\ensuremath{{\mathcal I}}_2^l$ hereafter; this should not cause confusion.  Throughout, we follow the convention that the confidence bound have been computed on the basis of the split.) To estimate the center of  the subgroup $l$, we use the average outcome $\hat{\mu}_{l,\textrm{mean}}=\mathbb{E}[\hat{\mu}_l(\ensuremath{{\mathbf x}})]$.  We define the {\em expected absolute deviation} in group $l$ to be
	$
	S_l=\mathbb{E}[v_l(\ensuremath{{\mathbf x}})],
	$
	where 
	\begin{equation}\label{equ:v_l}
		v_l(\ensuremath{{\mathbf x}})=\left(\hat{\mu}_{l,\textrm{mean}}-\hat{\mu}_l^{\textrm{up}}(\ensuremath{{\mathbf x}})\right)\mathbb{I}\left[\hat{\mu}_{l,\textrm{mean}}>\hat{\mu}_l^{\textrm{up}}(\ensuremath{{\mathbf x}})\right]
		+\left (\hat{\mu}_l^{\textrm{lo}}(\ensuremath{{\mathbf x}})-\hat{\mu}_{l,\textrm{mean}} \right)\mathbb{I}\left[\hat{\mu}_{l,\textrm{mean}}<\hat{\mu}_l^{\textrm{lo}}(\ensuremath{{\mathbf x}})\right].
	\end{equation}
	By definition,  $\hat{\mu}_l^{\textrm{up}}(\ensuremath{{\mathbf x}})$ is larger than 
	$\hat{\mu}_l^{\textrm{lo}}(\ensuremath{{\mathbf x}})$ provided the residual quantile $\hat{Q}_{1-\alpha}>0$. When the first indicator function is one, i.e. the average outcome (the group center) $\hat{\mu}_{l,\textrm{mean}}$ is larger than the upper bound $\hat{\mu}_l^{\textrm{up}}(\ensuremath{{\mathbf x}})$ at $\ensuremath{{\mathbf x}}$, we are confident that the outcome value at $\ensuremath{{\mathbf x}}$ is smaller than the group center (and perhaps smaller than the outcome value for many covariate vectors in $l$). Similarly, when the second indicator function is one, we are certain that the outcome value at $\ensuremath{{\mathbf x}}$ is larger than the group center (and perhaps larger than the outcome value for many covariate vectors in $l$).  It is worth noting that when  $\hat{C}_l(\ensuremath{{\mathbf x}})=\big[\hat{\mu}_l^{\textrm{lo}}(\ensuremath{{\mathbf x}}),\hat{\mu}_l^{\textrm{up}}(\ensuremath{{\mathbf x}})\big]$ contains the center $\hat{\mu}_{l,\textrm{mean}}$, both indicator functions are zero and $v_l(\ensuremath{{\mathbf x}})=0$. The quantity $v_l(\ensuremath{{\mathbf x}})$ evaluates the homogeneity in subgroup $l$ on the basis of the confidence interval for each 
	$\ensuremath{{\mathbf x}}$ in $l$. (It is more conservative than the mean discrepancy $|\hat{\mu}_l(\ensuremath{{\mathbf x}})-\hat{\mu}_{l,\textrm{mean}}|$ for partitioning, and hence provides greater protection against false discovery because of  uncertainty.)
	
	\begin{figure}
		\centering
		\includegraphics[width=1.\textwidth]{figs/conf_crit.eps}
		\caption{\textbf{Illustration of partitioning and impurity of the confident homogeneity.}  
			The regions shaded in red and gray (roughly) represent $W_l$ and $S_l$, respectively.
			Partitioning the heterogeneous covariate space (left panel) reduces its impurity of the confident homogeneity. The partition with the smaller impurity (middle panel) makes the heterogeneity across subgroups and the homogeneity within subgroups both stronger compared to others with the larger impurity (e.g., right panel). \label{fig:concept}}
		\vspace{-0.1in}
	\end{figure}
	
	However,  minimizing $S_l$ is not enough to maximize  subgroup homogeneity. If the intervals $\hat{C}_l(\ensuremath{{\mathbf x}})$ for all $\ensuremath{{\mathbf x}} \in l$ are very wide and contain the average outcome $\hat{\mu}_{l,\textrm{mean}}$,  homogeneity can be very low even though $S_l=\mathbb{E}[v_l(\ensuremath{{\mathbf x}})]$ is zero. To resolve this issue, when partitioning the covariate space we jointly minimize $S_l$ and the expected confidence interval width $W_l=\mathbb{E}\big[|\hat{C}_l(\ensuremath{{\mathbf x}})|\big]$. We formalize the robust partitioning problem as
	\begin{equation}
		\label{eqn:objective}
		\mathop{\rm minimize}_{\Pi} \sum_{l\in\Pi} \lambda W_l + (1-\lambda) S_l,
	\end{equation}
	where $\lambda\in[0,1]$ is a hyperparameter that balances the impact of $W_l$ and $S_l$.  We call the weighted sum, $\lambda W_l + (1-\lambda) S_l$, the impurity of the \textit{confident homogeneity} for subgroup $l$. Fig. \ref{fig:concept} illustrates how minimizing the impurity of the confident homogeneity improves both homogeneity within subgroups and heterogeneity across subgroups. There may be more than one partition that achieves the minimum; because a larger number of subgroups is harder to interpret, we will choose a minimizer with the smallest number of subgroups.
	
	
	


	
	\subsection{Robust Recursive Partitioning Method with Confident Homogeneity}
	We can now describe our robust recursive partitioning (R2P) method for solving the optimization problem  \eqref{eqn:objective}.
	We begin with the trivial partition  $\Pi=\{\ensuremath{{\mathcal X}}\}$.  
	We denote the set of subgroups whose objectives in \eqref{eqn:objective} can be potentially improved by $\Pi_c$.
	In the initialization step, we set $\Pi_c=\Pi$ and apply the SCR on $\ensuremath{{\mathcal D}}$ to obtain the confidence intervals in (\ref{eqn:conf_entire}). Based on these intervals, we compute $\hat{W}_\ensuremath{{\mathcal X}}$ and $\hat{S}_\ensuremath{{\mathcal X}}$.
More generally, using the intervals for subgroup $l$, we can estimate $W_l$ and $S_l$ by
		\[
		\hat{W}_l=\frac{1}{N_2^l}\sum_{i\in\ensuremath{{\mathcal I}}_2^l}|\hat{C}_l(\ensuremath{{\mathbf x}}_i)| \ \  \text{and} \ \  \hat{S}_l=\frac{1}{N_2^l}\sum_{i\in\ensuremath{{\mathcal I}}_2^l}v_l(\ensuremath{{\mathbf x}}_i),
		\]
		respectively, where $N_2^l$ is the number of samples in the validation set $\ensuremath{{\mathcal I}}_2^l$. 
	\begin{algorithm}[b]
		\caption{Robust Recursive Partitioning}\label{alg:algorithm}
		\begin{algorithmic}[1]
			\State \textbf{Input}: Samples $\ensuremath{{\mathcal D}}=\{(\ensuremath{{\mathbf x}}_i,y_i)\}_{i=1}^n$, miscoverage rate $\alpha\in(0,1)$, $\Pi=\{\ensuremath{{\mathcal X}}\}$
			\State \textbf{Initialize}: $\Pi_{c}=\Pi$, split $\ensuremath{{\mathcal D}}$ into $\ensuremath{{\mathcal I}}_1^\ensuremath{{\mathcal X}}$ and $\ensuremath{{\mathcal I}}_2^\ensuremath{{\mathcal X}}$, train $\hat{\mu}_\ensuremath{{\mathcal X}}$, compute its confidence interval $\hat{C}_\ensuremath{{\mathcal X}}$ using the split subsets, and obtain $\hat{W}_\ensuremath{{\mathcal X}}$ and $\hat{S}_\ensuremath{{\mathcal X}}$ using $\ensuremath{{\mathcal I}}^\ensuremath{{\mathcal X}}_2$
			\For{$l\in\Pi_c$}{}
			\State Obtain $\hat{W}^*_{l^\pm}$, $\hat{S}^*_{l^\pm}$, $i^*$, and $\phi^*$
			\If{$(1-\gamma)\left[\lambda\hat{W}_{l}+(1-\lambda)\hat{S}_l\right] \geq   \lambda\hat{W}^*_{l^\pm} + (1-\lambda)\hat{S}^*_{l^\pm}$}{}\Comment{\textit{Confident} criterion}
			\State Partition $l$ into $l^+(i^*,\phi^*)$ and $l^-(i^*,\phi^*)$
			\State $\Pi\leftarrow\Pi\setminus\{l\}$
			\State $\Pi\leftarrow\Pi\cup\{l^+(i^*,\phi^*),l^-(i^*,\phi^*)\}$ and $\Pi_c\leftarrow\Pi_c\cup\{l^+(i^*,\phi^*),l^-(i^*,\phi^*)\}$
			\EndIf
			\State $\Pi_c\leftarrow\Pi_c\setminus\{l\}$
			\EndFor
			\State \textbf{Output}: $\Pi$, $\hat{\mu}_\ensuremath{{\mathcal X}}$, and $\hat{C}_l$ for all $l\in\Pi$
		\end{algorithmic}
	\end{algorithm}
	
	After initialization, we recursively partition the covariate space by splitting the subgroups in $\Pi_c$ with respect to the criterion in \eqref{eqn:objective}. To split each subgroup $l\in\Pi_c$, we first consider the two disjoint subsets from subgroup $l$ given by
	$l_k^+(\phi)=\{\ensuremath{{\mathbf x}}\in l|x_k\geq\phi\} \textrm{ and } l_k^-(\phi)=\{\ensuremath{{\mathbf x}}\in l|x_k<\phi\}$,
	where $\phi\in(x_k^{l,\textrm{min}},x_k^{l,\textrm{max}})$ is the threshold for splitting, $x_k$ is the $k$-th covariate element, and $x_k^{l,\textrm{min}}$ and $x_k^{l,\textrm{max}}$ are the minimum and maximum values of the $k$-th covariate within the subgroup $l$, respectively. We then apply SCR to each of these subsets. Specifically, we split the samples corresponding to $l_k^+(\phi)$ and $l_k^-(\phi)$ into  training and validation sets:
	$
	\ensuremath{{\mathcal D}}_{l_k^+(\phi)} =\ensuremath{{\mathcal I}}_1^{l_k^+(\phi)} \cup \ensuremath{{\mathcal I}}_2^{l_k^+(\phi)} \textrm{ and }
	\ensuremath{{\mathcal D}}_{l_k^-(\phi)}=\ensuremath{{\mathcal I}}_1^{l_k^-(\phi)}\cup\ensuremath{{\mathcal I}}_2^{l_k^-(\phi)},
	$
	where the split subsets are  $\ensuremath{{\mathcal I}}_1^{l_k^+(\phi)}=\{(\ensuremath{{\mathbf x}}_i,y_i)\in\ensuremath{{\mathcal I}}_1^l:\ensuremath{{\mathbf x}}_i\in {l_k^+(\phi)}\}$ and $\ensuremath{{\mathcal I}}_2^{l_k^+(\phi)}=\{(\ensuremath{{\mathbf x}}_i,y_i)\in\ensuremath{{\mathcal I}}_2^l:\ensuremath{{\mathbf x}}_i\in {l_k^+(\phi)}\}$ (same for ${l_k^-(\phi)}$).  To compute residuals, we do not train new estimators for $l_k^+(\phi)$ and $l_k^-(\phi)$; instead we use the previously trained estimator $\hat{\mu}_\ensuremath{{\mathcal X}}$; this provides consistency of estimators across groups and within groups.  (It also avoids the enormous computational burden of  training new estimators for all the possible splits.) Using the residuals, we can construct the confidence intervals $\hat{C}_{l_k^+(\phi)}(\ensuremath{{\mathbf x}})$ and $\hat{C}_{l_k^-(\phi)}(\ensuremath{{\mathbf x}})$ and the associated quantities in the objective function, $\hat{W}_{l_k^+(\phi)},\hat{W}_{l_k^-(\phi)}, \hat{S}_{l_k^+(\phi)}$ and $\hat{S}_{l_k^-(\phi)}$. Then we find the optimal covariate $k^*_l$ and threshold $\phi^*_l$ for splitting subgroup $l$ as
	\[
	(k^*_l,\phi^*_l) = \mathop{\rm argmin}_{(k,\phi)} \lambda \left(\hat{W}_{l_k^+(\phi)}+\hat{W}_{l_k^-(\phi)} \right)+(1-\lambda)\left(\hat{S}_{l_k^+(\phi)}+\hat{S}_{l_k^-(\phi)} \right).
	\]
	For $(k^*_l,\phi^*_l)$, we compute $\hat{W}^*_{l^\pm}=\hat{W}_{l_{k^*}^+(\phi^*)}+\hat{W}_{l_{k^*}^-(\phi^*)}$ and $\hat{S}^*_{l^\pm}=\hat{S}_{l_{k^*}^+(\phi^*)}+\hat{S}_{l_{k^*}^-(\phi^*)}$. To improve the objective in \eqref{eqn:objective}, we split the subgroup $l$ into $l_{k^*}^+(\phi^*)$ and $l_{k^*}^-(\phi^*)$ only if the reduction in the impurity of the confident homogeneity is sufficiently large:
	\begin{equation}
		\label{eqn:confident_criterion}
		(1-\gamma)\left[\lambda\hat{W}_{l}+(1-\lambda)\hat{S}_l\right] \geq   \lambda\hat{W}^*_{l^\pm} + (1-\lambda)\hat{S}^*_{l^\pm}
	\end{equation}
	Here, $\gamma\in[0,1)$ is a hyperparameter for regularization. We refer to  \eqref{eqn:confident_criterion} as the \textit{confident criterion}.  With an appropriate choice of $\gamma$,  this criterion prevents overfitting, prevents the number of subgroups from becoming too large and prevents the size of each subgroup from becoming too small. Confident homogeneity does not necessarily improve as the group size shrinks because smaller groups lead to greater uncertainty.   This alleviates the issue of generalization to unseen data in HTE analysis \cite{athey2016recursive,tran2019learning}.   
	
	After the splitting decision, we remove $l$ from $\Pi_c$; if we have split $l$, we remove $l$ from $\Pi$ and add the two split sets to both $\Pi$ and $\Pi_c$.  We continue recursively until $\Pi_c$ is empty, at which point no further splitting is productive.  When the procedure stops, we will have obtained an estimator 
	$\mu_\ensuremath{{\mathcal X}}$ and a partition $\Pi$ and for each $l \in \Pi$ we will have  corresponding confidence intervals $\hat{C}_l(\ensuremath{{\mathbf x}})=\left[\hat{\mu}_\ensuremath{{\mathcal X}}(\ensuremath{{\mathbf x}})-\hat{Q}_{1-\alpha}^l,\hat{\mu}_\ensuremath{{\mathcal X}}(\ensuremath{{\mathbf x}})+\hat{Q}_{1-\alpha}^l\right]$, where $\hat{Q}_{1-\alpha}^l$ is the $(1-\alpha)(1+1/|\ensuremath{{\mathcal I}}_2^l|)$-th quantile of the set of the residuals on the validation set $\ensuremath{{\mathcal I}}_2^l$ using $\hat{\mu}_\ensuremath{{\mathcal X}}$, $\{|y_i-\hat{\mu}_\ensuremath{{\mathcal X}}(\ensuremath{{\mathbf x}}_i)|\}_{i\in\ensuremath{{\mathcal I}}_2^l}$.  The following theorem guarantees that  the R2P partition Algorithm~\ref{alg:algorithm} provides a valid confidence interval  $\hat{C}_l$ for each subgroup  $l \in \Pi$; this is exactly what a user would want in subgroup analysis.  The proof is  provided  in the supplementary material.
	
	\begin{theorem}\label{thm:thm1}
		Given a prescribed miscoverage rate $\alpha$, the created partition $\Pi$,  estimator $\hat{\mu}_\ensuremath{{\mathcal X}}$, and  confidence interval function $\hat{C}_l(\cdot)$ have the following property: for each $l \in \Pi$ and for new samples $(\ensuremath{{\mathbf x}},y)$ drawn from the truncated distribution $\ensuremath{{\mathcal P}}_{X,Y}^l$, we have
		$\mathbb{P}[y\in\hat{C}_l({\ensuremath{{\mathbf x}}})] \geq 1-\alpha$.
	\end{theorem}
	
	\vspace{-0.1in}
	\section{Robust Recursive Partitioning for Heterogeneous Treatment Effects} \vspace{-6pt}
	In this section, we extend the R2P method to the setting of HTE estimation, resulting in the \textit{R2P-HTE} design as detailed below.
	



	
	\begin{paragraph}{Heterogeneous Treatment Effect Model}
		We consider a setup with $n$ units (samples),
		For the unit $i\in\{1,2,...,n\}$, there exists a pair of potential outcomes, $Y_i(1)$ and $Y_i(0)$ that are independently drawn from an unknown distribution, where 0 and 1 represent whether the unit is treated or not, respectively. We define the treatment indicator as $t_i\in\{0,1\}$, where $t_i=1$ and $0$ indicate the unit $i$ is treated and untreated, respectively.
		The outcome for unit $i$ is realized as the potential outcome corresponding to its treatment indicator $y_i=Y_i(t_i)$. A dataset is given as $\ensuremath{{\mathcal D}}_{\textrm{HTE}}=\{(\ensuremath{{\mathbf x}}_i,t_i,y_i)\}_{i=1}^n$. The ITE for a given $\ensuremath{{\mathbf x}}$ is defined as $\tau(\ensuremath{{\mathbf x}})=\mathbb{E}[Y(1)-Y(0)|X=\ensuremath{{\mathbf x}}]$. Since the ITE is defined as the expected difference between the two potential outcomes $Y(1)$ and $Y(0)$, its estimator $\hat{\tau}(\ensuremath{{\mathbf x}})$ is given as the contrast between two regression models:  
		$\hat{\mu}^0(\ensuremath{{\mathbf x}})$ for the conditional non-treated outcome $\mathbb{E}[Y(0)|X=\ensuremath{{\mathbf x}}]$, and $\hat{\mu}^1(\ensuremath{{\mathbf x}})$ for the conditional treated outcome $\mathbb{E}[Y(1)|X=\ensuremath{{\mathbf x}}]$.
	\end{paragraph}
	
	\begin{paragraph}{R2P-HTE}
		We adapt R2P to HTE estimation by constructing the quantities $W_l$ and $S_l$ in \eqref{eqn:objective} based on the ITE estimator $\hat{\tau}(\ensuremath{{\mathbf x}})$. To construct an ITE estimator, many popular machine learning models have been considered in the literature \cite{athey2016recursive,wager2018estimation,alaa2017bayesian,yoon2018ganite,shalit2017estimating,zhang2020learning}. R2P-HTE can use one of these models to parameterize the outcome models $\hat{\mu}^0(\ensuremath{{\mathbf x}})$ and $\hat{\mu}^1(\ensuremath{{\mathbf x}})$. We set the target coverage rate of $\hat{\mu}^0(\ensuremath{{\mathbf x}})$ and $\hat{\mu}^1(\ensuremath{{\mathbf x}})$ as $\sqrt{1-\alpha}$.  As in the previous section, we can construct a confidence interval for each estimator by using the split conformal regression.
		Let us denote the $\sqrt{1-\alpha}$ confidence intervals for $Y(1)$ and $Y(0)$ by
		$
		\hat{C}^1(\ensuremath{{\mathbf x}})=\left[\hat{\mu}^1(\ensuremath{{\mathbf x}})-\hat{Q}^1_{\sqrt{1-\alpha}},\hat{\mu}^1(\ensuremath{{\mathbf x}})+\hat{Q}^1_{\sqrt{1-\alpha}}\right]
		$ and 
		$
		\hat{C}^0(\ensuremath{{\mathbf x}})=\left[\hat{\mu}^0(\ensuremath{{\mathbf x}})-\hat{Q}^0_{\sqrt{1-\alpha}},\hat{\mu}^0(\ensuremath{{\mathbf x}})+\hat{Q}^0_{\sqrt{1-\alpha}}\right],
		$
		respectively.
		We set the confidence interval for $\tau(\ensuremath{{\mathbf x}})$ to be
		\[
		\hat{C}^\tau(\ensuremath{{\mathbf x}})=\left[\hat{\mu}^1(\ensuremath{{\mathbf x}}) - \hat{\mu}^0(\ensuremath{{\mathbf x}})-\hat{Q}^1_{\sqrt{1-\alpha}}-\hat{Q}^0_{\sqrt{1-\alpha}},\hat{\mu}^1(\ensuremath{{\mathbf x}}) - \hat{\mu}^0(\ensuremath{{\mathbf x}})+\hat{Q}^1_{\sqrt{1-\alpha}}+\hat{Q}^0_{\sqrt{1-\alpha}}\right].
		\]
		This confidence interval ensures the coverage rate $1-\alpha$ for the estimated ITE $\hat{\tau}(\ensuremath{{\mathbf x}})=\hat{\mu}^1(\ensuremath{{\mathbf x}})-\hat{\mu}^0(\ensuremath{{\mathbf x}})$, because its upper endpoint is given as the difference between the upper endpoint of $\hat{C}^1$ and the lower endpoint of $\hat{C}^0$, and its lower endpoint is given as the difference between the lower endpoint of $\hat{C}^1$ and the upper endpoint of $\hat{C}^0$. If the coverage rates for $\hat{C}^1$ and $\hat{C}^0$ are $\sqrt{1-\alpha}$, the coverage rate for $\hat{C}^\tau$ will be $1-\alpha$.
		
		From the ITE estimator $\hat{\tau}(\ensuremath{{\mathbf x}})$ and its confidence interval $\hat{C}^\tau_l(\ensuremath{{\mathbf x}})$ for each subgroup $l$, we can calculate $W_l$ and $S_l$ and  adapt the R2P method to HTE estimation. The robust partitioning problem for HTE in \eqref{eqn:objective} is solved by applying the R2P method in Algorithm \ref{alg:algorithm}, with two minor changes: 1) each sample in the HTE dataset is a triple $(\ensuremath{{\mathbf x}}_i,t_i,y_i)$ consisting of the covariate vector, the treatment indicator, and the observed outcome; 2) the outcome model $\hat{\mu}(\ensuremath{{\mathbf x}})$ in R2P is replaced by the ITE estimator 
		$\hat{\tau}(\ensuremath{{\mathbf x}})$.  As before, we show that this procedure achieves the specified coverage guarantee.   The proof is  provided in the supplementary material.
		
		\begin{theorem}\label{thm2}
			Given a prescribed miscoverage rate $\alpha$, the created partition $\Pi$,  estimator $\hat{\tau}_\ensuremath{{\mathcal X}}$, and  confidence interval function $\hat{C}_l^\tau(\cdot)$ have the following property: for each $l \in \Pi$ and for new samples $(\ensuremath{{\mathbf x}},\tau)$ drawn from the truncated distribution of the subgroup $l$, we have  $\mathbb{P}[\tau\in\hat{C}_l^\tau({\ensuremath{{\mathbf x}}})] \geq 1-\alpha$.
		\end{theorem}
	\end{paragraph}
	
	
\begin{paragraph}{Well-identified subgroups and false discovery} Theorem \ref{thm2} guarantees that confidence intervals achieve the required finite sample coverage for the ITE estimates in each subgroup, regardless of how (in)accurate the underlying ITE estimator $\hat{\tau}_\ensuremath{{\mathcal X}}$ is. If the confidence intervals exhibit large overlap across the constructed subgroups, we can conclude that the constructed subgroups are not well-identified. Conversely, if the confidence intervals have little or no overlap across subgroups, we conclude that the subgroups are well-identified. Given the theoretical guarantee of the confidence intervals in R2P, the subgroups are robust against false discoveries if they are well-identified. \end{paragraph}
	


	
	
	
	\vspace{-0.1in}
	\section{Related Work} \vspace{-6pt}
	Subgroup analysis methods with recursive partitioning have been widely studied based on regression trees (RT) \cite{athey2016recursive,su2009subgroup,tran2019learning,johansson2018interpretable}.
	In these methods, the subgroups (i.e., leaves in the tree structure) are constructed and the individualized outcome or treatment effects are estimated by the corresponding sample mean estimator to the leaf for given covariates.
	To overcome the limitations of the traditional trees to represent the non-linearity such as interactions between treatment and covariates \cite{doove2014comparison},
	a parametric model is integrated into regression trees for subgroup analysis \cite{seibold2016model}.
	However, such approach can be used only for the limited types of estimator models, which is particularly undesirable since advanced causal inference models based on deep neural networks or multi-task Gaussian processes have been studied which outperform the traditional estimators \cite{alaa2017bayesian,yoon2018ganite,shalit2017estimating,zhang2020learning}.
	The global model interpretation method in \cite{yang2018global} can analyze the subgroup structure of arbitrary models but it depends on local model interpreters and does not consider the treatment effects.
	
	
	
	For recursive partitioning, various criteria have been proposed.
	In the traditional RT, the criterion based on the mean squared error between the sample mean estimations from the training samples and the test samples is used \cite{johansson2018interpretable}, and it is referred to as the adaptive criterion.
	Based on the adaptive criterion, an honest criterion is proposed in \cite{athey2016recursive} by splitting the training samples into the training set and the estimation set to eliminate the bias of the adaptive criterion.
	In addition, a generalization cost is introduced to the adaptive or honest criterion in \cite{tran2019learning} to encourage generalization of the analysis.
	The interaction measure between the treatment and covariates is used as a partitioning criterion in \cite{su2009subgroup} and the parameter instability of the parametric models is used in \cite{seibold2016model}.
	In \cite{yang2018global}, the contribution matrix of the samples from local model interpreters is used.
	Some of these criteria implicitly consider the confidence of the estimation by the variance, but most of them do not provide a valid confidence interval of the estimation for each subgroup.
	In \cite{johansson2018interpretable}, a conformal regression method based on regression trees that provides the confidence interval is proposed, but the adaptive criterion is used for partitioning without any consideration of the confidence interval.
	
	
	\vspace{-0.1in}
	\section{Experiments} \vspace{-6pt}
	In this section, we evaluate R2P-HTE by comparing its performance with state-of-the-art HTE methods.
	Specifically, we compare R2P-HTE with four baselines:
	standard regression trees for causal effects (CT-A) \cite{breiman1984classification}, conformal regression trees for causal effects (CCT) \cite{johansson2018interpretable}, causal trees with honest criterion (CT-H) \cite{athey2016recursive}, and causal trees with generalization costs (CT-L) \cite{tran2019learning}.
	We implement CCT and CT-A by modifying the conformal regression tree and conventional regression tree methods for causal effects.  
	Details of the baseline algorithms are provided in the supplementary material.
	For the ITE estimator of R2P-HTE, here abbreviate as R2P, we use the causal multi-task Gaussian process (CMGP) \cite{alaa2017bayesian}.
	Because individual ground-truth treatment effects can never be observed in real data, we use two synthetic and two semi-synthetic datasets.  The first synthetic dataset (Synthetic dataset A) is the simple one proposed in \cite{athey2016recursive}.  Because  dataset A possesses little of the homogeneity within subgroups that is often found in the real world, we offer a second synthetic dataset B that possesses  greater homogeneity within subgroups and greater heterogenity across subgroups  and has many more features than A.  B was inspired by the initial clinical trial of remdesivir \cite{wang2020remdesivir} as a treatment for COVID-19 and uses the patient features listed in that trial, but not the data. In specific, it represents a discovery that remdesivir results in a faster time to clinical improvement for the patients with the shorter time from symptom onset to starting trial. 
	Aside from inspiration and features, B is relevant to COVID-19 only in that the COVID-19 is known to be a disease which displays in very heterogeneous ways.  The two semi-synthetic datasets are based on real world data; the first uses the Infant Health and Development Program (IHDP) dataset \cite{hill2011bayesian} and the second uses the Collaborative Perinatal Project (CPP) dataset \cite{dorie2019automated}.  Details of all these datasets are provided in the supplementary material.
	For each experiment, we conduct 50 simulations.  
	
The ``optimal'' ground-truth of subgroups depends on multiple objectives, including homogeneity, heterogeneity, and the number of subgroups. In the literature, a commonly adopted metric is the variance, because greater heterogeneity across subgroups and homogeneity within each subgroup generally imply well-discriminated subgroups.
	We denote the set of test samples by $\ensuremath{{\mathcal D}}^{te}$ and the test samples that belong to subgroup $l$ as $\ensuremath{{\mathcal D}}^{te}_l$.  We define the mean and variance of the treatment effects of the test samples in subgroup $l$ as $\textrm{Mean}(\ensuremath{{\mathcal D}}^{te}_l)$ and $\textrm{Var}(\ensuremath{{\mathcal D}}^{te}_l)$, respectively.
	We define {\em heterogeneity across  subgroups} to be the variance of the mean of the treatment effects: 
	$V^{\textrm{across}}=\textrm{Var}(\{\textrm{Mean}(\ensuremath{{\mathcal D}}^{te}_l)\}_{l=1}^{L})$, where $L$ is the number of subgroups; we define the {\em average in-subgroup variance} to be 
	$
	V^{\textrm{in}}(\ensuremath{{\mathcal D}}^{te})=\frac{1}{L}\sum_{l=1}^L \textrm{Var}(\ensuremath{{\mathcal D}}^{te}_l)
	$.
	We also provide  the average number of subgroups for better understanding of the results.
	We set the miscoverage rate to be $\alpha=0.05$, so we demand a 95\% ITE coverage rate.  (We do not report actual coverage rates here because all methods achieve the target coverage rate, but they are reported in the supplementary material, along with other details.)
	
	
	\begin{table}[ht]
		\caption{\textbf{Comparison of methods}~~ For the measures $V^{\textrm{across}}$ and $V^{\textrm{in}}$, and for the widths of confidence intervals, we highlight the best results in bold.}
		\footnotesize
		\label{table:experiments}
		\begin{center}
			\setlength\tabcolsep{3pt}
			\begin{tabular}{c|c|c|c|c|c|c|c|c}
				\toprule
				& \multicolumn{4}{c|}{Synthetic dataset A} & \multicolumn{4}{c}{Synthetic dataset B} \\
				\midrule
				& $V^{\textrm{across}}$ & $V^{\textrm{in}}$  & \# SGs   & CI width      & $V^{\textrm{across}}$ & $V^{\textrm{in}}$  & \# SGs   & CI width      \\
				\midrule
			
				\textsc{R2P}      & \textbf{0.22$\pm$.01} & \textbf{0.03$\pm$.001} & 4.9$\pm$.16 & \textbf{0.08$\pm$.003}
				& \textbf{2.39$\pm$.04} & \textbf{0.12$\pm$.01} & 5.0$\pm$.16 & \textbf{0.88$\pm$.06} \\
				\textsc{CCT} 	  & 0.18$\pm$.02 & 0.05$\pm$.01 & 4.4$\pm$.24 & 7.42$\pm$.48
				& 1.97$\pm$.14 & 0.58$\pm$.15 & 5.0$\pm$.13 & 5.95$\pm$.59 \\
				\textsc{CT-A}     & 0.19$\pm$.02 & 0.04$\pm$.01 & 4.7$\pm$.21 & 3.96$\pm$.16
				& 2.24$\pm$.06 & 0.30$\pm$.05 & 5.1$\pm$.15 & 2.77$\pm$.20  \\
				\textsc{CT-H}     & 0.12$\pm$.03 & 0.11$\pm$.02 & 3.1$\pm$.39 & 4.39$\pm$.22
				& 2.01$\pm$.13 & 0.53$\pm$.13 & 4.5$\pm$.15 & 3.38$\pm$.32  \\
				\textsc{CT-L}     & 0.12$\pm$.02 & 0.10$\pm$.02 & 2.9$\pm$.35 & 5.22$\pm$.02
				& 0.80$\pm$.26 & 1.77$\pm$.27 & 3.1$\pm$.28 & 6.92$\pm$.53 \\
				\midrule
				\midrule
				& \multicolumn{4}{c|}{IHDP dataset}& \multicolumn{4}{c}{CPP dataset} \\
				\midrule
				& $V^{\textrm{across}}$ & $V^{\textrm{in}}$  & \# SGs   & CI width     & $V^{\textrm{across}}$ & $V^{\textrm{in}}$  & \# SGs   & CI width     \\
				\midrule
			
				\textsc{R2P}      & \textbf{0.46$\pm$.04} & \textbf{0.38$\pm$.03} & 4.1$\pm$.12 & \textbf{1.27$\pm$.22}
				& \textbf{0.06$\pm$.02} & \textbf{0.10$\pm$.01} & 5.7$\pm$.30 & \textbf{1.11$\pm$.13}\\
				\textsc{CCT} 	  & 0.30$\pm$.04 & 0.43$\pm$.05 & 4.3$\pm$.13 & 5.70$\pm$.23
				& 0.03$\pm$.02 & 0.12$\pm$.01 & 6.4$\pm$.20 & 3.60$\pm$.12 \\
				\textsc{CT-A}     & 0.31$\pm$.04 & 0.57$\pm$.05 & 4.1$\pm$.08 & 3.71$\pm$.08
				& 0.03$\pm$.01 & 0.12$\pm$.01 & 6.6$\pm$.18 & 2.45$\pm$.06 \\
				\textsc{CT-H}     & 0.28$\pm$.05 & 0.56$\pm$.05 & 3.8$\pm$.14 & 3.76$\pm$.14
				& 0.01$\pm$.00 & 0.14$\pm$.01 & 5.2$\pm$.23 & 2.67$\pm$.06 \\
				\textsc{CT-L}     & 0.27$\pm$.06 & 0.64$\pm$.05 & 2.8$\pm$.23 & 4.75$\pm$.15
				& 0.01$\pm$.01 & 0.14$\pm$.01 & 2.9$\pm$.29 & 3.23$\pm$.07 \\
				\bottomrule
			\end{tabular}
		\end{center}
	\end{table}
	\vspace{-6pt}
	
	
	\begin{paragraph}{Results}
		Table \ref{table:experiments} reports the performances of R2P and the baselines for all four  datasets.  Keep in mind that {\em larger} $V^{\textrm{across}}$ means  greater heterogeneity  {\em across}  subgroups, while {\em smaller} $V^{\textrm{in}}$ means  greater  homogeneity {\em within}  subgroups.  
		As the table shows, R2P displays by far the best performance on all four datasets:  the greatest heterogeneity across  subgroups, the greatest homogeneity within subgroups, and the narrowest confidence intervals.  It accomplishes all this while identifying comparable numbers of subgroups.  We conclude that R2P identifies subgroups more effectively than any of the other methods, and allows fewer false discoveries (as was illustrated in the Introduction).  The performance  of R2P reflects one of its strengths: the ability to use {\em any} method of estimating ITE.
		
		The  effectiveness of R2P can also be seen in Fig. \ref{fig:boxplot_synthetic}, which provides, for R2P and each of the four baseline algorithms, boxplots of the distribution of the treatment effects for each identified subgroup for Synthetic dataset B.  We see that that R2P identifies subgroups reliably: different subgroups have very different average treatment effects and their distributions are non-overlapping or well-discriminated.  All the other methods are unreliable: false discovery  occurs for all four baseline methods, and occurs frequently for three of the four.  
	\end{paragraph}
	
	
	\begin{figure}[t]
		\vspace{-0.5em}
		\centering
		\includegraphics[width=1.\linewidth]{figs/sim/boxplot_synthetic.eps}
		\caption{Treatment effects for subgroups identified by each algorithm when applied to Synthetic dataset B. Each box represents the range between the 25th and 75th percentiles of the treatment effects on the test samples; each whisker represents the range between the 5th and 95th percentiles.
			\label{fig:boxplot_synthetic}}
		\vspace{-1.5em}
	\end{figure}
	
	\vspace{-6pt}
	\begin{paragraph}{Gain from subgroup analysis}
		
		\begin{wraptable}{r}{7.1cm}
		
			\caption{Normalized $V^{\textrm{in}}$ of R2P\label{table:gain_subgroup}}
			\small
			\centering
			\setlength\tabcolsep{3pt}
			\begin{tabular}{c|c|c|c}
				\toprule  
				Synth A & Synth B & IHDP & CPP \\
				\midrule
				0.110$\pm$.005 & 0.046$\pm$.003 & 0.459$\pm$.033 & 0.691$\pm$.076 \\
				\bottomrule
			\end{tabular}
		
		\end{wraptable}
		To indicate the gain from subgroup analysis obtained by R2P, and hence to indicate the effectiveness of recursive partitioning, we compare $V^{\textrm{in}}$, the homogeneity within subgroups obtained by  R2P in Table \ref{table:experiments}, against the homogeneity within the entire population, $V^{\textrm{pop}}$. We divide $V^{\textrm{in}}$ by $V^{\textrm{pop}}$ to obtain the normalized $V^{\textrm{in}}$ in Table \ref{table:gain_subgroup}. Subgroup analysis with R2P reduces the average in-subgroup variance by 89\% and more than 95\% on Synthetic dataset A and B, respectively. For the semi-synthetic datasets, it reduces the average in-subgroup variance by more than 50\% and 30\% on the IHDP and CPP datasets, respectively. (Keep in mind that R2P constructs subgroups in a way that produces {\em both} strong heterogeneity across subgroups and strong homogeneity within subgroups.) 
	
		
		
	\end{paragraph}
	
	
	
	\vspace{-0.1in}
	\section{Conclusion} \vspace{-6pt}
	In this paper, we have studied robust HTE analysis based on recursive partitioning.  The algorithm proposed, R2P-HTE,  
	recursively partitions the entire population by taking into account both  heterogeneity across subgroups and homogeneity within subgroups, using the novel criterion of confident homogeneity that is based on the  quantification of uncertainty of the ITE estimation.  R2P-HTE robustly constructs subgroups and also provides confidence guarantees  for each subgroup.  Experiments for synthetic and semi-synthetic datasets (the latter based on real data) demonstrate that R2P-HTE outperforms state-of-the-art baseline algorithms in every dimension: greater heterogeneity across subgroups, greater homogeneity within subgroups, and  narrower confidence intervals.  One of the strengths of R2P-HTE is that it can employ {\em any} method for interpretable ITE estimation, including improved methods that will undoubtedly be developed in the future. 
	
	
	
	\section*{Broader Impact} The understanding of treatment effects plays an important role in many areas, and especially in medicine and public policy.  In both areas, it is often the case that the same treatment has different effects on different groups; hence subgroup analysis is called for.  In medicine, subgroup analysis may make it possible to identify groups of patients (defined by covariates such as age, body mass index, blood pressure, etc.) suffering from a particular disease for whom a particular drug is effective and safe and other groups for whom the same drug is ineffective and unsafe.  Similarly, subgroup analysis may make it possible to identify groups of patients for whom one course of treatment (e.g. a particular mode of radiotherapy or chemotherapy) is preferable (more likely to be successful with fewer side effects) to another.  In public policy, subgroup analysis may make it possible to identify groups of people or geographic regions for which particular interventions (e.g., providing mosquito nets to combat malaria) are likely to be successful or unsuccessful.  The method for subgroup analysis that is developed in this paper, R2P, is an enormous improvement over state-of-the-art methods and therefore has the potential to make  enormous and widespread impact.  Moreover, because R2P can make use of improvements in the underlying estimation methods, this impact may grow over time.
	

		\section*{Acknowledgments}
		This work was supported by GlaxoSmithKline (GSK), the US Office of Naval Research (ONR), and the National Science Foundation (NSF): grant numbers 1407712, 1462245, 1524417, 1533983, 1722516. CS acknowledges the funding support from Kneron, Inc. We thank all reviewers for their comments and suggestions.



