\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.


\section{Introduction}

Data augmentation is an effective way to improve the performance of CNN models for image recognition tasks, particularly when its policy is optimized for the target dataset. Conventional data augmentation for images consists of image transformation operations, such as random cropping and flipping, and color enhancing including modification of color intensities \cite{Krizhevsky2012,He2016b}. However, designing good data augmentation strategies requires profound understanding of the target data and operations.
For example, CutOut \cite{DeVries2017} randomly erases a patch region of each image and is known to improve performance on the CIFAR-10 dataset, but is also reported to degrade the performance on other datasets, e.g., ImageNet \cite{Lopes2019}.

Therefore, automatically designing effective augmentation strategies according to target data and tasks is desirable to improve the performance of image recognition models. One approach to augment existing data is to generate new data samples by interpolating several training images \cite{Hauberg2015,Devries2017,Zhang2018} or by using conditional generative models \cite{Antoniou2018,Baluja2018}. However, this approach requires a large amount of labeled data \cite{Rather2017} and sometimes fails to improve the performance \cite{Ravuri2019}, even if powerful conditional generative models are used.
On the other hand, some methods improve the performance by efficiently selecting effective combinations of image transformation operations from exponentially large candidate pools \cite{Rather2017,Cubuk2018a,Zhang2020}. In particular, AutoAugment\xspace \cite{Cubuk2018a} and its family \cite{Ho2019,Lim2019,Hataya2019a,Lin2020,Cubuk2019,Berthelot2019a,Lin2020} optimized combinations of operations to improve validation performance and achieved state-of-the-art results.

These methods involve a bi-level optimization: the inner process optimizes parameters of a CNN on training data using a given combination of operations, and the outer process optimizes the combination of operations to maximize the validation performance. Particularly, the inner loop, i.e., training of a CNN, is usually expensive. Therefore, prior works use proxy tasks that adopt small subsets of training datasets or small models \cite{Cubuk2018a,Ho2019,Lim2019,Hataya2019a} or reduce the search space \cite{Cubuk2019,Berthelot2019a,Lin2020} to keep the entire training feasible.
Such approaches may result in sub-optimal solutions. We will further review and formalize this problem in \Cref{sec:generalizing_dao}.

In this paper, we tackle the original bi-level optimization problem directly without using proxy tasks or reducing the search space. We propose Meta Approach to Data Augmentation Optimization\xspace (MADAO\xspace), which optimizes CNNs and augmentation policies simultaneously by using gradient-based optimization. Here, policies are updated so that they directly increase CNNs' performance.
Na\"ively applying gradient-based optimization to this bi-level optimization requires differentiation through the inner optimization process \cite{Finn2017b} or computation of the inverse Hessian matrix \cite{Bengio2000}, both of which suffer from large space complexity. These problems are fatal because data augmentation optimization needs to handle large networks, e.g., CNNs for ImageNet. We bypass these issues by using the implicit gradient method with Neumann series approximation.
Thanks to these approximations, MADAO\xspace is simple with little overhead as shown in \Cref{ls:bilevel_optimization}. Notably, this simplicity allows MADAO\xspace to scale to problems of ImageNet size, which has been nearly impossible for existing bi-level optimization methods \cite{hutter19}.

We empirically demonstrate that MADAO\xspace learns effective data augmentation policies and achieves performance comparable or even superior to existing methods on benchmark datasets for image classification: CIFAR-10, CIFAR-100, SVHN and ImageNet, as well as four fine-grained datasets. In addition, we show that MADAO\xspace improves the performance of self-supervised learning on ImageNet. All of the reported results have been achieved \textit{without using dataset-specific configurations}.

\newcommand{\highlig}[1]{\colorbox[rgb]{1.0, 0.8, 0.8}{#1}}
\begin{listing}[t]
\centering

\begin{minted}[%
linenos,
xleftmargin=20pt,
breaklines,
frame=lines,
mathescape,
escapeinside=||,
]{python}
cnn.initialize() # parameterized by ${\bm \theta}$
|\highlig{policy.initialize()}| # parameterized by ${\bm \phi}$
for epoch in range(num_epochs):
    for train_data, val_data in data_loader:
        for i in range(num_inner_iters):
            input = |\highlig{policy(}|train_data[i]|\highlig{)}| # augment data by policy
            criterion = cnn.train(input) # referred to $f$ in the text
            cnn.update(criterion)
        vcriterion = cnn.val(val_data) # referred to $g$ in the text
        |\highlig{policy.update(vcriterion)}|
\end{minted}
\vspace{-0.8\baselineskip}
\caption{Our proposed method MADAO\xspace in a Python-like pseudo code. MADAO\xspace optimizes \mintinline[]{python}{policy} as well as \mintinline[]{python}{cnn} by gradient descent with \textit{slight modifications} (\highlig{\texttt{highlighted}}) and \textit{little overhead} to the standard optimization protocol.}
\label{ls:bilevel_optimization}
\end{listing}

\section{Generalizing Data Augmentation Optimization} \label{sec:generalizing_dao}

In \Cref{sub:optimization_da}, we describe the preliminaries of data augmentation optimization of AutoAugment\xspace family, and then review prior methods by generalizing the problem in \Cref{sub:generalization}.

\subsection{Designing Data Augmentation Space}\label{sub:optimization_da}

Let us define a set of input images $\mathcal{X}$ and a set of operations $\mathcal{S}$ consisting of data augmentation operations such as rotation and color inversion. In the AutoAugment\xspace family, each image ${\bm{x}}\in\mathcal{X}\subset[0, 1]^D$ is augmented by an operation $O: \mathcal{X}\to\mathcal{X}$ with a probability of $p_O\in[0, 1]$ and a magnitude of $\mu_O\in[0, 1]$ as illustrated in \Cref{fig:subpolicies}. The magnitude parameter can correspond for example to the degree of rotation, while some operations, such as inversion, have no magnitude parameter. By applying $K$ consecutive augmentation operations, each image results in $2^K$ possible images, that is, the number of images virtually increases. This formulation makes the size of the search space $(\lvert\mathcal{S}\rvert\times[0, 1]\times[0, 1])^{K}$, where $\lvert\mathcal{S}\rvert$ is the size of the operation set.

Operations and accompanied parameters need to be selected so that they minimize the validation criterion, such as the error rate. Usually, this selection is performed heuristically \cite{Krizhevsky2012,He2016b}. However, \citeauthor{Cubuk2018a} showed that data-driven optimization surpasses handcrafted selection \cite{Cubuk2018a}.

\begin{figure}
    \centering
    \includegraphics[width=0.65\linewidth]{subpolicy2}
    \vspace{-0.7\baselineskip}
    \caption{In AutoAugment\xspace family, a \mintinline[]{python}{policy} consists of operations. Each operation, e.g., $\texttt{Rotate}$ or $\texttt{Invert}$, augments an image with a probability of $p$ and a magnitude of $\mu$. As can been seen, performing multiple operations virtually increases the diversity of images. The operations are selected from given operation sets. Our proposed method MADAO\xspace learns to select suitable operations and their parameters, and simultaneously optimize a CNN by using gradient descent. This figure shows a 2-stage augmentation.}
    \label{fig:subpolicies}
\end{figure}

\subsection{Generalizing  AutoAugment Family}\label{sub:generalization}

Let ${\bm \theta}\in\mathbb{R}^M$ denote parameters of a CNN model and ${\bm \phi}\in\mathbb{R}^N$ denote parameters of a policy for augmentation, i.e., selection of operations from the operation set and their accompanied parameters $\{(\mu_O, p_O); O\in\mathcal{S}\}$. Let the empirical risk be $f({\bm \theta}, {\bm \phi}; \mathcal{D}_T)$ and the validation criterion be $g({\bm \theta}; \mathcal{D}_V)$. $\mathcal{D}_T$ and $\mathcal{D}_V$ are training and validation datasets.
In the case of classification task, $\mathcal{D}\coloneqq\{({\bm{x}}_i, y_i)\}\in\mathcal{X}\times \{1, 2, \dots, C\}$, where $C$ is the number of classes.

Optimization of data augmentation policy in AutoAugment\xspace family methods can be generalized as

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:bi_level_optimization}
    \argmin_{{\bm \phi}}g(\argmin_{{\bm \theta}}f({\bm \theta}, {\bm \phi}; \mathcal{D}_T); \mathcal{D}_V),
\end{equation}

\noindent that is, optimizing CNNs on training data with policies that minimize validation criteria on validation data.

Na\"ively solving this bi-level minimization problem takes a long time because CNN training $\min_{{\bm \theta}}f({\bm \theta}, {\bm \phi})$ is costly and the number of possible combinations of augmentation operations and their parameters is infeasibly large. Therefore, prior works tried to alleviate this problem in several ways. One direction is to reduce the search space over augmentation policies, i.e., dimension and range of the parameter ${\bm \phi}$ \cite{Cubuk2019,Berthelot2019a,Lin2020}. For example, RandAugment\xspace \cite{Cubuk2019} randomly samples operations from the operation set $\mathcal{S}$ and shares $\mu_{O}$ among all operations. This reduction changes the outer problem in \Cref{eq:bi_level_optimization} from $\argmin_{\bm \phi} g$ to $\argmin_\mu g$, which makes it possible to use a simple searching process, such as grid search. OHL-AutoAug \cite{Lin2020} enables online searching using policy gradient \cite{Williams1992} by restricting the search space only to a limited range.

On the other hand, some methods use proxy tasks that approximate \Cref{eq:bi_level_optimization} to obtain (sub-) optimal policy ${\bm \phi}'$ to reduce the searching time of the inner optimization. The obtained policy ${\bm \phi}'$ is then used to train a CNN as $\min_{{\bm \theta}}f({\bm \theta}, {\bm \phi}')$ \cite{Cubuk2018a,Ho2019,Lim2019,Hataya2019a} in an ``offline'' manner.
For instance, AutoAugment\xspace employs a proxy task $f'$ that approximates the original inner problem $f$ as

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:proxy_task}
     {\bm \phi}'=\argmin_{{\bm \phi}}g(\argmin_{{\bm \theta}'}f'({\bm \theta}', {\bm \phi}'; \mathcal{D}'); \mathcal{D}_V),
\end{equation}

\noindent with a smaller dataset $\lvert\mathcal{D}'\rvert\ll\lvert\mathcal{D}_T\rvert$ and a smaller network $\dim{\bm \theta}'\ll\dim{\bm \theta}$ for efficiency. The outer problem $g$ is optimized by black-box optimization techniques, such as reinforcement learning \cite{Cubuk2018a}. Fast AutoAugment\xspace \cite{Lim2019} and Faster AutoAugment\xspace \cite{Hataya2019a} approximate \Cref{eq:bi_level_optimization} as minimizing distance of distributions between augmented images and original images without directly minimizing $f$, which allows faster searching.

To summarize, prior methods indirectly solve the bi-level optimization problem in \Cref{eq:bi_level_optimization}, as displayed in \Cref{tab:ourmethod}. We instead propose to tackle this problem directly.

\begin{table}[t]
    \centering
    \caption{Comparison of prior methods and our proposal, MADAO\xspace. \textbf{MADAO\xspace can efficiently search \textit{full space} for policies on \textit{full datasets} with \textit{full CNN models}}, e.g., on ImageNet with ResNet-50.}
    \label{tab:ourmethod}
    \vspace{-0.5\baselineskip}
    \begin{tabular}{lcc}
       \toprule
           Method                            & Direct Inner Problem             & Direct Outer Problem \\  \midrule
           AutoAugment\xspace \cite{Cubuk2018a}          &                                  & \checkmark           \\
           Fast AutoAugment\xspace \cite{Lim2019}        &                                  & \checkmark           \\
           \midrule
           OHL-AutoAug \cite{Lin2020}        & \checkmark           & \\
           RandAugment\xspace \cite{Cubuk2019}           & \checkmark           & \\
           \midrule
           MADAO\xspace (ours)            & \checkmark           & \checkmark\\
       \bottomrule
    \end{tabular}
\end{table}

\section{Method} \label{sec:method}

In this paper, we propose to directly optimize the bi-level optimization problem in \Cref{eq:bi_level_optimization}. In other words, we optimize CNNs and augmentation policies \textit{simultaneously}, i.e., in an \textit{online} manner, without reducing the search space or using proxy tasks. With this simultaneous optimization, policies are expected to augment images according to the learning state of CNN models. \Cref{ls:bilevel_optimization} is a simple depiction of our approach, which we call Meta Approach to Data Augmentation Optimization\xspace, or in short, MADAO\xspace.

\subsection{Optimizing Policies by Gradient Descent}

MADAO\xspace directly optimizes the bi-level problem \Cref{eq:bi_level_optimization} using gradient descent. To this end, we assume that  $f$ and $g$ are differentiable w.r.t.\xspace ${\bm \theta}$. Taking $C$-category classification as an example, cross entropy $-\mathbb{E}_{{\bm{x}}_i, y_i}\log [{\bm{h}}_{{\bm \theta}}({\bm{x}}_i)]_{y_i}$ can be used as $f$ and $g$, but error rate $\mathbb{E}_{{\bm{x}}_i, y_i}\mathbbm{1}(\argmax {\bm{h}}_{{\bm \theta}}({\bm{x}}_i) \neq y_i)$ cannot. Here, $\mathbbm{1}$ is the indicator function, and ${\bm{h}}_{{\bm \theta}}: \mathbb{R}^D\to\mathbb{R}^C$ is a CNN with a softmax output layer.

Gradient-based optimization of \Cref{eq:bi_level_optimization} requires $\nabla_{{\bm \phi}}g$ for iterative updating. Since the data augmentation implicitly affects the validation criterion, in other words, data augmentation is not used for validation, we obtain

\vspace{-0.6\baselineskip}
\begin{equation}\label{eq:grad_phi_g}
    \frac{\partial g}{\partial {\bm \phi}}=\frac{\partial g}{\partial {\bm \theta}}\frac{\partial {\bm \theta}}{\partial {\bm \phi}}.
\end{equation}

Because of the requirement of $g$, $\nabla_{{\bm \theta}}g$ can be obtained. On the other hand, the exact computation of $\nabla_{{\bm \phi}}{\bm \theta}$ has a large space complexity, as we will describe in \Cref{eq:unroll_steps}. Yet, if this gradient was available, the policies could be optimized by gradient descent.

\subsection{Approximating Gradients of Policy and Inverse Hessian}

To obtain $\nabla_{{\bm \phi}}{\bm \theta}$ without suffering from a large space complexity, we can use the Implicit Function Theorem. If there exists a fixed point $({\bm \theta}^{\star}, {\bm \phi}^{\star})$ that satisfies $\nabla_{{\bm \theta}}f({\bm \theta}^{\star}, {\bm \phi}^{\star})=0$, then there exists a function $\hat{\theta}$ around ${\bm \phi}^{\star}$ such that $\hat{{\bm \theta}}({\bm \phi}^\star)={\bm \theta}^{\star}$. If this condition holds, we also obtain

\vspace{-0.6\baselineskip}
\begin{equation}\label{eq:ift}
    \frac{\partial \hat{{\bm \theta}}}{\partial {\bm \phi}} = -\left( \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \theta}^{\top}} \right)^{-1} \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \phi}^{\top}}\Big|_{{\bm \theta}^\star, {\bm \phi}^\star}.
\end{equation}

We can obtain an approximated gradient using this property. Unfortunately, $M=\dim{\bm \theta}$ is usually large; therefore, computing the inverse Hessian matrix $(\nabla^2_{{\bm \theta}}f)^{-1}$ is prohibitively expensive as it usually requires $\mathcal{O}(M^3)$ computations. To avoid computing the inverse Hessian matrix, we use iterative methods based on the Neumann series, which boast better scalability than conjugate gradients in various problems \cite{Koh2017,Liao2018,Lorraine2019}.

The Neumann series ${\bm{I}}+{\bm{A}}+{\bm{A}}^2+\dots=\sum_i^J {\bm{A}}^j\to({\bm{I}}-{\bm{A}})^{-1} ~(J\to\infty)$ holds with a given squared matrix ${\bm{A}}$ if $\|{\bm{A}}\|<1$. Using this property, \Cref{eq:grad_phi_g} can be approximated with a positive integer $J$ as

\vspace{-0.75\baselineskip}
\begin{equation}
    \frac{\partial g}{\partial {\bm \phi}} =
   -\frac{\partial g}{\partial {\bm \theta}}
    \left( \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \theta}^{\top}}\right)^{-1}
    \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \phi}^{\top}}
    \approx
   -\frac{\partial g}{\partial {\bm \theta}}
    \sum_{j=0}^{J} \left({\bm{I}}-\frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \theta}^{\top}} \right)^j
    \frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \phi}^{\top}}.
\end{equation}

We regularize the norm by simply introducing a scalar $\alpha\in\mathbb{R}^+$ as ${\bm{I}}-\alpha\nabla_{{\bm \theta}}^2 f$ as \cite{Lorraine2019}. We use $\alpha=10^{-3}$ with $J=5$ in the experiments. This Neumann series approximation can also help us avoid explicitly storing the Hessian matrix, whose space complexity is $\mathcal{O}(M^2)$, by using Hessian-vector products derived from the following identity:

\vspace{-0.75\baselineskip}
\begin{equation}
    {\bm{v}}^{\top}\left(\frac{\partial^2 f}{\partial {\bm \theta}\partial {\bm \theta}^{\top}}\right)
   =\frac{\partial}{\partial{\bm \theta}}\left({\bm{v}}^{\top}\frac{\partial f}{\partial{\bm \theta}}\right)\hphantom{xxxx}(\forall{\bm{v}}\in\mathbb{R}^M).
\end{equation}

This right-hand side can be used instead of storing the Hessian matrix and only has the space complexity of $\mathcal{O}(2M)$.

\subsection{Differentiable Data Augmentation}\label{sub:dda}

To differentiate through ${\bm \phi}$, MADAO\xspace adopts the differentiable data augmentation pipeline following \cite{Hataya2019a}. As described in \Cref{sub:optimization_da}, each image ${\bm{x}}$ is augmented with an operation $O$ with a magnitude of $\mu_O$ and a probability of $p_O$, which can be written as

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:operation_process}
    {\bm{x}} \mapsto
    \begin{cases}
        O({\bm{x}}; \mu_O)  &\text{ with probability of } p_O \\
        {\bm{x}}            &\text{ with probability of } 1-p_O.
    \end{cases}
\end{equation}

This right-hand side can be written as $bO({\bm{x}}; \mu_O)+(1-b){\bm{x}}$ with a binary stochastic variable $b\sim\text{Bern}(b; p_O)$. Although the original Bernoulli distribution $\text{Bern}(b; p_O)$ is not differentiable w.r.t.\xspace $p_O$, Gumbel trick \cite{Jang2016} relaxes this restriction to enable backpropagation to update $p_O$. Similarly, some color-enhancing operations are non-differentiable w.r.t.\xspace the magnitude $\mu_O$, so the straight-through estimator \cite{Bengio2013} is used for such operations. We clamp $\mu_O$ and $p_O$ by a sigmoid function to limit their range to [0, 1].

MADAO\xspace uses a different method to select operations compared to Faster AutoAugment\xspace in order to accelerate training.
MADAO\xspace selects operations using categorical distribution parameterized by a weight parameter ${\bm \pi}\in[0, 1]^{\lvert\mathcal{S}\rvert}$, where $\sum_i\pi_i=1$. Since the original categorical distribution is non-differentiable as Bernoulli distribution, we use Gumbel-softmax with a temperature of $\tau\in\mathbb{R}^{+}$. This distribution, referred to as $\text{RelaxCat}({\bm \pi}; \tau)$, samples one-hot-like vectors as $\tau\longrightarrow 0$. Using this distribution, an operation is selected and applied as

\vspace{-1.2\baselineskip}
\begin{align}
    {\bm{x}} &\mapsto \frac{u_i}{\text{SG}(u_i)}O_i({\bm{x}}; \mu_{O_i}, p_{O_i}), \label{eq:transformation} \\
    i &= \argmax{\bm{u}}, \\
    {\bm{u}} &\sim \text{RelaxCat}({\bm \pi}; \tau).
\end{align}

Here, $\text{SG}$ is the stop gradient operation, and thus, $\frac{u_i}{\mathrm{SG}(u_i)}=1$ so that the transformation \Cref{eq:transformation} keeps the range in $(0, 1)$.
Different from this approach, Faster AutoAugment\xspace applies all operations and takes the weighted sum of the outputs to approximate this selection.

\subsection{Connection to Gradient-based Hyperparameter Optimization}

As can been seen, \Cref{eq:bi_level_optimization} is a hyperparameter optimization (HO) problem to neural networks. Traditional HO methods, such as grid search, random search \cite{Bergstra2012} and Bayesian optimization \cite{Snoek2012}, have poor scalability to increasing dimensionality of hyperparameters \cite{hutter19}. For that reason, gradient-based HO attracts attention. From HO viewpoint, policy parameters ${\bm \phi}$ are hyperparameters.

As shown in \Cref{eq:grad_phi_g}, we need to obtain $\nabla_{{\bm \phi}}{\bm \theta}$ to optimize the outer problem $g$.  The inner optimization process (\Cref{ls:bilevel_optimization} Line 5-8) can be rewritten as

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:unrolling}
{\bm \theta}_T
={\bm \theta}_{T-1}-\eta\nabla_{{\bm \theta}} f({\bm \theta}_{T-1}, {\bm \phi})
=\cdots=
{\bm \theta}_0-\eta\nabla_{{\bm \theta}}\sum_{t=0}^{T-1} f({\bm \theta}_{t}, {\bm \phi})
\eqqcolon {\bm \theta}_T({\bm \theta}_0, {\bm \phi})
\end{equation}

\noindent after $T$ SGD steps with learning rate of $\eta$. One approach to obtain $\nabla_{{\bm \phi}}{\bm \theta}$ is to unroll the steps in \Cref{eq:unrolling} as  \cite{Maclaurin2015,Franceschi2018,Shaban2018}:

\vspace{-0.75\baselineskip}
\begin{equation}\label{eq:unroll_steps}
    \nabla_{{\bm \phi}}{\bm \theta}\approx\nabla_{{\bm \phi}}{\bm \theta}_T({\bm \theta}_0, {\bm \phi}).
\end{equation}

This unrolling requires to cache ${\bm \theta}_0, \dots, {\bm \theta}_{T-1}$, and thus, the space complexity becomes $\mathcal{O}(TM)$, which might be prohibitive for large neural networks, while MADAO\xspace requires $\mathcal{O}(2M)$.

Alternatively, implicit gradient yields $\nabla_{{\bm \phi}}{\bm \theta}$ as explained in \Cref{eq:ift}. This approach needs the inverse Hessian matrix $(\nabla^2_{{\bm \theta}}f)^{-1}$ \cite{Bengio2000}, but this computation is infeasible for modern neural networks with millions of parameters. Iterative methods, such as conjugate gradient \cite{Do2009,Domke2012,Pedregosa2016a,Rajeswaran} or Neumann series approximation \cite{Lorraine2019}, effectively compensate for this issue by approximating this inverse matrix in gradient hyperparameter optimization. Such iterative approximation methods using the Neumann series are also used in approximating influence function \cite{Koh2017} and enabling RNNs to handle long sequences \cite{Liao2018}. We exploit the knowledge from these prior works and adopt the implicit gradient method with Neumann series approximation to efficiently handle large-scale datasets and CNNs.

\section{Experiments and Results} \label{sec:experiments}

This section describes the empirical results of our proposed method in supervised learning for image classification and self-supervised learning tasks. For classification tasks, we use CIFAR-10, CIFAR-100 \cite{Krizhevsky2009}, SVHN \cite{Netzer2011} and ImageNet (ILSVRC-2012) \cite{Russakovsky2015}. In addition, we also used fine-grained classification datasets: Oxford 102 Flowers \cite{Nilsback08}, Oxford-IITT Pets \cite{parkhi12a}, FGVC Aircraft \cite{maji13} and Stanford Cars \cite{Krause2013}. For the self-supervised learning task, we used ImageNet. In all experiments except those on ImageNet, we set 10 \% of the original training data aside as validation data $\mathcal{D}_V$ and report error rates on test data. For ImageNet, we used 1 \% of the training data as validation data. Note that this data split means that we use less training data than prior works and that changes the performance of baseline models. More details about experiment configurations are in \Cref{ap:exp_details}.

We used 14 operations for augmentation: \texttt{ShearX}, \texttt{ShearY}, \texttt{TranslateX}, \texttt{TranslateY}, \texttt{Rotate}, \texttt{Invert}, \texttt{AutoContrast}, \texttt{Equalize}, \texttt{Solarize}, \texttt{Color}, \texttt{Posterize}, \texttt{Contrast}, \texttt{Brightness} and \texttt{Sharpness} (see \Cref{ap:operations} for more details.) To make these operations differentiable, we implemented them using PyTorch \cite{NIPS2019_9015} and kornia \cite{eriba2019kornia}. We scale the magnitude of operations from $0$ to $1$ so that the magnitude of $0$ means no change and the magnitude of $1$ corresponds to the strongest level of the particular augmentation operation. We initialized the parameters of magnitude and probability with $\text{sigmoid}(0.5)\approx0.62$ and the parameters for operation selection to be equal. That means that MADAO in its initial state is nearly equivalent to RandAugment\xspace with the magnitude of $0.62$. We set the number of augmentation stages $K=2$ in all experiments below (see \Cref{fig:subpolicies}).

For training, we used SGD with a momentum for CNN optimization and RMSprop optimizer with learning rate of $10^{-2}$ for policy optimization, following \cite{Lorraine2019}. \Cref{eq:ift} requires the existence of fixed points that satisfies $\nabla_{{\bm \theta}}f({\bm \theta}, {\bm \phi})=0$. Following \cite{Lorraine2019}, we assume that $s$ iterations, which corresponds to \mintinline[]{python}{num_inner_iters} in \Cref{ls:bilevel_optimization}, makes the parameters hold the condition. To further encourage parameters to satisfy the condition, we also used warm up, i.e., training CNNs without augmentation for the first $w$ epochs. We set $s=30$ and $w=20$ by default. Validation criterion $g$ is the same loss function that is used as the training loss function $f$, e.g., cross-entropy, for simplicity.

Training of WideResNet 28-2 on CIFAR-10 for 200 epochs with MADAO\xspace takes 1.7 hours, while training without MADAO\xspace takes 1.0 hours in our environment. As regards memory consumption, training with MADAO\xspace requires 3.38 GB of GPU memory, while training without MADAO\xspace requires 1.74 GB in this setting. We tune the magnitude parameter of RandAugment\xspace using three runs of random search and report the test error rate for the run with the lowest validation error rate.

\subsection*{CIFAR-10, CIFAR-100 and SVHN}

\Cref{tab:results_small} presents test error rates on CIFAR-10, CIFAR-100 and SVHN with various CNN models: WideResNet-28-2, WideResNet-40-2, WideResNet-28-10 \cite{Zagoruyko} and ResNet-18 \cite{He2016b}. We show the average scores of three runs. For comparison, we present the results of RandAugment\xspace and the default augmentation: Cutout \cite{DeVries2017}, random cropping, and random horizontal flipping (except SVHN), following \cite{Cubuk2019}, which we refer to as Baseline. For CIFAR-10 and CIFAR-100, we also present the results of AutoAugment\xspace, which are trained on the split training set using the policy provided by the authors \footnote{\url{https://github.com/tensorflow/models/tree/master/research/autoaugment}}. RandAugment\xspace and AutoAugment\xspace are selected here as representative methods that limit the search space and use proxy tasks, as discussed in \cref{sub:generalization}.

As can been seen, MADAO\xspace achieves performance superior to baseline methods, which demonstrates the effectiveness of directly solving the bi-level problem rather than using proxy tasks as AutoAugment\xspace or a reduced search space as RandAugment\xspace.

\subsection*{ImageNet}

\Cref{tab:results_imagenet} shows the results on ImageNet with ResNet-50 \cite{He2016b} trained in supervised and self-supervised learning manners. In the supervised learning setting, we trained models for 180 epochs. In addition, to showcase the effectiveness of MADAO\xspace in other tasks than supervised learning, we apply MADAO\xspace to contrastive self-supervised learning \cite{he2019moco} for 100 epochs of training. We report the results of the linear classification protocol.

Most importantly, MADAO\xspace can scale to an ImageNet-size problem without using proxy tasks or reducing search space. We hypothesize that the slightly inferior performance of MADAO\xspace compared to RandAugment\xspace could be attributed to the fact that ImageNet has 1,000 categories and a single policy might be insufficient for such diverse data. We also believe that class- or instance-conditional data augmentation might be required to further improve performance, which we leave as an open problem.

\subsection*{Fine-grained Classification}

To showcase the ability of MADAO\xspace to augment images according to target datasets, we conducted experiments on four fine-grained datasets with ResNet-18 \cite{He2016b}. \Cref{tab:results_finedgrained} shows the average test error rates over three runs. These datasets are nearly ten to twenty times smaller than CIFARs and SVHN, yet MADAO\xspace drastically improves the performance without using dataset-specific hyperparameters.
This improvement emphasizes the importance of searching for good policies according to full target datasets from the full search space. MADAO\xspace can capture the characteristics of each fine-grained dataset and generate tailored policies for each dataset from a large search space.

\begin{table}[tb]
    \centering
    \caption{Test error rates on CIFAR-10, CIFAR-100 and SVHN. For AutoAugment${}^{\ast}$, we train CNN models with the official AutoAugment\xspace policy for CIFAR-10 provided by the authors. \textbf{MADAO\xspace uses the same hyperparameters for all configurations and achieves superior performance to baselines, which demonstrates the effectiveness of our method.}}
    \vspace{-0.5\baselineskip}
    \label{tab:results_small}
    \begin{tabular}{llcccc}
       \toprule
       Dataset   &  Model            & Baseline  & AutoAugment${}^{\ast}$   & RandAugment\xspace      & MADAO\xspace (ours) \\
       \cmidrule(r){1-2} \cmidrule(rl){3-6}
       CIFAR-10  &  WideResNet 28-2  & 4.8       &    4.7              & 4.7          &  4.6          \\
                 &  WideResNet 40-2  & 4.5       &    4.3              & 4.3          &  4.2          \\
                 &  WideResNet 28-10 & 3.4       &    3.1              & 3.1          &  3.1          \\
                 &  ResNet-18        & 5.2       &    4.5              & 4.5          &  4.4          \\
       \cmidrule(r){1-2} \cmidrule(rl){3-6}
       CIFAR-100 &  WideResNet 28-2  & 24.8      &   24.0              & 24.0         &  23.9           \\
       \cmidrule(r){1-2} \cmidrule(rl){3-6}
       SVHN      &  WideResNet 28-2  & 2.9       &    N/A              & 2.6           & 2.5             \\
       \bottomrule

    \end{tabular}
\end{table}

\begin{table}[tb]
    \centering
    \caption{ImageNet test error rates (Top-1/Top-5) on supervised and self-supervised learning.}
    \vspace{-0.5\baselineskip}
    \label{tab:results_imagenet}
    \begin{tabular}{llccc}
       \toprule
       Task                 &  Model            & Standard & RandAugment\xspace       & MADAO\xspace (ours) \\
       \cmidrule(r){1-2} \cmidrule(rl){3-5}
       Supervised Learning       &  ResNet-50        & 23.7/ 6.9  &  22.5/ 6.4    & 23.0/ 6.6             \\
       \midrule
       Self-Supervised Learning  &  ResNet-50        & 42.1/19.4 &  39.6/17.6   & 39.6/17.7            \\
       \bottomrule
    \end{tabular}
\end{table}

\begin{table}[tb]
    \centering
    \caption{Test error rates on fine-grained datasets. MADAO\xspace outperforms RandAugment\xspace with a large margin. \textbf{Note that we use the same hyperparemeters as in the experiments in \Cref{tab:results_small}.}}
    \vspace{-0.5\baselineskip}
    \label{tab:results_finedgrained}
    \begin{tabular}{ccccccc}
       \toprule
       Dataset   &  Model            & \# Classes  & \# Training Set & Standard & RandAugment\xspace       & MADAO\xspace (ours) \\
       \cmidrule(r){1-4} \cmidrule(l){5-7}
       Flower    &  ResNet-18        & 102         &  2,040          & 10.7      &  \zero8.5   &  \zero7.5             \\
       Pet       &  ResNet-18        & \zero37     &  3,680          & 15.0      &  13.4       &  12.5                 \\
       Aircraft  &  ResNet-18        & \zero70     &  6,667          & 12.8      &  \zero9.7   &  \zero8.7            \\
       Car       &  ResNet-18        & 196         &  8,144          & 14.2      &  11.3       &  10.8                 \\
       \bottomrule
    \end{tabular}

\end{table}

\section{Analysis}

\subsection{How Policies Develop during Training}

In \Cref{fig:policy_development}, we present the development of policies during training on fine-grained datasets. As can be seen, each dataset has its specific operations that are selected, which could be thought as reflecting the characteristics of each dataset. Besides, the way of selection changes according to the learning phase. Note that the first 20 epochs are set to warm-up and the augmentation policy parameters are not updated. We show further observations in \Cref{ap:additional_experiments}.

\begin{figure}[tb]
    \centering
    \includegraphics[width=0.9\linewidth]{policy_development.pdf}
    \vspace{-1\baselineskip}
    \caption{The development of operation selection probabilities at the first augmentation stage during training on fine-grained datasets for 200 epochs. Note that we set the first 20 epochs as a warm-up period during which parameters are not updated.}
    \label{fig:policy_development}
\end{figure}

\begin{wrapfigure}[11]{r}{0.4\linewidth}
    \centering
    \vspace{-3.5\baselineskip}
    \includegraphics[width=\linewidth]{ablations.pdf}
    \vspace{-2\baselineskip}
    \caption{The relationship between the number of inner steps and the final test error rates on CIFAR-10 with WideResNet-28-2. We present the mean and standard deviation over three runs.}
    \label{fig:ablations}
\end{wrapfigure}

\subsection{How Inner Steps Affect the Performance}\label{sub:num_inner_optim}

\Cref{fig:ablations} presents test error rates with various numbers of inner update steps $s$ using WideResNet-28-2 on CIFAR-10. CNNs yield the best performance when $s=30$ and the results indicate that there exists a trade-off between ``exploration and exploitation'' of obtained policies: a small number of inner steps might not correctly evaluate the current policies, while running a large number of inner steps might fail to explore better strategies. Importantly, unrolled-based implementations would require to store $s$ model caches, which is infeasible for $s=30$ with modern CNNs. On the other hand, MADAO\xspace can efficiently handle a large $s$.

\section{Conclusion}

In this paper, we have proposed MADAO\xspace, a novel approach to optimize an image recognition model and its data augmentation policy simultaneously. To efficiently achieve this goal, we use the implicit gradient method with Neumann series approximation. As a result, the overhead of MADAO\xspace to the standard CNN training, w.r.t., time and memory, is marginal, which enables ImageNet-scale training. Empirically, we demonstrate on various tasks that MADAO\xspace achieves superior performance to prior works without restricting search space or using sub-optimal proxy tasks.

Data augmentation is known to boost the performance in various visual representation learning settings, such as semi-supervised learning \cite{Berthelot2019a,Suzuki2020}, domain generalization \cite{Volpi2018}, and self-supervised learning \cite{Chen2020}. We believe that our method can be introduced into these representation learning methods and efficiently enhance their performance, as we showcased with self-supervised learning in this paper.

\clearpage
\makeatletter
\if@preprint

\else
\section*{Potential Impacts of Our Work}

MADAO\xspace boosts performance of image recognition models with little overhead and minimal hyperparameter tuning. These advantages help to reduce energy consumption, which is a known issue of large-scale hyperparameter optimization, such as neural architecture search \cite{strubell2019}. MADAO\xspace showcases superior performance, especially on fine-grained datasets. This property enables further application of image recognition models to various domains, but it may also help in misuse of image recognition.

\fi
\makeatother

\printbibliography