EconBase
← Back to paper

On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization

The exact contents of citations.db main_text.text for this paper — one flattened LaTeX string, title through conclusion, appendix excluded, unmodified except for removing email addresses. This is what our citation measures are computed over.

42,387 characters

On Efficient Estimation of Distributional Treatment Effects under Covariate-Adaptive Randomization



\twocolumn[
\icmltitle{On Efficient Estimation of Distributional Treatment Effects\\ under Covariate-Adaptive Randomization}



\icmlsetsymbol{equal}{*}

\begin{icmlauthorlist}
\icmlauthor{Undral Byambadalai}{comp}
\icmlauthor{Tomu Hirata}{sch}
\icmlauthor{Tatsushi Oka}{yyy}
\icmlauthor{Shota Yasui}{comp}

\end{icmlauthorlist}

\icmlaffiliation{yyy}{Department of Economics, Keio University, Tokyo, Japan}
\icmlaffiliation{comp}{CyberAgent, Inc., Tokyo, Japan}
\icmlaffiliation{sch}{Databricks Japan, Inc., Tokyo, Japan}

\icmlcorrespondingauthor{Undral Byambadalai}{undral\textunderscore [email removed]}

\icmlkeywords{distributional treatment effects, field experiments, covariate-adaptive randomization, regression adjustment, machine learning}

\vskip 0.3in
]



\printAffiliationsAndNotice{}

\begin{abstract}
This paper focuses on the estimation of distributional treatment effects in randomized experiments that use covariate-adaptive randomization (CAR). These include designs such as Efron's biased-coin design and stratified block randomization, where participants are first grouped into strata based on baseline covariates and assigned treatments within each stratum to ensure balance across groups. In practice, datasets often contain additional covariates beyond the strata indicators. We propose a flexible distribution regression framework that leverages off-the-shelf machine learning methods to incorporate these additional covariates, enhancing the precision of distributional treatment effect estimates. We establish the asymptotic distribution of the proposed estimator and introduce a valid inference procedure. Furthermore, we derive the semiparametric efficiency bound for distributional treatment effects under CAR and demonstrate that our regression-adjusted estimator attains this bound. Simulation studies and empirical analyses of microcredit programs highlight the practical advantages of our method.
\end{abstract}

\section{Introduction}
Randomized experiments have been fundamental in uncovering the impact of interventions and shaping policy decisions since the seminal work by \citet{fisher1935design}. The use of randomized experiments to estimate causal effects has been extensively adopted across diverse scientific fields \cite{rubin1974estimating, heckman1997making, imai2005get,imbens2015causal}
and has also become a widely accepted practice in the technology industry \cite{tang2010overlapping, bakshy2014designing, kohavi2020trustworthy}.

In this paper, we examine the estimation of distributional treatment effects in randomized experiments that employ covariate-adaptive randomization (CAR). Under CAR, individuals are first partitioned into \emph{strata} based on similar covariates, and treatments are then assigned within each stratum to ensure balance across groups. The CAR framework encompasses various randomization schemes, from stratified block randomization to Efron's biased coin design \cite{imbens2015causal}, with simple random sampling as a special case.

We introduce a regression adjustment method for estimating distributional treatment effects under CAR when auxiliary data beyond stratum indicators are available. Incorporating this data enhances estimation precision. While experimental analyses often focus on the average treatment effect (ATE), relying solely on the ATE can overlook important insights. Examining distributional treatment effects provides a more comprehensive understanding of the treatment impact by capturing changes in the entire outcome distribution.

Our approach employs a distribution regression framework, leveraging the Neyman-orthogonal moment condition \cite{chernozhukov2018debiased} to ensure first-order insensitivity to nuisance parameter estimation. These nuisance parameters—conditional outcome distributions given pre-treatment covariates—are estimated using machine learning techniques such as random forests, neural networks, and gradient boosting, enabling flexibility in handling complex and high-dimensional data. Incorporating cross-fitting further strengthens robustness against estimation errors.

Randomization schemes under CAR are extensively used across various disciplines. In clinical trials, stratified block randomization ensures balanced treatment allocation across key covariates such as age, gender, and disease severity \cite{rosenberger2015randomization}. In social experiments, researchers often stratify by geographical regions or socioeconomic characteristics \cite{duflo2007using, bruhn2009pursuit}, while in the technology sector, covariate-adaptive designs enhance precision in A/B testing \cite{xie2016improving}.

Our paper makes several key contributions. First, we extend the applicability of regression adjustment under CAR to estimate distributional treatment effects. While regression adjustment is widely used to reduce variance in ATE estimation under simple random sampling \cite{freedman2008regression, lin2013agnostic} and has recently been studied in the context of CAR \cite{rafi2023efficient, cytrynbaum2024covariate, wang2023model}, our work advances this framework to accommodate distributional treatment effects. Unlike \citet{jiang2023regression}, who focus on regression adjustment for quantile treatment effects and assume continuous outcomes,
our method is applicable to both discrete and mixed discrete-continuous outcomes.

Second, we establish the limit distribution of our estimator within the asymptotic framework for CAR, extending beyond the standard i.i.d. treatment assignment structure in causal inference. Third, we derive the semiparametric efficiency bound for distributional treatment effects under CAR and demonstrate that our estimator attains this bound. Finally, through simulations and an empirical analysis of microcredit programs, we illustrate the effectiveness of our approach in practical settings.

The rest of the paper is organized as follows. Section \ref{sec:literature} reviews the literature, and Section \ref{sec:setup} provides the setup. Section \ref{sec:dte} introduces distributional treatment effect parameters, their identification, and estimation. Section \ref{sec:asymptotics} presents asymptotic results. Section \ref{sec:experiments} discusses findings from simulated and real data. Section \ref{sec:conclusion} concludes. Appendix contains notations, proofs, and additional experimental results.


\section{Related Literature}\label{sec:literature}
\begin{figure*}[!h]
\begin{tikzpicture}
    \tikzstyle{block} = [rectangle, draw,
    text centered, minimum width=2cm, minimum height=1.5cm]
    \tikzstyle{label} = [text centered, font=\bfseries]
    \tikzstyle{stratum} = [rectangle, draw, text centered, minimum width=5cm, minimum height=1.5cm, fill=gray!20]
    \tikzstyle{wide_block} = [rectangle, draw, text centered, minimum width=5cm, minimum height=1.5cm]

            
   \node[above] at (6.5,8) {\textbf{Illustration: Comparison of Randomization Methods}};

    \node[label] at (0, 7.5) {Sample with Strata};
    \node[stratum] at (0, 6) {};
    \node at (0, 6.3) {Stratum 1};
    \node at (0, 5.7) {\textbf{30} subjects};

    \node[stratum] at (0, 4.3) {};
    \node at (0, 4.6) {Stratum 2};
    \node at (0, 4.0) {\textbf{70} subjects};

    \node[label] at (6, 7.5) {Simple Random Sampling};

    \path[fill=red!20] (6-2.5, 5.25) rectangle (6-0.83, 6.75);
    \path[fill=blue!20] (6-0.83, 5.25) rectangle (6+2.5, 6.75);
    \node[wide_block] at (6, 6) {};
    \node at (6-1.67, 6.3) {Treatment};
    \node at (6-1.67, 5.7) {\textbf{10}};
    \node at (6+0.8, 6.3) {Control};
    \node at (6+0.8, 5.7) {\textbf{20}};
    \draw (6-0.83, 5.25) -- (6-0.83, 6.75);

    \path[fill=red!20] (6-2.5, 3.55) rectangle (6+0.71, 5.05);
    \path[fill=blue!20] (6+0.71, 3.55) rectangle (6+2.5, 5.05);
    \node[wide_block] at (6, 4.3) {};
    \node at (6-1.0, 4.6) {Treatment};
    \node at (6-1.0, 4.0) {\textbf{40}};
    \node at (6+1.5, 4.6) {Control};
    \node at (6+1.5, 4.0) {\textbf{30}};
    \draw (6+0.71, 3.55) -- (6+0.71, 5.05);

    \node[label] at (12, 7.5) {Stratified Block Randomization};

    \path[fill=red!20] (12-2.5, 5.25) rectangle (12, 6.75);
    \path[fill=blue!20] (12, 5.25) rectangle (12+2.5, 6.75);
    \node[wide_block] at (12, 6) {};
    \node at (12-1.3, 6.3) {Treatment};
    \node at (12-1.3, 5.7) {\textbf{15}};
    \node at (12+1.3, 6.3) {Control};
    \node at (12+1.3, 5.7) {\textbf{15}};
    \draw (12, 5.25) -- (12, 6.75);

    \path[fill=red!20] (12-2.5, 3.55) rectangle (12, 5.05);
    \path[fill=blue!20] (12, 3.55) rectangle (12+2.5, 5.05);
    \node[wide_block] at (12, 4.3) {};
    \node at (12-1.3, 4.6) {Treatment};
    \node at (12-1.3, 4.0) {\textbf{35}};
    \node at (12+1.3, 4.6) {Control};
    \node at (12+1.3, 4.0) {\textbf{35}};
    \draw (12, 3.55) -- (12, 5.05);
  \end{tikzpicture}
\caption{An illustration of treatment and control assignments under simple random sampling (SRS) and stratified block randomization (SBR). While both allocate 50 subjects per group, SRS may result in imbalanced group compositions, whereas SBR preserves the proportional representation of strata in each group, matching the overall sample distribution.} \label{fig:complete-vs-stratified}
\end{figure*}

\paragraph{Distributional Treatment Effects} Distributional and quantile treatment effects have long been recognized as important parameters to estimate beyond the mean effects. The quantile treatment effect was first introduced by \citet{doksum1974empirical} and \citet{lehmann1975nonparametrics}. Subsequently, estimation and inference methods for distributional and quantile treatment effects have been developed and applied in econometrics, statistics and machine learning community, including \citet{heckman1997making, imbens1997estimating, abadie2002bootstrap, abadie2002instrumental, chernozhukov2005iv, koenker2005quantile, bitler2006mean, athey2006identification, firpo2007efficient, chernozhukov2013inference, koenker2017handbook, callaway2018quantile, callaway2019quantile, chernozhukov2019generic, ge2020conditional, zhou2022estimating, park2021conditional, kallus2023robust,
oka2023heterogeneous, naf2024causal, xu2025quantile}, among others. Most of this work explore the conditional distributional and quantile treatment effects. \citet{oka2024regression} and \citet{byambadalai24a} consider the estimation of unconditional distributional treatment effects but under simple random sampling. \citet{kallus2024localized} address problems in which nuisance parameters depend on the target parameter itself, as seen in cases like the quantile treatment effect (QTE) and local QTE. In contrast, our estimation of the distributional treatment effect involves nuisance parameters that correspond to conditional means, which can be effectively estimated using machine learning algorithms.

\paragraph{Conditional Average Treatment Effects}
An alternative method for examining heterogeneity in treatment effects is to condition on observed variables and estimate the Conditional Average Treatment Effect (CATE) \cite{imai2013estimating, athey2016recursive, johansson2016learning, shalit2017estimating, alaa2017bayesian, wager2018estimation, chernozhukov2018debiased, kunzel2019metalearners, shi2019adapting, nie2021quasi, guo2023estimating, sverdrup2023proximal, van2023causal}.
CATE quantifies the ATE within subgroups defined by observed attributes, such as gender, age, or prior platform use, thereby capturing the heterogeneity that can be explained by observable information. In contrast, our approach is designed to measure unobserved heterogeneity and can be extended to estimate distributional parameters conditional on observed data.

\vspace{-0.25cm}
\paragraph{Regression adjustment under covariate-adaptive randomization}
The literature on utilizing pre-treatment covariates to reduce variance in estimating the ATE under simple random sampling is extensive, beginning with \citet{fisher1932statistical} and followed by contributions from \citet{cochran1977sampling, yang2001efficiency, rosenbaum2002covariance, freedman2008regression, freedman2008regression2, tsiatis2008covariate, rosenblum2010simple, lin2013agnostic, berk2013covariance, ding2019decomposing}, among others.

In the context of covariate-adaptive randomization, some recent studies have explored regression adjustment for estimating the ATE. Recent work by \citet{cytrynbaum2024covariate} derives an asymptotically optimal linear covariate adjustment tailored to a given stratification. Similarly, \citet{rafi2023efficient} investigates regression adjustment and establishes the semiparametric efficiency bound for estimating the ATE under covariate-adaptive randomization. \citet{bai2024covariate} examines covariate adjustment within a ``matched pairs'' design, where each stratum consists of two observations, with one randomly assigned to treatment. In biostatistics, \citet{bannick2023general} and \citet{tu2023unified} analyze general approaches to covariate adjustment under covariate-adaptive randomization, while \citet{wang2023model} considers parameters defined by estimating equations. Notably, most of these studies concentrate on ATE estimation, with the exception of \citet{jiang2023regression}, who investigates the estimation of quantile treatment effects in the same setting.

\vspace{-0.25cm}
\paragraph{Semiparametric Estimation}
Our work builds on the semiparametric estimation literature, which addresses the challenge of estimating low-dimensional parameters in the presence of high-dimensional nuisance parameters. This literature includes fundamental contributions from \citet{robinson1988root, bickel1993efficient, newey1994asymptotic, robins1995semiparametric}, and more recent developments by \citet{chernozhukov2018debiased, ichimura2022influence}. Our setup can be characterized as a semiparametric problem with Neyman-orthogonal moment conditions,
as outlined in \citet{neyman1959optimal, chernozhukov2022locally}.

\section{Setup and Notation} \label{sec:setup}
We consider randomized controlled trials (RCT) with covariate-adaptive randomization when there are multiple treatment arms. In covariate-adaptive randomization, individuals are first grouped into \emph{strata} based on similar values of their baseline covariates. Within each stratum, treatments are assigned to achieve ``balance'' across groups, typically using complete randomization within each stratum. This approach ensures that the treatment allocation accounts for covariate distributions within strata, thereby improving comparability between treatment groups.

Figure \ref{fig:complete-vs-stratified} provides a simple illustration. Consider an experiment involving 100 subjects, where 50\% are assigned to a treatment group and 50\% to a control group. Suppose the subjects are divided into two strata based on baseline covariates (e.g., Stratum 1 represents younger individuals, and Stratum 2 represents older individuals). Stratum 1 comprises 30 subjects, while Stratum 2 comprises 70. Under simple random sampling (SRS), subjects are randomly assigned to treatment or control groups without regard to strata. In contrast, stratified block randomization (SBR) independently assigns subjects within each stratum, ensuring proportional representation. While both methods allocate exactly 50 subjects to each group, SRS does not guarantee that the composition of treatment groups reflects the overall sample distribution. For example, Figure \ref{fig:complete-vs-stratified} depicts one possible outcome, where Stratum 2 (older individuals) is over-represented in the treatment group under SRS. By contrast, SBR achieves balance within each stratum and maintains the same composition in treatment and control groups as in the full sample.

Let $Y_i \in \mathcal{Y} \subset \mathbb{R}$ denote the observed outcome of interest for the $i$th unit and $W_i \in \mathcal{W} := \{1,\dots, |\mathcal W|\}$ denote the index of the treatment received by the $i$th unit for each unit $i \in [n]:= \{1,\ldots,n\}$, where $n$ denotes the sample size. Also, let $S_i\in \mathcal S:=\{1,..., S\}$ denote the stratum variable and $X_i\in\mathcal X \subset \mathbb R^{d_x}$ denote the extra covariates besides $S_i$. We allow $X_i$ and $S_i$ to be dependent. To rule out empty stratum, we assume the probability of individuals being assigned to each stratum is positive, i.e., $p(s) := \mathbb{P}(S_i=s) >0$ for every $s\in\mathcal S$. We adapt the potential outcome framework \cite{rubin1974estimating, imbens2015causal}, and let $Y_i(w)$ denote the potential outcome under treatment $w\in \mathcal W$ for the $i$th unit. The observed outcome and potential outcomes are related to treatment assignment by the relationship $Y_i = Y_i(W_i)$.

In order to describe the treatment assignment mechanism, we let $\pi_w(s):=\mathbb{P}(W_i=w|S_i=s)\in(0,1)$ be the target assignment probability for treatment $w$ in stratum $s$, which may vary across strata. We define sub-sample  as $n_w(s) := \sum_{i=1}^{n} \mathbbm{1}\{W_i=w, S_i=s\}$ and $n(s):= \sum_{i=1}^{n} \mathbbm{1}\{S_i=s\}$, where $\mathbbm{1}\{\cdot\}$ denotes the indicator function.
The empirical analogues of
$p(s)$ and $\pi_{w}(s)$
are given by
$\widehat p(s):=
n(s)/n$ and $\widehat\pi_{w}(s) := n_w(s)/n(s)$, respectively.
In line with \citet{bugni2019inference},
we impose the following assumptions on the treatment assignment mechanism.

\begin{assumption} \label{ass:dgp_treatment_assignment}
We have

(i) $\{Y_i(1), \dots, Y_i(|\mathcal{W}|), S_i, X_i\}_{i=1}^{n}$ are i.i.d.

(ii) $\{Y_i(1), \dots, Y_i(|\mathcal{W}|), X_i\}_{i=1}^{n} \rotatebox[origin=c]{90}{$\models$} \{W_i\}_{i=1}^{n} | \{S_i\}_{i=1}^{n}$.

(iii) $\widehat\pi_w(s) = \pi_w(s) + o_p(1)$ for every $(w, s) \in \mathcal{W} \times \mathcal S$.
\end{assumption}
Assumption \ref{ass:dgp_treatment_assignment} (i) allows for cross-sectional dependence among treatment statuses $\{W_i\}_{i=1}^{n}$, thereby accomodating many covariate-adaptive randomization schemes. Assumption \ref{ass:dgp_treatment_assignment} (ii) states that the assignment is independent of potential outcomes and pre-treatment covariates conditional on strata. Assumption \ref{ass:dgp_treatment_assignment} (iii) states the assignment probabilities converge to the target assignment probabilities as sample size increases.

Common randomization schemes satisfying Assumption \ref{ass:dgp_treatment_assignment} include simple random sampling, stratified block randomization, biased-coin design \cite{efron1971forcing}, and adaptive biased-coin design \cite{wei1978adaptive}. In Section \ref{sec:microfinance}, we reanalyze a field experiment on microcredit programs, where stratification occurs at the provincial level, and treatments are randomly assigned within each province based on target probabilities.


\section{Distributional Treatment Effects}\label{sec:dte}
The parameters of interest are based on the distribution functions  of potential outcomes, denoted by
\begin{align*}
F_{Y(w)}(y):=\mathbb{P}\big(Y(w) \leq y\big),
\end{align*}
for $y\in\mathcal{Y}$ and $w\in\mathcal{W}$.

First, we define the \emph{distributional treatment effect} (DTE) between treatments $w, w'\in \mathcal{W}$ as
\begin{align*}
\Delta^{DTE}_{w, w'}(y) := F_{Y(w)}(y) - F_{Y(w')}(y),
\end{align*}
for $y\in\mathcal{Y}$. The DTE measures the difference between the distribution functions of the potential outcomes.

We also define the \emph{probability treatment effect} (PTE) between treatments $w, w'\in \mathcal{W}$ as
\begin{align*}
\Delta^{PTE}_{w, w'}(y_{j}) := f_{Y(w)}(y_{j}) - f_{Y(w')}(y_{j}),
\end{align*}
for each $j = 1, \dots, J$,
where,
given a set of points
$\mathcal{Y}_J:=\{y_1, \cdots, y_J\} \subset \mathcal{Y}$ with $y_0=-\infty$,
the probability mass function $f_{Y(w)}(\cdot)$ is defined as
\begin{align*}
f_{Y(w)}(y_j):&= \mathbb{P}\big (y_{j-1} < Y(w) \leq y_j \big) \\
& = F_{Y(w)}(y_j) - F_{Y(w)}(y_{j-1}).
\end{align*}
Each probability $f_{Y(w)}(y_j)$, which we refer to as the \emph{bin probability}, can be obtained as a difference between the values of the distribution function at $y_j$ and $y_{j-1}$. The PTE quantifies the difference in probabilities across bins, effectively capturing differences in “histograms” of the potential outcomes.

\subsection{Identification}
Although the potential outcomes $\{Y(w): w \in \mathcal{W}\}$ are unobserved variables, the conditional distribution functions of these outcomes, $F_{Y(w)}(\cdot|S=s)$, can be identified. This is because, under Assumption \ref{ass:dgp_treatment_assignment}, $F_{Y(w)}(\cdot|S=s)$ coincides with $F_{Y}(\cdot|W=w,S=s)$, the conditional distribution function of the observed outcome given treatment $w$ within stratum $s$. By the law of total probability, the distribution function $F_{Y(w)}(y)$ is then expressed as,
for any $y \in \mathcal{Y}$,
\begin{align}
F_{Y(w)}(y) = \sum_{s=1}^{S} p(s) F_{Y(w)}(y|S=s).
\end{align}
Hence, $F_{Y(w)}(y)$ is identifiable under Assumption \ref{ass:dgp_treatment_assignment}.


\subsection{Empirical estimator}
Under the RCT setting, the distribution function can be calculated without conditioning on pre-treatment covariates, unlike observational studies.
Under CAR, we define an empirical estimator of the distribution function $F_{Y(w)}(y)$ for $y\in\mathcal{Y}$ and treatment $w\in\mathcal{W}$ as:
\begin{align} \label{eq:empirical-cdf}
    \widehat F^{emp}_{Y(w)}(y)
     := \frac{1}{n}\sum_{i=1}^{n} \frac{ \mathbbm{1}\{W_i=w\}\cdot \mathbbm{1}\{Y_i \leq y\}}{\widehat \pi_w(S_i)}.
\end{align}
This estimator aggregates the empirical distribution functions across strata, and takes the form of an inverse-propensity weighting (IPW) estimator \citet{rosenbaum1984reducing}. The empirical estimator for the DTE is then formed as
\begin{align*}
    \widehat{\Delta}^{DTE, emp}_{w,w'}(y) := \widehat F_{Y(w)}^{emp} (y) - \widehat F_{Y(w')}^{emp} (y).
\end{align*}
While the estimator is unbiased and consistent, its efficiency can be enhanced by leveraging pre-treatment covariates.

\subsection{Regression-adjusted estimator} To incorporate pre-treatment covariates $X$, we adopt the distribution regression framework, treating the conditional distribution function $\mu_w(y, s,x):=F_{Y(w)}(y|S=s, X=x)$ as the mean regression for a binary outcome
$\mathbbm{1}\{Y(w) \leq y\} $. Specifically, for each $y \in \mathcal{Y}$ and $w \in \mathcal{W}$, we can write
\begin{align*}
   \mu_w(y, S,X) = \mathbb{E}[\mathbbm{1}\{Y(w)\leq y\}|S, X].
\end{align*}
The conditional mean function can be estimated at each location $y \in \mathcal{Y}$ using supervised learning algorithms, such as LASSO, random forests, boosted trees, or deep neural networks. Let $\widehat \mu_w(\cdot) $ be an estimator for $\mu_w(\cdot)$.

The regression-adjusted estimator of $F_{Y(w)}(y)$ for each $w\in\mathcal{W}$ and $y\in\mathcal{Y}$ is
then defined as
\begin{align} \label{eq:reg-adj-estimator}
 \widehat F_{Y(w)}^{adj} (y) = \frac{1}{n} \sum_{i=1}^{n}  \Psi_i(y),
\end{align}
where
\begin{align*}
\Psi_i(y) := & \frac{ \mathbbm{1}\{W_i=w\} \cdot \big(\mathbbm{1}\{Y_i \leq y\} - \widehat \mu_w(y, S_i, X_i) \big) }{\widehat \pi_w(S_i)} \\
 & + \widehat \mu_w(y, S_i,X_i).
\end{align*}

The regression-adjusted estimator for DTE can then be formed as
\begin{align}\label{eq:dte-adj}
    \widehat{\Delta}^{DTE, adj}_{w,w'}(y) := \widehat F_{Y(w)}^{adj} (y) - \widehat F_{Y(w')}^{adj} (y).
\end{align}

The estimator takes the form of the well-known augmented inverse-propensity weighting (AIPW) estimator, relying on a doubly-robust moment condition \cite{robins1994estimation, robins1995semiparametric}. Specifically, the moment condition exhibits the Neyman orthogonality property \cite{chernozhukov2018debiased, chernozhukov2022locally}, making our estimator first-order insensitive to estimation errors in the nuisance parameters $\mu_w(\cdot)$. To further improve robustness, we apply cross-fitting as outlined in \citet{chernozhukov2018debiased}. The estimation procedure is outlined in Algorithm \ref{alg:reg-adj-estimator}.

The empirical and adjusted estimators for the PTE can be defined in a similar fashion. See Appendix \ref{app:other-parameters} for the details.

\begin{algorithm}[!h]
\caption{ML regression-adjusted DTE estimator with cross-fitting}
\label{alg:reg-adj-estimator}
\begin{algorithmic}
\STATE \textbf{Input:} Data $\{(Y_i, W_i, X_i, S_i)\}_{i=1}^n$ split randomly into $L$ roughly equal-sized folds ($L > 1$);  $\mathcal{M}$ a supervised learning algorithm
\FOR{level $y \in \mathcal{Y}$}
\FOR{(treatment $w {\in} \mathcal{W}$, stratum $s {\in} \mathcal{S}$, fold $\ell {=} 1$ to$L$)}
\STATE Train $\mathcal{M}$ on data excluding fold $\ell$, using observations in treatment group $w$ within stratum $s$.
\STATE Use $\mathcal{M}$ to obtain predictions $\hat{\mu}_w(y, S_i, X_i)$ for all observations in stratum s for fold $\ell$.
\ENDFOR
\STATE Compute $\widehat{F}^{adj}_{Y(w)}(y)$ according to Eq.~\eqref{eq:reg-adj-estimator}.
\STATE Compute $\widehat\Delta^{DTE, adj}_{w,w'}(y)$ using  Eq.~\eqref{eq:dte-adj}
for $w, w' \in \mathcal{W}$.
\ENDFOR
\STATE \textbf{Output:}
DTE estimator $\{\widehat\Delta^{DTE, adj}_{w,w'}(y) \}_{y\in\mathcal Y}$.
\end{algorithmic}
\end{algorithm}


\section{Asymptotic Distribution of the DTE estimator}\label{sec:asymptotics}
In this section, we derive the asymptotic distribution of the proposed estimator, which enables statistical inference and the construction of confidence intervals. Additionally, we establish the semiparametric efficiency bound for the DTE and demonstrate that our estimator achieves this bound under the specified assumptions. We begin by introducing some additional notation to formalize our results. Let  $\|\cdot\|_{P,q}$ denote the $L_q(P)$ norm, and $L^{\infty}(\mathcal Y)$ be the space of uniformly bounded functions mapping an arbitrary index set $\mathcal{Y}$ to the real line.

\vspace{0.1cm}
\begin{assumption} \label{ass:asymptotic-dist}
(i) For $w\in\mathcal{W}$, $\delta_w(y,s,X_i) := \widehat{\mu}_w(y,s,X_i) - \mu_w(y,s,X_i)$, we have
\begin{align*}
\sup_{y \in \mathcal{Y},s\in \mathcal{S}}&\biggl|\frac{\sum_{i\in I_w(s)}\delta_w(y,s,X_i)}{n_w(s)} - \frac{\sum_{i \in I_{w'}(s)}\delta_{w}(y,s,X_i)}{n_{w'}(s)}\biggr| \\&= o_p(n^{-1/2}),
\end{align*}
where $I_w(s) := \{i\in [n]: W_i = w, S_i=s\}$.

(ii) For $w\in\mathcal{W}$, let $\mathcal{F}_w = \{\mu_w(y,s,x): y\in \mathcal{Y} \}$ with an envelope $F_{w}(s,x)$. Then, $\max_{s \in \mathcal{S}}\mathbb{E}[|F_{w}(S_i,X_i)|^q|S_i=s]<\infty$ for $q > 2$ and there exist fixed constants $(\alpha,v)>0$ such that
\begin{align*}
\sup_Q N\left(\varepsilon||F_w||_{Q,2}, \mathcal{F}_w, L_2(Q) \right) \leq \left(\frac{\alpha}{\varepsilon}\right)^{v}, \quad \forall \varepsilon \in (0,1],
\end{align*}
where $N(\cdot)$ denotes the covering number and the supremum is taken over all finitely discrete probability measures $Q$.
\end{assumption}
\vspace{0.1cm}

Assumption \ref{ass:asymptotic-dist}(i) provides a high-level condition on the estimation of $\widehat{\mu}_w(y,s,X_i)$. Assumption \ref{ass:asymptotic-dist}(ii) imposes mild regularity conditions on $\mu_w(y,s,X_i)$. Specifically, Assumption \ref{ass:asymptotic-dist}(ii) holds automatically when $\mathcal{Y}$ is a finite set.

We now establish the weak convergence of our proposed estimator in the following theorem, which serves as the theoretical foundation for statistical inference. To that end, we define the following terms. Letting $\mu_w(y, s) := \mathbb{E}[\mu_w(y, S_i, X_i)|S_i=s]$, we define
\begin{align*}
 \zeta_{i}(y) &: = \mu_w(y,S_i) - \mu_{w'}(y,S_i),  \\
 \phi_{i}(y, w)  &: =
\frac{\mathbbm{1}\{Y_i(w)\leq y\} - \mu_w(y,s)}{\pi_w(s)} \\
& \ \  + \left(1-\frac{1}{\pi_w(s)}\right) \left(\mu_w(y,s,X_i)-\mu_w(y,s)\right) \\
& \ \   - \left(\mu_{w'}(y,s,X_i) -\mu_{w'}(y,s)\right).
\end{align*}

\vspace{0.1cm}
\begin{theorem}[Asymptotic Distribution]\label{thm:asymptotics}
Suppose Assumption \ref{ass:dgp_treatment_assignment} and \ref{ass:asymptotic-dist} hold. Then, for every $w, w'\in \mathcal{W}$, in $L^{\infty}(\mathcal Y)$, uniformly over $y\in\mathcal Y$, the regression-adjusted estimator defined in Algorithm \ref{alg:reg-adj-estimator} satisfies
\begin{align*}
\sqrt{n} \big (\widehat \Delta^{DTE, adj}_{w, w'}(y) - \Delta^{DTE}_{w, w'}(y) \big) \rightsquigarrow \mathcal{G}(y),
\end{align*}
where $\mathcal{G}(y)$ is a Gaussian process with covariance kernel $\Omega(y, y')$, which is given by
\begin{align*}
    & \Omega(y, y') :=  \Omega_{1}(y, y', w) +\Omega_{1}(y, y', w') + \Omega_{2}(y, y'),
\end{align*}
with
$
\Omega_{1}(y, y', w)  :=
\mathbb{E}[\pi_w(S_i)\phi_i(y, w)\phi_{i}(y', w)]
$ and
$\Omega_{2}(y, y')  :=
\mathbb{E}[\xi_{i}(y)\xi_{i}(y')].$
\end{theorem}
\vspace{0.1cm}

We next derive the semiparametric efficiency bound and show our estimator achieves this bound in the following theorem. This implies that the asymptotic variance of any regular, consistent, and asymptotically normal estimator of DTE cannot be smaller than this variance.

\vspace{0.1cm}
\begin{theorem}[Semiparametric Efficiency Bound]\label{thm:semi-eff-bound} \

\vspace{-\baselineskip}
\begin{itemize}[noitemsep]
    \item[(a)]
    Under Assumption \ref{ass:dgp_treatment_assignment},
    the semiparametric efficiency bound for $\Delta^{DTE}_{w, w'}(y)$ for a given $y\in\mathcal{Y}$ is $\Omega(y)$, which is defined by
    \begin{align*}
    & \Omega(y) := \Omega_{1}(y, y, w) +\Omega_{1}(y, y, w') + \Omega_{2}(y, y),
    \end{align*}
    where $\Omega_{1}(\cdot)$ and $\Omega_{2}(\cdot)$ are defined in Theorem \ref{thm:asymptotics}.
    \item[(b)]
    Under Assumptions \ref{ass:dgp_treatment_assignment} and \ref{ass:asymptotic-dist}, the regression-adjusted estimator $\widehat\Delta^{DTE, adj}_{w, w'}(y)$
    attains the semiparametric efficiency bound.
\end{itemize}
\vspace{-\baselineskip}
\end{theorem}
\vspace{0.1cm}

As a corollary to this theorem, it follows that the variance of the regression-adjusted estimator with known nuisance functions is smaller than that of the empirical estimator, since the latter can be regarded as a special case in which the adjustment terms $\widehat\mu_w(\cdot)$ are set to zero.

\vspace{0.1cm}
\begin{corollary}[Variance reduction]\label{cor:variance-reduction}
Under Assumption \ref{ass:dgp_treatment_assignment}, for each $y\in\mathcal Y$,
\begin{align*}
    Var\big (\tilde\Delta^{DTE, adj}_{w, w'}(y)\big ) \leq Var\big (\widehat\Delta^{DTE,emp}_{w, w'}(y)\big ),
\end{align*}
where $\tilde\Delta^{DTE, adj}_{w, w'}(y)$ is the regression-adjusted DTE estimator that incorporates known adjustment terms.
\end{corollary}

Theorem \ref{thm:semi-eff-bound} and Corollary \ref{cor:variance-reduction} indicate that, asymptotically, regression adjustment enhances the precision of the DTE estimates compared to the unadjusted empirical estimator. In the following section, we evaluate the finite sample performance of our estimators using both simulated and real datasets.

\section{Experiments} \label{sec:experiments}
\subsection{Simulation Study} \label{sec:simulation}
\begin{figure*}
    \centering
    \includegraphics[width=\linewidth]{metrics_simulation_continuous.png}
    \caption{RMSE, average length and coverage probability of 95\% confidence intervals (CI) on simulated data ($n=1{,}000$). Linear adjustment uses linear regression, and machine learning (ML) adjustment uses gradient boosting, both with 2-fold cross-fitting. Number of simulations is 1{,}000.}
    \label{fig:metrics_simulation_continuous}
\end{figure*}

\begin{figure*}[!htbp]
\centering
\includegraphics[width=0.4\linewidth]{entrevenueDTE_simple.boot.png}
\includegraphics[width=0.4\linewidth]{entrevenueDTE_adj.boot.png}
\includegraphics[width=0.4\linewidth]{entrevenuePTE_simple.boot.png}
\includegraphics[width=0.4\linewidth]{entrevenuePTE_adj.boot.png}
\caption{The Impacts of Microfinance: Distributional Treatment Effect (DTE) and Probability Treatment Effect (PTE) of joint liability lending on enterprise revenue (in thousand Mongolian Tugriks). The left panels depict empirical estimates, while the right panels present regression-adjusted estimates obtained using gradient boosting with 10-fold cross-fitting. Shaded regions and error bars represent 95\% confidence intervals. Sample size: $n=611$.} \label{fig:microfinance}
\end{figure*}

In this section, we examine the finite sample performance of the estimators through a simulation study. The design consists of four strata ($S = 4$) constructed by partitioning the support of $Z_i \sim U(0,1)$ into $S$ equal-length intervals, where $S_i$ indicates the interval containing $Z_i$. For each unit $i$, we draw a 20-dimensional covariate vector $X_i = (X_{1,i}, \dots, X_{20,i})^\top$ from a multivariate normal distribution $\mathcal N(0, I_{20\times20})$. The treatment indicator $W_i$ follows a Bernoulli distribution with probability 0.5 within each stratum, maintaining a constant target proportion of treated units ($W_i = 1$) across strata with $\pi_w(s) = 0.5$ for all $s \in \mathcal{S}$.
We generate the outcome variable $Y_i$ as follows:
\begin{align*}
Y_i= b(X_i) + c(X_i)W_i +\gamma Z_i +u_i,
\end{align*}
where $b(X_i)$ is the scaled \citet{friedman1991multivariate} function  given by $b(X_i)=\sin(\pi X_{i1} X_{i2}) +2(X_{i3} -0.5)^2+X_{i4}+0.5X_{i5}$, the treatment effect function is $c(X_i)=0.1(X_{i1}+\log(1+\exp(X_{i2})))$ and $\gamma=0.1$.
The error term is $u_i\sim \mathcal N(0,1)$. This data generating process introduces a complex, highly nonlinear relationship between covariates and the outcome, while including many covariates that do not affect the outcome.
The heterogeneous treatment effects are  correlated with specific covariates. This setup is  a modified version of the setups in \citet{nie2021quasi} and  \citet{guo2021machine}.

We draw a sample of size 5,000 from the data-generating process and estimate the DTE at quantiles $\{0.1, \dots, 0.9\}$ using empirical and regression-adjusted estimators. To approximate the ground truth, a separate sample of size $10^6$ is drawn, and the DTE is computed at the same locations. Linear and machine learning (ML) regression adjustments are implemented using linear regression and gradient boosting, with 2-fold cross-fitting. Linear regression is included as a baseline for comparison, as it has traditionally been widely used for regression adjustment in estimating average treatment effects.

Figure \ref{fig:metrics_simulation_continuous} presents the results from 1,000 simulation runs with a sample size of $n=1{,}000$. Both the root mean squared error (RMSE) and the length of the confidence intervals are smaller for the linear adjustment method, and even further reduced with ML adjustment, compared to the empirical estimator. The confidence intervals are calculated using sample estimates of the asymptotic variance. The empirical estimator achieves a 95\% confidence interval coverage close to the nominal level of 0.95. In contrast, regression-adjusted estimators show slight over-coverage, with coverage rates ranging from approximately 0.96 to 0.97 for linear adjustment and 0.97 to 0.99 for ML adjustment.
Additional results on RMSE reductions across quantiles for sample sizes
$n \in \{1{,}000, 5{,}000, 10{,}000\}$ are provided in Appendix Section \ref{app:simulation}.

\subsection{Real Data: The Impacts of Microfinance} \label{sec:microfinance}
We reexamine the randomized field experiment conducted by \citet{attanasio2015impacts} in 2008 to evaluate the impact of a joint-liability microcredit program targeting women. The study took place in northern Mongolia, encompassing 40 villages across five provinces. Randomization occurred at the village level, with three treatment groups: joint-liability (group) lending, individual lending, and a control group. Stratification was carried out at the provincial level to ensure balance ($S_i\in\{1, \dots, 5\}$), as complete randomization could have led to some provinces containing only some particular treatment or control villages. This type of geographic stratification is commonly employed in social sciences, medicine and other disciplines to facilitate robust comparisons between treatment groups.

The experiment is one of many studies to evaluate the effectiveness of microcredit as a tool for alleviating global poverty. Since the launch of microcredit programs by the Grameen Bank in Bangladesh, which earned its founder Muhammad Yunus the Nobel Peace Prize in 2006, skepticism about their impact has grown in later years, fueled by the publication of findings on the subject. A primary goal of microcredit programs is to promote investment in and expansion of small-scale enterprises. However, \citet{attanasio2015impacts} found no significant average impact of these programs on enterprise revenue, profit, or other income sources.

In this paper, we revisit their analysis to estimate the distributional treatment effects of the lending program on enterprise revenue ($Y_i$), aiming to uncover potential heterogeneity beyond the average effects across the distribution. Our focus is on the comparison between joint-liability lending ($W_i=2$) and the control group ($W_i=1$), as this was the central concern of the original study. Notably, group lending was pioneered by the Grameen Bank in Bangladesh during the 1970s.

Figure \ref{fig:microfinance} depicts the distributional and probability treatment effects of joint-liability lending on enterprise revenue. The outcome is measured in thousands of Mongolian Tugriks (MNT), with an exchange rate of 1 USD = 1150 MNT at the time of the study. We compute the DTE and PTE for $y\in\{0, 10, \dots, 200\}$ accounting for the stratified design. For regression adjustment, we use gradient boosting with 10-fold cross-fitting, with pre-treatment covariates ($X_i$) including enterprise revenue prior to the experiment, household size, education level, age, etc. The full list of covariates can be found in Table \ref{tab:microfinance_covariates} in the Appendix.

The top-left panel of Figure \ref{fig:microfinance} presents the empirical DTE, while the top-right panel shows the regression-adjusted DTE. The shaded areas represent the 95\% pointwise confidence band computed using multiplier bootstrap \cite{gine1984some, belloni2017program} with 1000 repetitions. Although the sample size in the experiment is modest ($n=611$), the regression adjustment reduces the standard errors by 1\% to 13\% across revenue levels, with an average reduction of 7\%. The  bottom-left panel of Figure \ref{fig:microfinance} presents the empirical PTE, and the bottom-right panel depicts the regression-adjusted PTE on enterprise revenue. For the PTE, the effectiveness of regression adjustment is limited by the small sample size, and it does not consistently reduce standard errors compared to the empirical estimator.

The analysis of DTE and PTE indicates that regression adjustment reduces the standard error by approximately 10\% when estimating the probability of revenue being zero, leading to a statistically significant negative effect. Specifically, the probability of revenue being zero decreases by 10 percentage points (pp), with a standard error of 4.6 pp.

Overall, the evidence points to joint-liability lending mainly helping individuals move out of the ``zero-revenue'' trap, with little discernible effect elsewhere in the distribution. Economically, this pattern implies that the program acts more as a hedge against downside risk than as a catalyst for broad productivity gains, so welfare improvements are likely to flow from reduced business failure and smoother consumption rather than from large increases in aggregate output. Because the confidence bands remain wide beyond that lowest bin—even after flexible covariate adjustment—larger samples or richer, more predictive covariates will be needed to sharpen estimates across the rest of the revenue range.

\section{Conclusion} \label{sec:conclusion}
We introduce a novel regression adjustment method designed to efficiently estimate distributional treatment effects under covariate-adaptive randomization. Our framework supports high-dimensional settings with many pre-treatment covariates and enables flexible modeling through the integration of off-the-shelf machine learning techniques for regression adjustment.

Despite its strengths, our method has certain limitations. First, it assumes experimental data with perfect compliance and no interference. While this setup is appropriate for some applications, it may limit applicability in contexts where these assumptions do not hold. Second, the effectiveness of our approach relies on the availability of pre-treatment covariates that are highly predictive of the outcome. Although we leverage flexible machine learning techniques to enhance prediction quality and achieve greater variance reduction compared to linear regression, the potential for variance reduction diminishes when covariates provide limited predictive power. Third, in scenarios with a large number of strata, improving precision using pre-treatment covariates becomes increasingly challenging, particularly when the sample size is modest. These limitations point to several promising directions for future research, such as designing methods to handle imperfect compliance and network effects, as well as enhancing estimation efficiency in settings with a large number of strata and locations. Additionally, extending recent advances in kernel mean embeddings to characterize outcome distributions (e.g., \citet{park2021conditional, naf2024causal}) to the CAR framework represents a promising direction for future research.

\section*{Acknowledgements}
We extend our gratitude to the four anonymous reviewers and the program chairs for their insightful comments and discussions, which significantly enhanced the quality of this paper. Additionally, Oka acknowledges the financial support from JSPS KAKENHI Grant Number 24K04821.

\section*{Impact Statement}
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.


\bibliography{dte-car}
\bibliographystyle{icml2025}


\newpage