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.
51,139 characters
Active Adaptive Experimental Design for Treatment Effect Estimation with Covariate Choice
\twocolumn[
\icmltitle{\emph{Active} Adaptive Experimental Design for Treatment Effect Estimation\\
with Covariate Choice}
\begin{icmlauthorlist}
\icmlauthor{Masahiro Kato}{mizuho}
\icmlauthor{Akihiro Oga}{mizuho}
\icmlauthor{Wataru Komatsubara}{mizuho}
\icmlauthor{Ryo Inokuchi}{mizuho}
\end{icmlauthorlist}
\icmlaffiliation{mizuho}{Mizuho-DL Financial Technology Co., Ltd., Tokyo, Japan}
\icmlcorrespondingauthor{Masahiro Kato}{[email removed]}
\icmlkeywords{Treatment effect estimation, Adaptive experimental design, Active learning, Causal inference}
\vskip 0.3in
]
\printAffiliationsAndNotice{}
\begin{abstract}
This study designs an adaptive experiment for efficiently estimating \emph{average treatment effects} (ATEs). In each round of our adaptive experiment, an experimenter sequentially samples an experimental unit, assigns a treatment, and observes the corresponding outcome immediately. At the end of the experiment, the experimenter estimates an ATE using the gathered samples. The objective is to estimate the ATE with a smaller asymptotic variance. Existing studies have designed experiments that adaptively optimize the propensity score (treatment-assignment probability). As a generalization of such an approach, we propose optimizing the covariate density as well as the propensity score. First, we derive the efficient covariate density and propensity score that minimize the semiparametric efficiency bound and find that optimizing both covariate density and propensity score minimizes the semiparametric efficiency bound more effectively than optimizing only the propensity score. Next, we design an adaptive experiment using the efficient covariate density and propensity score sequentially estimated during the experiment. Lastly, we propose an ATE estimator whose asymptotic variance aligns with the minimized semiparametric efficiency bound.
\end{abstract}
\section{Introduction}
\emph{Experimental approaches} play a pivotal role in uncovering causality in various fields of science and industrial applications, such as epidemiology and economics \citep{ChowChangPong,Hahn2011,ChowChang201112,fda}. This study focuses on cases with \emph{binary treatments}. Typically, one of these treatments corresponds to the \emph{treatment}, and the other to the \emph{control} \citep{imbens_rubin_2015}. Among the quantities representing causal effects, this study focuses on an \emph{average treatment effect} (ATE), which defines a causal effect as the difference between the expected outcomes of binary treatments.
\emph{Treatments} are an abstraction of an experimenter's choices, including drugs, online advertisements, and economic policies, which are also referred to differently, such as actions and arms. A difference in outcomes of binary treatments (treatment effect) is often considered a causal effect \citep{imbens_rubin_2015}. However, due to the counterfactual property, \emph{individual treatment effects} cannot be observed. Therefore, we use ATEs to capture causality.
A fundamental experimental approach for estimating ATEs is a \emph{randomized control trial} (RCT). In an RCT, one of the binary treatments is randomly assigned to each experimental unit to obtain an unbiased estimator of ATEs. While RCTs are the gold standard and reliable \citep{Hariton2018}, they often require large sample sizes, which can be costly. To reduce the sample size as much as possible, or to reduce an estimation error given a sample size, \emph{adaptive experimental designs} have gained increasing attention across various fields \citep{Laan2008TheCA,Hahn2011,ChowChang201112,villar,fda}.
For estimating ATEs efficiently (with smaller asymptotic variances), \citet{Laan2008TheCA} and \citet{Hahn2011} design adaptive experiments that optimize the propensity score (treatment-assignment probability). They show that by optimizing the propensity score, we can minimize the asymptotic variance of ATE estimators. Note that an asymptotic variance can be interpreted as an asymptotic mean squared error (MSE) of ATE estimation. Furthermore, in hypothesis testing, as an asymptotic variance decreases, the sample size required for hypothesis testing also decreases. \citet{Meehan2022}, \citet{Kato2020adaptive}, and \citet{cook2023semiparametric} design refined experiments based on this approach.
As a generalization of this approach, we consider a novel setting where an experimenter can sample experimental units based on their covariates and covariate density set by the experimenter. We show that by optimizing the covariate density as well as the propensity score, we can further minimize asymptotic variances compared to adaptive experiments that only optimize the propensity score. Then, we design an adaptive experiment that optimizes efficient covariate density and propensity score by using past observations.
Firstly, we derive the semiparametric efficiency bound, a lower bound of the asymptotic variance, which is a functional of a covariate density and a propensity score. Based on this efficiency bound, we compute the efficient covariate density and propensity score that minimize the lower bound. Then, we design an adaptive experiment that sequentially estimates the efficient covariate density and propensity score. In the proposed experiment, an experimental unit is sampled following an adaptively estimated efficient covariate density and randomly assigned treatments following an estimated efficient propensity score. At the experiment's end, the experimenter obtains an estimator whose asymptotic variance aligns with the semiparametric efficiency bound. We employ the \emph{Augmented Inverse Probability Weighting} (AIPW) estimator with covariate shift adaptation using Importance Weighting (AIPWIW). We refer to the setting of an adaptive experiment as \emph{active adaptive experimental design with covariate choice} and our designed experiment as the \emph{Active-Adaptive-Sampling (AAS)-AIPWIW experiment}.
Our method is inspired by three lines of work: adaptive experimental design for efficient ATE estimation, active learning using a covariate shift, and off-policy evaluation under a covariate shift. \citet{Laan2008TheCA}, \citet{Hahn2011}, \citet{Meehan2022}, and \citet{Kato2020adaptive} design adaptive experiments for estimating ATEs with smaller asymptotic variances by optimizing propensity scores. Our result generalizes their work by incorporating covariate density optimization. In active learning of regression models, \citet{Sugiyama2006} demonstrates that the asymptotic variance of regression coefficient estimators can be reduced by appropriately shifting the covariate density. That study deals with a non-adaptive experiment and regression with importance weighting \citep{Shimodaira2000}. In contrast, our target is ATE estimation, and we adaptively optimize the covariate density to estimate ATEs. Note that this approach for active learning is an application of importance sampling \citep{mcbook}. In causal inference under a covariate shift, \citet{Uehara2020} presents an efficient estimator of policy values, a generalization of ATEs, using double machine learning \citep[DML,][]{ChernozhukovVictor2018Dmlf}.
Intuitively, we can reduce the asymptotic variance of the ATE estimator by using more units for a treatment with a larger variance. In existing studies, an experimenter assigns a treatment to an experimental unit with a higher propensity score if it has a larger variance. In this study, we also optimize the choice of experimental units based on their covariates in addition to the propensity score. Under this approach, we use more experimental units if their conditional variances of outcomes are large, conditioned on their covariates. See Figure~\ref{fig:concept_figure2} in Section~\ref{sec:efficiency_prob}. This covariate choice allows us to reduce the asymptotic variance of an ATE estimator.
Our framework introduces a novel aspect of optimizing covariate density, which has not been fully explored in the literature on adaptive experiments for ATE estimation. While active learning focuses on optimizing the covariate density, adaptive experiments typically focus on optimizing the propensity score. We refer to our setting as an \emph{active adaptive experiment} for ATE estimation.
\paragraph{Contributions} Compared to existing work, our contributions are summarized as follows:
\begin{itemize}[topsep=0pt, itemsep=0pt, partopsep=0pt, leftmargin=*]
\item We propose a novel framework for an adaptive experiment by optimizing the covariate density, as well as the propensity score (Section~\ref{sec:problem});
\item We develop the semiparametric efficiency bound for the ATE in our context (Section~\ref{sec:semiparametric_efficiency});
\item Based on this lower bound, we derive efficient covariate density and propensity score (Section~\ref{sec:efficiency_prob});
\item Using these efficient probabilities, we design an adaptive experiment for efficient ATE estimation (Section~\ref{sec:opt_experiment}).
\end{itemize}
\section{Problem Setting}
\label{sec:problem}
This section provides our problem formulation. We define potential outcomes and observations separately by following the Neyman–Rubin causal model \citep{Neyman1923,rubin1974}. Then, we discuss estimating the ATE from the observations gathered from an adaptive experiment.
\subsection{Potential Outcomes}
There is a binary treatment $a \in\{1, 0\}$. Let us define the corresponding potential outcome by $Y(a)$.
Let $X\in \mathcal{X}\subset \mathbb{R}^d$ be a $d$-dimensional covariate, where $\mathcal{X}$ is the space.
For any $x\in\mathcal{X}$, let $P_0(x)$ be the true conditional distribution of $(Y_1, Y_0)$ given $X = x$ whose density is given as $r^1(y(1), y(0)\mid x)$.
Let $\mathbb{E}[\cdot \mid X = x]$ and $\mathrm{Var}\big(\cdot \mid X = x\big)$ be expectation and variance operators over $P_0(x)$.
Let us denote the true mean and variance of $Y(a)$ conditioned by $X = x \in\mathcal{X}$ by $\mu_0(a)(x) = \mathbb{E}[Y(a)\mid X=x]$ and
$\sigma^2_0(a)(x) = \mathrm{Var}(Y(a) \mid X=x)$, respectively.
We cannot observe both $Y(1)$ and $Y(0)$ simultaneously. We only observe $Y(A)$ for a treatment-assignment indicator $A \in \{1, 0\}$. We define how we observe data in Section~\ref{sec:aae}.
We make the following regularity assumption.
\begin{assumption}
\label{asm:dist}
There exist known universal constants $\underline{C}$ and $\overline{C}$ such that $0 < \underline{C} < \overline{C} < \infty$ and for any $x\in\mathcal{X}$, under $P_0(x)$, $\big|\mu_0(a)(x)\big| < \overline{C}$ and $\underline{C} < \sigma^2_0(a)(x) < \overline{C}$ hold.
\end{assumption}
\subsection{Average Treatment Effect}
As a causal effect, we are interested in the difference $Y(1) -Y(0)$, referred to as an ITE in some cases. However, because one of $Y(1)$ and $Y(0)$ is unobservable owing to the counterfactual property, we cannot observe it directly. Therefore, we focus on the expected value of $Y(1) -Y(0)$.
Let $q(x)$ be the density of an evaluation covariate distribution. Our interest is in the ATE over $q(x)$, defined as
\begin{align*}
\theta_0 &\coloneqq \mathbb{E}_{X\sim q(x),Y(a) \sim r^a(y\mid X)}[Y(1) - Y(0)],
\end{align*}
where $\mathbb{E}_{X\sim q(x),Y(a) \sim r^a(y\mid X)}$ is expectation over $X$ generated from $q(x)$ and $Y(a)$ given $X$ generated from $r^a(y\mid X)$. Note that $\theta_0 = \mathbb{E}_{X\sim q(x)}[\mathbb{E}[Y(1) - Y(0)\mid X]] = \mathbb{E}_{X\sim q(x)}\left[\mu_0(1)(X) - \mu_0(0)(X)\right]$ holds, where $\mathbb{E}_{X\sim q(x)}$ denotes expectation over $X$ generated from $q(x)$.
Our interest is to estimate the ATE with a smaller asymptotic variance via the following adaptive experiment.
For simplicity, $q(x)$ is assumed to be known to the experimenter. This assumption can be mitigated by considering different settings. In Appendix~\ref{sec:rejection}, we introduce another setting where we can observe covariates $X_t$ generated from $q(x)$ and decide whether we use it by the rejection sampling.
\begin{assumption}
\label{asm:covariate}
The covariate density $q(x)$ is known.
\end{assumption}
\subsection{Active Adaptive Experiments}
\label{sec:aae}
Consider an adaptive experiment with $T$ rounds, $[T] \coloneqq\{1,2,\dots, T\}$. In the experiment, in each $t\in[T]$, an experimenter decides a covariate density $p_t(x)$ and observes an experimental unit with covariates $X_t \sim p_t(x)$. Given the observed covariates $X_t$, the experimenter decides propensity score $w_t(a\mid X_t)$ and assigns treatment $A_t \in \{1, 0\}$, where $A_t \in \{1, 0\}$ is a \emph{treatment indicator}. Then, the experimenter observes an outcome $Y_t=\mathbbm{1}[A_t = 1]Y_{t}(1) + \mathbbm{1}[A_t = 0]Y_{t}(0)$, where $(Y_t(1), Y_t(0))$ is an i.i.d. copy of $(Y(1), Y(0))$. After round $T$, the experimenter estimates $\theta_0$ by using the observations $\{(Y_{t}(1), Y_{t}(0)\}^T_{t=1}$.
In our adaptive experiment, an experimenter can optimize the probabilities of $p_t(x)$ and $w_t(a\mid X_t)$ to obtain an estimator of $\theta_0$ with a smaller asymptotic variance.
In summary, the experimenter takes the following steps:
\begin{description}[topsep=0pt, itemsep=0pt, partopsep=0pt, leftmargin=*]
\item[Step~1:] based on past observations, an experimenter decides a covariate density $p_t$ and a propensity score $w_t$;
\item[Step~2:] the experimenter observes covariate $X_t\sim p_t(x)$;
\item[Step~3:] based on past information, the experimenter decides a propensity score $w_t(a\mid X_t)$;
\item[Step~4:] the experimenter assigns treatment $A_t = a$ with probability $w_t(a\mid X_t)$;
\item[Step~5:] the experimenter observes outcome $Y_t=\mathbbm{1}[A_t = 1]Y_{t}(1) + \mathbbm{1}[A_t = 0]Y_{t}(0)$, where $Y_{t}(a)\sim r^a(y\mid X_t)$;
\item[Step~6:] after iterating Steps~1--5 for $t=1,2,\dots, T$, in round $T$, we estimate the ATE.
\end{description}
We assume that $r^a(y\mid x)$ is invariant across rounds, and $Y_t(a)$ are drawn from the density independently of the other round's random variables; that is, $\{(Y_{t}(1), Y_{t}(0)\}^T_{t=1}$ is i.i.d., but $p_t(a)$ and $w_t(a\mid x)$ can take different values across rounds based on past observations.
Thus, the adaptive experiments yields observations $\{(Y_t, A_t, X_t)\}^T_{t=1}$, where $(Y_t, A_t, X_t)$ and $(Y_s, A_s, X_s)$ are correlated over time for $t \neq s$, that is, the samples are not i.i.d. Let $\mathcal{F}_{t-1}=\{X_{t-1}, A_{t-1}, Y_{t-1}, \dots, X_{1}, A_1, Y_{1}\}$ be the history. The probability $w_t(a\mid x)$ is a function of a covariate $X_t$, an action $A_t$, and history $\mathcal{F}_{t-1}$.
We refer to an experiment with this setting as an \emph{active adaptive experiment with covariate choice}.
\begin{remark}[Stable unit treatment value assumption (SUTVA)]
\label{rem:sutva}
The DGP implies SUTVA; that is, $r^a(y\mid x)$ is invariant against $A_t$ across $t\in[T]$ \citep{Angrist1996}.
\end{remark}
\begin{figure}[t]
\centering
\includegraphics[width=80mm]{Figure/concet_figure1-eps-converted-to.pdf}
\vspace{-5mm}
\caption{Active adaptive experiment.}
\label{fig:exp3}
\vspace{-5mm}
\end{figure}
\section{Semiparametric Efficiency Bound}
\label{sec:semiparametric_efficiency}
This section provides a lower bound for the asymptotic variance of ATE estimators. We focus on semiparametric efficiency bound for a class of regular estimators. The semiparametric efficiency bound generalize the Cramer-Rao lower bound, where the latter considers parametric models, while the former considers non- and semiparametric models.
The semiparametric efficiency bound is a lower bound for a class of regular estimators, which are $\sqrt{T}$-consistent estimators for all DGPs and remain so under perturbations of size $1/\sqrt{T}$ to the DGP
(see \citet{Vaart1998}).
For simplicity, suppose that $(Y_t, A_t, X_t)$ is an i.i.d. copy of $(Y, A, X)$.
Let $r^a(y\mid x)$ be a density of $Y(a)$ conditioned on $X = x$, $w(d\mid x)$ be a probability of $A$ conditioned on $X = x$, and $p(x)$ be a density of $X$.
Suppose that $(Y, A, X)$ is generated as $(Y, A, X) \sim f(y, a, x)$, where
\begin{align*}
f(y, a, x) = \prod_{d\in\{1, 0\}}\left(r^d(y(d)\mid x) w(d\mid x)\right)^{a=d}p(x),
\end{align*}
where $r^d(y(d)\mid x)$ is the marginal distribution of $Y(d)$.
Note that in our setting, we can manipulate $w(a\mid x)$ and $p(x)$ and update them during an experiment.
Furthermore, we \emph{hypothetically} assume access to a set of covariates sampled from $q(x)$. Then, to derive the semiparametric efficiency bound, we consider a stratified sampling scheme, where two independent datasets $\mathcal{D}_T \coloneqq \{(Y_t, A_t, X_t)\}^T_{t=1}$ and $\widetilde{\mathcal{D}}_S \coloneqq \{\widetilde{X}_s\}^S_{s=1}$ are available. Note that $(Y_t, A_t, X_t) \stackrel{\mathrm{i.i.d}}{\sim} \prod_{d\in\{1, 0\}}\Big\{r^d(y\mid x) w(d\mid x)\Big\}^{a = d}p(x)$ and $\widetilde{X}_s \stackrel{\mathrm{i.i.d}}{\sim} q(x)$. First, we derive the semiparametric efficiency bound under a situation where $\mathcal{D}_T$ and $\widetilde{\mathcal{D}}_S$ are available. Then, since we know $q(x)$ and can generate infinitely many samples from $q(x)$, we focus on the semiparametric efficiency bound at a limit of $S\to\infty$.
Under this DGP and $q(x)$,
we obtain the following theorem. The proof is shown in Appendix~\ref{appdx:thm:lower}.
\begin{theorem}
\label{thm:lower}
Under Assumption~\ref{asm:covariate}, given $\mathcal{D}_T$ and $\widetilde{\mathcal{D}}_S$, as $S\to\infty$, the semiparametric efficiency bound for $\theta_0$ with fully nonparametric models is
\begin{align*}
&\tau(w, p)=\mathbb{E}_{X\sim q(x)}\left[\left(\frac{\sigma^2_0(1)(X)}{w(1\mid X)} + \frac{\sigma^2_0(0)(X)}{w(0\mid X)}\right)\frac{q(X)}{p(X)}\right].
\end{align*}
\end{theorem}
Theorem~\ref{thm:lower} is a lower bound; that is, a theoretically best asymptotic variance. It means that even though we minimize the lower bound, it is still unclear whether an estimator whose asymptotic variance aligns with the lower bound exists. In Section~\ref{sec:opt_experiment}, we actually design an adaptive experiment under which an ATE estimator's asymptotic variance aligns with the lower bound minimized over $w$ and $p$.
This semiparametric efficiency bound can be regarded as a functional of $p$ and $w$. During our adaptive experiments, we can minimize the bound regarding $p$ and $w$.
\section{Efficient Probabilities}
\label{sec:efficiency_prob}
This section investigates $w$ and $p$ that minimizes $\tau(w, p)$ by considering the following functional optimization problem:
\begin{align}
\label{eq:functional}
(w^*, p^*) := \operatorname*{arg\,min}_{(w, p)\in\mathcal{W}\times \mathcal{P}}\tau(w, p),
\end{align}
where $\mathcal{W}$ is a set of all measurable functions
$w:\mathcal{X}\to \Delta \coloneqq \big\{ u = (u_1\ u_2\ \dots\ u_K)^\top \in [0, 1]^K \mid \sum^K_{k=1}u_k = 1 \big\}$ with its $a$-th element $w(a\mid x)$, and $\mathcal{P}$ is a set of all measurable functions $p:\mathcal{X} \to \mathbb{R}^+$ such that $\int_{x\in\mathcal{X}} p(x) \mathrm{d}x = 1$. We refer to the minimizers (stationary point) $p^*$ and $w^*$ as \emph{efficient} covariate density and propensity score since the semiparametric efficiency bound is minimized under them.
\subsection{Efficient Propensity Score}
\label{sec:eff_prop}
First, we find an efficient propensity score that minimizes the semiparametric efficiency bound given some $p(x)$. Given $p$, a stationary point regarding $w$ is given as the point-wise optimizer for each $x \in \mathcal{X}$. That is, given $x$, for $w(x) \coloneqq \{w(a\mid x)\}_{a\in[K]} \in (0, 1)^K$, we have $w^*(x) = \operatorname*{arg\,max}_{w\in\Delta} \Big\{\frac{\sigma^2_0(1)(x)}{w(1\mid x)} + \frac{\sigma^2_0(0)(x)}{w(0\mid x)}\Big\}\frac{q^2(x)}{p(x)}$. Then,
for each $x\in\mathcal{X}$, the efficient propensity score is
\begin{align*}
&w^*(a\mid x) = \frac{\sigma_0(a)(x)}{\sigma_0(1)(x) + \sigma_0(0)(x)}\qquad \forall a \in \{1, 0\}.
\end{align*}
This efficient propensity score is referred to as the \emph{Neyman allocation}, which has been traditionally employed in designing experiments with binary treatments without covariate choice \citep{Neyman1934OnTT,Laan2008TheCA,Hahn2011,Kato2020adaptive,adusumilli2022minimax,Dai2023}.
\subsection{Efficient Covariate Density}
Next, we find an efficient covariate density that minimizes the semiparametric efficiency bound for fixed $w$.
\begin{theorem}
\label{thm:covariate_prob}
For each $x\in\mathcal{X}$, the minimizer $p^*$ of the semiparametric efficiency bound $\tau(w, p)$ with a fixed $w$ is
\begin{align*}
&p^*(x) = \frac{\sqrt{\frac{\sigma^2_0(1)(x)}{w(1\mid x)} + \frac{\sigma^2_0(0)(x)}{w(0\mid x)}}}{\mathbb{E}_{X\sim q(x)}\left[\sqrt{\frac{\sigma^2_0(1)(X)}{w(1\mid X)} + \frac{\sigma^2_0(0)(X)}{w(0\mid X)}}\right]}q(x).
\end{align*}
\end{theorem}
By substituting $w^*$ obtained in Section~\ref{sec:eff_prop}, we have
\[p^*(x) = \frac{\sigma_0(1)(x) + \sigma_0(0)(x)}{\mathbb{E}_{X\sim q(x)}\left[\sigma_0(1)(x) + \sigma_0(0)(x)\right]}q(x).\]
This result implies that we can reduce the efficiency bound by drawing samples whose sums of conditional standard deviations of potential outcomes, $\sigma_0(1)(x) + \sigma_0(0)(x)$, are large. That is, we assign treatments to experimental units with high-variance outcomes more than those with less-variance outcomes to accurately estimate the ATE.
For example, in clinical trials, to reduce the ATE, we gather elderly with high blood pressure more than healthy youth because the former may have more volatile outcomes. We illustrate this intuition in Figure~\ref{fig:concept_figure2}.
\begin{figure}[t]
\centering
\includegraphics[width=80mm]{Figure/concet_figure2-eps-converted-to.pdf}
\vspace{-5mm}
\caption{Efficient covariate density. The distribution with the blue line represents the covariate density of interest, while that with the dotted black line represents the covariate density in an experiment.}
\label{fig:concept_figure2}
\vspace{-3mm}
\end{figure}
\subsection{Minimized Semiparametric Efficiency Bound}
The minimizer of $\tau(w, p)$ is given as $w^*$ and $p^*$. By substituting them, we can obtain the minimized semiparametric efficiency bound $\tau^* \coloneqq \tau(w^*, p^*)$ as follows:
\begin{align*}
\tau^* &=
\left(\mathbb{E}_{X \sim q(x)}\left[\sqrt{\frac{\sigma^2_0(1)(X)}{w^*(1\mid X)} + \frac{\sigma^2_0(0)(X)}{w^*(0\mid X)}}\right]\right)^2\\
&=
\Big(\mathbb{E}_{X \sim q(x)}\left[\sigma_0(1)(X) + \sigma_0(0)(X)\right]\Big)^2.
\end{align*}
\subsection{Efficiency Gain}
This section investigates how the efficient probability assignment probability and covariate density improve the semiparametric efficiency bound (efficiency gain).
\paragraph{Efficient propensity score} As a result of the optimization, for each $x\in\mathcal{X}$, for any $w \in \mathcal{W}$, we have
\begin{align*}
&\frac{\sigma^2_0(1)(x)}{w(1\mid x)} + \frac{\sigma^2_0(0)(x)}{w(0\mid x)}\\
&\geq \frac{\sigma^2_0(1)(x)}{w^*(1\mid x)} + \frac{\sigma^2_0(0)(x)}{w^*(0\mid x)} = \big(\sigma_0(1)(x) + \sigma_0(0)(x)\big)^2.
\end{align*}
We confirm how the lower bound is minimized by considering the uniform sampling $w^{\mathrm{Unif}}(a\mid x) = 1/2$ for each $a\in\{1, 0\}$ and $x\in\mathcal{X}$. We have
\begin{align*}
&\frac{\sigma^2_0(1)(x)}{w^{\mathrm{Unif}}(1\mid x)} + \frac{\sigma^2_0(0)(x)}{w^{\mathrm{Unif}}(0\mid x)}= 2\Big(\sigma^2_0(1)(x) + \sigma^2_0(0)(x)\Big)\\
&\geq \big(\sigma_0(1)(x) + \sigma_0(0)(x)\big)^2= \frac{\sigma^2_0(1)(x)}{w^*(1\mid x)} + \frac{\sigma^2_0(0)(x)}{w^*(0\mid x)}.
\end{align*}
The difference between them is $\big(\sigma_0(1)(x) - \sigma_0(0)(x)\big)^2$. Thus, as the difference between $\sigma_0(1)(x)$ and $\sigma_0(0)(x)$ becomes large, the efficiency gain also becomes large.
\paragraph{Efficient covariate density}
Let us fix the propensity score as $w^*(a\mid x) = \frac{\sigma_0(1)(x)}{\sigma_0(1)(x) + \sigma_0(0)(x)}$ and consider a case where the covariate density is fixed at $q(x)$.
When $(Y, A, X)$ is generated from a density $\prod_{d\in[K]}\big\{r^a(y\mid x) w^*(d\mid x)\big\}^{d=a}q(x)$,
the semiparametric efficiency bound becomes
\begin{align*}
\widetilde{\tau} &\coloneqq \mathbb{E}_{X\sim q(x)}\Bigg[\frac{\sigma^2_0(1)(X)}{w^*(1\mid X)} + \frac{\sigma^2_0(0)(X)}{w^*(0\mid X)}\Bigg]\\
&= \mathbb{E}_{X\sim q(x)}\left[\big(\sigma_0(1)(X) + \sigma_0(0)(X)\big)^2\right].
\end{align*}
That is, any regular estimators cannot obtain a smaller asymptotic variance than the lower bound when samples are generated from $\prod_{d\in[K]}\Big\{r^a(y\mid x) w^*(d\mid x)\Big\}^{d=a}q(x)$.
Next, we compare the efficiency bound $\tau^*$ when samples are generated from $\prod_{d\in[K]}\Big\{r^a(y\mid x) w^*(d\mid x)\Big\}^{d=a}p^*(x)$.
From the Jensen inequality, it follows that
\begin{align*}
\tau^* &= \Big(\mathbb{E}_{X\sim q(x)}\left[\sigma_0(1)(X) + \sigma_0(0)(X)\right]\Big)^2\\
&\leq
\mathbb{E}_{X\sim q(x)}\left[\big(\sigma_0(1)(X) + \sigma_0(0)(X)\big)^2\right]= \widetilde{\tau}.
\end{align*}
The equality holds when $\big(\sigma_0(1)(x) + \sigma_0(0)(x)\big)^2$ is the same across $x \in \mathcal{X}$. Therefore, when $\big(\sigma_0(1)(x) + \sigma_0(0)(x)\big)^2$ varies across $x$, we can reduce the semiparametric efficiency bound by optimizing $p(x)$; that is $\tau^* < \widetilde{\tau}$.
\section{Active Adaptive Experimental Design for Treatment Effect Estimation}
\label{sec:opt_experiment}
We design an adaptive experiment with covariate choice, referred to as the \emph{Active-Adaptive-Sampling-and-Augmented-Inverse-Probability-Weighting} (AAS-AIPWIS) experiment because our approach resembles active learning, which adaptively optimizes covariate density and the propensity score, and then estimate the ATE by the AIPWIW estimator. We first define our experiment in Section~\ref{sec:aaed} and then show the asymptotic property in Section~\ref{sec:asymp_dist}.
\subsection{AAS-AIPWIS Experiment}
\label{sec:aaed}
This section provides the definition of our experiment, which consists of the \emph{sampling phase} and \emph{estimation phase}. In round $t \in [T]$, an experimenter engages in the sampling phase, where the experimenter optimizes the covariate density and the propensity score using past observations up to round $t$; then, the experimenter observes covariates and assigns a treatment following the efficient probabilities. After round $T$, the estimation phase starts, where we estimate the ATE using $\{(Y_t, A_t, X_t)\}^T_{t=1}$. We describe the details below, and its pseudo-code is provided in Algorithm~\ref{alg}.
\paragraph{Sampling phase}
Define the following estimators:
\begin{itemize}[topsep=0pt, itemsep=0pt, partopsep=0pt, leftmargin=*]
\item \textbf{Conditional mean estimator.} Let $\widehat{\mu}_t(a)(x)$ be an $\mathcal{F}_{t-1}$-measurable estimator of $\mu_0(a)(x)$ for each $a\in\{1, 0\}$ and $x\in\mathcal{X}$ such that $|\widehat{\mu}_t(a)(x)| < \overline{C}$ holds, and it converges to $\mu_0(a)(x)$ almost surely.
\item \textbf{Conditional second-moment estimator.} Let $\widehat{\nu}_t(a)(x)$ be an $\mathcal{F}_{t-1}$-measurable estimator of $\nu_0(a)(x) \coloneqq \mathbb{E}[Y^2_t(a)\mid x]$ for each $a\in\{1, 0\}$ and $x\in\mathcal{X}$.
\item \textbf{Conditional variance estimator.} Let $\widehat{\sigma}^2_t(a)(x)$ be an $\mathcal{F}_{t-1}$-measurable bounded estimator of $\sigma^2_0(a)(x)$ for each $a\in\{1, 0\}$ and $x\in\mathcal{X}$, defined as
\begin{align*}
\widehat{\sigma}^2_t(a)(x) = \mathrm{thre}\left(\widehat{\nu}_t(a)(x) - \widehat{\mu}^2_t(a)(x); \underline{C}, \overline{C}\right),
\end{align*}
where $\mathrm{thre}(A; a, b) = A\ \mathrm{if}\ a \leq A \leq b,\ \ \mathrm{thre}(A; a, b) = a\ \mathrm{if}\ A < a,\ \mathrm{and}\ \mathrm{thre}(A; a, b) = b\ \mathrm{if}\ A > b$.
\end{itemize}
Specifically, we construct $\widehat{\mu}_t(a)(x)$ and $\widehat{\nu}_t(a)(x)$ only by using past observations $\mathcal{F}_{t-1} = \{(Y_s, A_s, X_s)\}^{t-1}_{s=1}$.
For $\widehat{\mu}_t(a)(x)$ and $\widehat{\nu}_t(a)(x)$, we can use nonparametric estimators, such as the nearest neighbor regression estimator and kernel regression estimator, which have been proven to converge to the true function almost surely under a bounded sampling probability $\widehat{w}_t$ by \citet{yang2002} and \citet{qian2016kernel}. In our simulation studies, we employ the Nadaraya-Watson kernel regression \citep{Nadaraya1964,Watson1964}, whose consistency has been proved by \citet{qian2016kernel} in an adaptive experiment when the propensity score is lower bounded by a universal constant.
Then, we estimate the efficient propensity score as
\begin{align}
\label{eq:est_efficient_prob}
&\widehat{p}_t(x) \coloneqq \frac{\widehat{\sigma}_t(1)(x) + \widehat{\sigma}_t(0)(x)}{\mathbb{E}_{X\sim q(x)}\left[\widehat{\sigma}_t(1)(x) + \widehat{\sigma}_t(0)(x)\right]}q(x)\nonumber\\
&\widehat{w}_t(a\mid x) \coloneqq \frac{\widehat{\sigma}_t(a)(x)}{\widehat{\sigma}_t(1)(x) + \widehat{\sigma}_t(0)(x)},\\
&\mathrm{where}\ \ \widehat{\theta}_t(x) \coloneqq \widehat{\mu}_t(1)(x) - \widehat{\mu}_t(0)(x).\nonumber
\end{align}
We assumed that we can directly compute $\mathbb{E}_{X\sim q(x)}\left[\widehat{\sigma}_t(1)(X) + \widehat{\sigma}_t(0)(X)\right]$ by using the known $q(x)$. For example, we generate enough samples from $q(x)$ and compute the expectation.
Instead of the true expectation, by using past observations, we can approximate it as $\mathbb{E}_{q(x)}\left[\widehat{\sigma}_t(1)(X) + \widehat{\sigma}_t(0)(X)\right] \approx$
\begin{align*}
&\frac{1}{t - 1}\sum^{t-1}_{s=1}\big(\widehat{\sigma}_t(1)(X_s) + \widehat{\sigma}_t(0)(X_s)\big) \frac{q(X_s)}{\widehat{p}_s(X_s)}.
\end{align*}
\paragraph{Estimation phase}
In the estimation phase, we estimate the ATE $\theta_0$ by using the AIPWIW estimator, defined as
\begin{align}
\label{eq:AIPWIW}
\widehat{\theta}_T \coloneqq \frac{1}{T}\sum^T_{t=1}\Psi_t(Y_t, A_t, X_t; \widehat{w}_t, \widehat{p}_t),
\end{align}
where $\Psi_t(Y_t, A_t, X_t; \widehat{w}_t, \widehat{p}_t) \coloneqq$
\begin{align*}
&\Bigg(\frac{\mathbbm{1}[A_t = 1]\big(Y_t(1) - \widehat{\mu}_t(1)(X_t)\big)}{\widehat{w}_t(1\mid X_t)} \\
&\ \ \ \ \ \ \ \ \ \ \ \ - \frac{\mathbbm{1}[A_t = 0]\big(Y_t(0) - \widehat{\mu}_t(0)(X_t)\big)}{\widehat{w}_t(0\mid X_t)}\Bigg)\frac{q(X_t)}{\widehat{p}_t(X_t)}\\
\nonumber
&\ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ \ + \mathbb{E}_{\widetilde{X}\sim q(x)}\left[\widehat{\theta}_t(\widetilde{X})\right],
\end{align*}
where $\mathbb{E}_{\widetilde{X}\sim q(x)}\left[\widehat{\theta}_t(\widetilde{X})\right]$ is computed using the known $q(x)$. This estimator is a variant of the \emph{Adaptive AIPW (A2IPW) estimator} proposed in \citet{Kato2020adaptive} with importance weighting \citep{Uehara2020}.
This estimator removes a sample selection bias caused by the propensity score, which assigns treatments based on covariates. Furthermore, by employing conditional mean estimators $\widehat{\mu}_t$ constructed from past observations, we can employ $\mathbb{E}\left[\Psi_t(Y_t, A_t, X_t; \widehat{w}_t, \widehat{p}_t)\mid \mathcal{F}{t-1}\right] = \theta_0$. This method is a form of sample-splitting technique, as discussed in \citep{Laan2008TheCA}, frequently utilized in semiparametric analysis and recently refined as DML \citep{ChernozhukovVictor2018Dmlf}.
\subsection{Asymptotic Distribution and Efficiency}
\label{sec:asymp_dist}
This section provides an asymptotic distribution of $\widehat{\theta}_T$. First, we make the following assumptions.
\begin{assumption}
\label{asm:subgaussian}
Random variables $Y_t(a) - \mu_a(a)(x)$ and $X_t - \mathbb{E}[X_t]$ are mean-zero sub-Gaussian.
\end{assumption}
\begin{assumption}
\label{asm:nuisance_consistency}
For all $a\in\{1, 0\}$ and any $x\in\mathcal{X}$,
\begin{align*}
\widehat{\mu}_t(a)(x) \xrightarrow{\mathrm{a.s.}} \mu_0(a)(x),\quad \widehat{\sigma}^2_t(a)(x) \xrightarrow{\mathrm{a.s.}} \sigma^2_0(a)(x)
\end{align*}
hold as $t\to \infty$.
\end{assumption}
This assumption also implies $\widehat{p}_t(x) \xrightarrow{\mathrm{a.s.}} p^*(x)$.
\begin{theorem}
\label{thm:asymp_dist}
Consider the AAS-AIPWIS experiment.
Suppose that Assumptions~\ref{asm:dist}--\ref{asm:covariate} and \ref{asm:subgaussian}--\ref{asm:nuisance_consistency} hold. Then,
\begin{align*}
\sqrt{T}\left(\widehat{\theta}_T - \theta_0\right) \xrightarrow{\mathrm{d}}\mathcal{N}(0, \tau^*)
\end{align*}
holds as $T \to \infty$.
\end{theorem}
Thus, the asymptotic variance of $\widehat{\theta}_T$ aligns with the semiparametric efficiency bound minimized concerning the covariate density $p(x)$ and the propensity score $w(a \mid x)$. This result implies that our proposed experiment successfully reduces the asymptotic variance, and we cannot achieve a smaller variance within a class of regular estimators.
In the proof, we use the sample-splitting technique developed by \citet{klaassen1987}, \citet{Laan2008TheCA}, \citet{Luedtke2016}, \citet{ChernozhukovVictor2018Dmlf}, and \citet{Kato2021adr}. With this technique, we can derive the asymptotic distribution only with the convergence rate condition in Assumption~\ref{asm:nuisance_consistency}, without assuming the Donsker condition.
We also note that $\widehat{\theta}_T$ is consistent for $\theta_0$, even without assuming Assumption~\ref{asm:subgaussian}.
\begin{corollary}
Suppose that Assumptions~\ref{asm:dist}--\ref{asm:covariate} and \ref{asm:subgaussian} hold. Then, $
\widehat{\theta}_T - \theta_0$ holds as $T \to \infty$.
\end{corollary}
This property is based on the double robustness of the AIPW estimator $\widehat{\theta}_T$, which guarantees the consistency of $\widehat{\theta}_T$ if either $\mu_0(a)(x)$ or $w^*(a \mid x)$ is consistently estimated. In our case, since we know the true value of $w^*(a \mid x)$, the consistency is guaranteed by the double robustness. For this property, even if we cannot estimate $\mu_0(a)(x)$ consistently, we can still consistently estimate $\theta_0$.
Note that the asymptotic variances in the lower bound and our estimator in Theorem~\ref{thm:asymp_dist} do not depend on the dimension $d$. The dimensionality may affect whether the nuisance estimators $\widehat{\mu}_t$ and $\widehat{\sigma}^2_t$ satisfy Assumption~\ref{asm:nuisance_consistency}. For example, if $d$ is significantly large, we might use arguments about sparsity or benign overfitting to guarantee that Assumption~\ref{asm:nuisance_consistency} holds. Once Assumption~\ref{asm:nuisance_consistency} holds, the asymptotic variance of $\widehat{\theta}_T$ does not depend on $d$.
\begin{algorithm}[tb]
\caption{AAS-AIPWIS experiment}
\label{alg}
\begin{algorithmic}
\STATE {\bfseries Parameter:} Positive constants $C_{\mu}$ and $C_{\sigma^2}$.
\STATE {\bfseries Initialization:}
\FOR{$t\in\{1, 2\}$}
\STATE Assign $A_t=1$ at $t = 1$ and $A_t = 0$ at $t = 2$. For each $a\in\{1, 0\}$, set $\widehat{w}_{t}(a\mid x) = 1/2$.
\ENDFOR
\FOR{$t=K+1$ to $T$}
\STATE Construct $\widehat{p}_t(x)$ following (\ref{eq:est_efficient_prob}).
\STATE Draw $X_t$ from $\widehat{p}_t(x)$.
\STATE Construct $\widehat{w}_{t}(1\mid X_t)$ by using the variance estimators.
\STATE Draw $\xi_t$ from the uniform distribution on $[0,1]$.
\STATE $A_t = 1$ if $\xi_t \leq \widehat{w}_{t}(1\mid X_t)$ and $A_t = 0$ otherwise.
\ENDFOR
\STATE Estimate $\theta_0$ by using the AIPWIW estimator $\widehat{\theta}_T$ in (\ref{eq:AIPWIW}).
\end{algorithmic}
\end{algorithm}
\section{Simulation Studies}
This section presents simulation studies. There are two objectives for the simulation. First, through simulation studies, we clarify the problem setting. Second, for the problem, we confirm the soundness of our designed experiment.
\begin{figure}[t]
\centering
\includegraphics[width=70mm]{Figure/Fig1-eps-converted-to.pdf}
\vspace{-5mm}
\caption{Heterogeneous variances.}
\label{fig:hetero}
\vspace{1mm}
\includegraphics[width=70mm]{Figure/Fig2-eps-converted-to.pdf}
\vspace{-5mm}
\caption{Comparison between $q(x)$ and $p^*(x)$ in the experiment with the covariates following the Gaussian distribution.}
\label{fig:gaussian}
\vspace{1mm}
\includegraphics[width=70mm]{Figure/Fig3-eps-converted-to.pdf}
\vspace{-5mm}
\caption{Comparison between $q(x)$ and $p^*(x)$ in the experiment with the covariates following the Uniform distribution.}
\vspace{1mm}
\label{fig:uniform}
\vspace{-5mm}
\end{figure}
\begin{figure}[t]
\centering
\includegraphics[width=70mm]{Figure/SimRes1-eps-converted-to.pdf}
\vspace{-5mm}
\caption{Results of simulation studies with covariates following with a Gaussian distribution $\mathcal{N}(1, 25)$.}
\label{fig:gaussian_res}
\vspace{1mm}
\includegraphics[width=70mm]{Figure/SimRes2-eps-converted-to.pdf}
\vspace{-5mm}
\caption{Results of simulation studies with covariates following with a Uniform distribution $\mathcal{U}(-10, 10)$.}
\label{fig:uniform_res}
\vspace{-5mm}
\end{figure}
\paragraph{Basic setup}
Let $T = 2000$. We specify the means as $\mathbb{E}_{X\sim q(x)}[Y(1)] = 10$ and $\mathbb{E}_{X\sim q(x)}[Y(0)] = 7$, under which the ATE is $\theta_0 = 3$. Let covariates $x$ be one-dimensional. For any $x$, the conditional means are
\begin{align*}
&\mu_0(1)(x) \coloneqq C_1\left(- x + 3x^2 - 1\right),\\
&\mu_0(0)(x) \coloneqq C_0\left(0.1x + 0.2\right),
\end{align*}
where $C_1$ and $C_0$ are parameters so that $\mathbb{E}_{X\sim q(x)}[\mu_0(1)(X)] = \mathbb{E}_{X\sim q(x)}[Y(1)] = 10$ and $\mathbb{E}_{X\sim q(x)}[\mu_0(0)(X)] = \mathbb{E}_{X\sim q(x)}[Y(0)] = 7$.
\paragraph{Heterogeneous variances}
We define $\sigma^2_0(a)(x)$ as follows:
\begin{align*}
&\sigma^2_0(1)(x) \coloneqq 2 + 1.2\mathrm{sin}(2x) +(x + x^2)/25\\
&\sigma^2_0(0)(x) \coloneqq 2 + 0.8\mathrm{cos}(x / 2) +x^2 / 50.
\end{align*}
For $x\in[-10, 10]$, we plot $\sigma^2_0(a)(x)$ in Figure~\ref{fig:hetero}.
\paragraph{Alternative experiments}
We compare our designed experiment with the deterministic RCT (DRCT), the RCT, and the Adaptive-Sampling-and-AIPW (AS-AIPW) experiment. In the deterministic RCT, we deterministically assign treatment $1$ if $t$ is odd and assign treatment $0$ if $t$ is even. In the RCT, we randomly assign treatments with probability $w(1\mid x) = w(0\mid x) = 0.5$ for any $x$. We estimate the ATE by using the sample average in the DRCT and the RCT. In the AS-AIPW experiment, we optimize the propensity score but do not optimize the covariate density (experimental units are just generated from $q(x)$); then, we estimate the ATE with the AIPW estimator without importance weighting. The AS-AIPW estimator is basically the same method proposed in \citet{Laan2008TheCA}, \citet{Hahn2011}, and \citet{Kato2020adaptive}, though we simplified them for simplicity.
Additionally, we conduct the AS-AIPW and AAS-AIPWIW experiments under the true $w^*$ and $p^*$; that is, they are known at the beginning of an experiment. We refer to the experiments as the AS-AIPW (Oracle) and the AAS-AIPWIW (Oracle). Although we use the true values for $w^*$ and $p^*$ in the sampling phase, we estimate $\mu(a)(x)$ in the estimation phase to construct the AIPW estimator.
\textbf{Performance measure.}
In each of our simulations, we conduct $200$ experiments and compute the performance measures. In this section, in the y-axis of Figures~\ref{fig:gaussian_res}--\ref{fig:uniform_res}, we report the empirical MSEs for $\theta_0$; that is, $\frac{1}{200}\sum^{200}_{i=1}(\theta_0 - \hat{\theta}^{(i)}_T)^2$, where $\hat{\theta}^{(i)}_T$ denotes the estimator of ATE in the $i$-th trial. Note that the asymptotic MSEs correspond to the asymptotic variances.
\textbf{Covariates following with a Gaussian distribution.}
Let $q(x)$ be the density of the Gaussian distribution $\mathcal{N}(1, 25)$ with mean $1$ and variance $25$. For $q(x)$, we plot $q(x)$ and $p^*(x)$ in Figure~\ref{fig:gaussian}. We show the results in Figure~\ref{fig:gaussian_res}.
\paragraph{Covariates following with a Uniform distribution}
Let $q(x)$ be a density of the Uniform distribution with support $[-10, 10]$. For $q(x)$, we plot $q(x)$ and $p^*(x)$ in Figure~\ref{fig:uniform}. We show the results of simulation studies in Figure~\ref{fig:uniform_res}.
\paragraph{Results}
From Figures~\ref{fig:gaussian_res} and \ref{fig:uniform_res}, we find that our experiments successfully reduce the MSE compared to the others. Surprisingly, the AS-AIPW and the AAS-AIPW show almost identical or better performances than the oracle ones. Such a paradox often occurs in semiparametric analysis, as pointed out by \citep{Henmi2004paradox,Hitomi_Nishiyama_Okui_2008}.
\section{Related Work and Discussion}
In this section, we discuss related topics.
\subsection{Related Work}
We provide a literature review of related work in the field of adaptive experimental design for efficiently estimating ATEs.
Adaptive sampling has long been investigated with various purposes \citep{Wald1945,Robbins1952,Solomon1970, Rosenberger2001,bai2024primer}. Pioneering contributions for ATE estimation by \citet{Laan2008TheCA} and \citet{Hahn2011} lay the foundation for this field.
Building upon \citet{Laan2008TheCA} and \citet{Hahn2011}, \citet{Kato2020adaptive} and \citet{Meehan2022} propose refined methods by extending the existing methodologies. These studies derive the semiparametric efficiency bound for ATE estimators and specify the efficient propensity score that minimizes this bound.
This approach is further refined by \citet{cook2023semiparametric}, who proves the asymptotic normality of the A2IPW estimator under weaker assumptions than those of \citet{Kato2020adaptive}. They also propose an adaptive experiment that ensures finite-sample stability by truncating the propensity score \citep{Waudby-Smith2024}. Furthermore, they derive tighter concentration inequalities using results from \citet{Howard2020TimeuniformNN}, \citet{Waudby-Smith2023}, and \citet{Waudby-Smith2023_nonpara}. \citet{gupta2021efficient} designs an adaptive experiment involving instrumental variables. There are other studies discussing related topics \citep{Blackwell2022, Westort2021}.
Recent work by \citet{Dai2023} introduces an approach based on online stochastic projected gradient descent for experimental design when the optimal treatment allocation rule depends on unknown parameters. Their study rigorously addresses the issue of variance estimation, previously considered negligible at the order of $O(1/\sqrt{T})$ under asymptotic approximation. They propose a method to effectively estimate variance and conduct experimental design while accounting for the impact of variance estimation errors. \citet{zhao2023adaptive} also investigates the variance estimation problem in Neyman allocation.
The AIPW estimator is popular in adaptive experiments. Utilizing its variance reduction and martingale properties, the A2IPW estimator is introduced in \citet{Kato2020adaptive} based on the arguments in \citet{Laan2008TheCA} and \citet{ChernozhukovVictor2018Dmlf}. \citet{Kato2021adr} and \citet{li2023double} propose DML for adaptive experiments. \citet{hadad2019} and \citet{Zhan2021} also focus on this topic, presenting different estimators under a different setting. There are also several approaches for statistical inference in adaptive experiments \citep{Zhang2020, Bibaut2021}.
In the context of active learning, \citet{Sugiyama2006} demonstrates that an appropriate covariate density choice can reduce the asymptotic variance of regression coefficients. This setting is extended to the pool-based setting by \citet{Sugiyama2009}, which is relevant to our context. \citet{Wu2019pool} further extends these approaches to adaptive experiments, contrasting with earlier non-adaptive experiments where observations are gathered without parameter estimation during the experiment. Independently, \citet{zrnic2024active} explores active learning for causal inference.
\citet{Uehara2020} considers off-policy evaluation, a generalization of ATE estimation, under a covariate shift. They show that ATE estimation with importance weighting, as proposed by \citet{Shimodaira2000}, with nonparametrically estimated density-ratio, may not yield an estimator whose asymptotic variance aligns with the semiparametric efficiency bound.
Best-Arm Identification (BAI), another related topic, is a type of stochastic multi-armed bandit (MAB) problem \citep{Thompson1933, Lai1985}, aimed at identifying the best treatment (arm) \citep{Audibert2010, Kaufman2016complexity, Russo2016}. BAI has two settings: fixed-confidence \citep{Garivier2016} and fixed-budget. The fixed-confidence setting is related to sequential testing \citep{bechhofer1968sequential, Garivier2016}, while the fixed-budget setting is BAI with a fixed sample size \citep{Bubeck2009, Bubeck2011, degenne2023existence}. \citet{kato2024locally, kato2024worstcase, kato2024agna} investigate and generalize the Neyman allocation in fixed-budget BAI, proving its asymptotic optimality.
In the MAB problem, covariates are referred to as contexts. While various algorithms with contexts have been investigated in the MAB problem, it is still under exploration how to utilize contexts in BAI. In the fixed-confidence setting, \citet{russac2021abn}, \citet{kato2021role}, and \citet{simchi2024experimentation} investigate the BAI with covariates, which is referred to as contextual information. In the fixed-budget setting, \citet{kato2024adaptivepolicylearning} designs an adaptive experiment for policy learning.
\subsection{Sample Size Calculation}
In statistical hypothesis testing, sample size calculation is important. Here, we discuss how to decide $T$ in our method. Let the null and alternative hypotheses be $H_0: \theta_0 = \mu$ and $H_1(\Delta): \theta_1 = \mu + \Delta$, where $\mu \in \mathbb{R}$ and $\Delta > 0$ are some constants. Let $\alpha \in (0, 1)$ be a significance level such that $\mathbb{P}(\mathrm{reject}\ H_0) = \alpha$, and $\beta \in (0, 1)$ be the Type II error level. Let $R^{\alpha}_T$ be a rejection region defined based on $\alpha$. Then, we are interested in deciding the minimum sample size for hypothesis testing, which is defined as $T^*_{\alpha,\beta}(\Delta) \coloneqq \min\big\{T: \mathbb{P}_{H_1(\Delta)}\big(\widehat{\theta}_T \in R^{\alpha}_T\big) \geq 1 - \beta \big\}$ for some estimator $\widehat{\theta}_T$ of $\theta_0$. Let us assume that $\sqrt{T}\big(\widehat{\theta}_T - \mu\big) \xrightarrow{\mathrm{d}} \mathcal{N}(0, V)$ under $H_0$ and $\sqrt{T}\big(\widehat{\theta}_T - \mu - \Delta\big) \xrightarrow{\mathrm{d}} \mathcal{N}(0, V)$ under $H_1$. Then, when using $t$-statistics for hypothesis testing, we have $T^*_{\alpha,\beta}(\Delta) = \frac{V}{\Delta^2}\big(z_{1-\alpha/2} - z_{\beta}\big)$, where $z_{c}$ is the $c$-quantile of the standard normal distribution \citep{Balsubramani2016,Kato2020adaptive}. For example, in the AAS-AIPWIW, we substitute $\tau^*$ into $V$. This result implies that we can reduce the sample size $T$ for hypothesis testing by reducing the asymptotic variance. For the details of sample size calculation, see Appendix~G in \citet{Kato2020adaptive}.
\subsection{Extensions}
We can consider several extensions of our designed experiment. In Appendix~\ref{appdx:extension}, we propose three extensions: (i) rejection sampling, (ii) batch design, and (iii) experiments with multiple treatments. In the rejection sampling setting, we do not assume that $q(x)$ is known, but we sequentially observe an experimental unit generated from the unknown $q(x)$. By employing rejection sampling, we can decide whether to accept the observed experimental unit and if we accept it, we assign a treatment to it. This setting is feasible without knowing $q(x)$. In (ii), we can update $\widehat{p}_t$ and $\widehat{w}_t$ at certain rounds. For example, we can update them at $t = 100, 200, 300, \dots$ and need to keep using the same $\widehat{p}_t$ and $\widehat{w}_t$ between the two update timings. In (iii), we consider there are treatments more than or equal to three.
In this study, we assumed that $q(x)$ is known. For example, if the population we are interested in consists of the people of a certain country, we can obtain information about their distribution through government statistics. Then, in an experiment, we consider gathering people with different distributions using the method we propose. In applications, we can also consider estimating $q(x)$ using observations from a pilot experiment. If there are $M$ samples from $q(x)$, it is possible to estimate $q(x)$ using those samples. In this case, if $M$ diverges independently of $T$, the current results can be directly applicable (asymptotics of $M\to \infty$ and $T\to \infty$). However, if the divergence rate of $M$ depends on $T$, we need to slightly modify our results and methods for each setting of applications.
\section{Conclusion}
In this study, we designed an adaptive experiment to estimate ATEs efficiently. In our adaptive experiment, we can optimize covariate densities and propensity scores in each round based on past observations obtained until that round. To derive the efficient covariate density and propensity score, we focused on the semiparametric efficiency bound for the ATE. Based on the derived efficient probabilities, we designed an adaptive experiment using the AIPW estimator with importance weighting. Under the designed experiment, we demonstrated that the asymptotic variance of the AIPW estimator aligns with the efficiency bound, minimized with respect to the covariate density and propensity score.
\section*{Impact Statement}
This paper presents an innovative approach to machine learning, specifically focusing on the efficient estimation of ATEs through adaptive experimental design. Our work has the potential to significantly impact various sectors, including healthcare, economics, and social sciences. However, if not applied carefully, the optimization techniques could lead to biased sampling or unethical manipulation of experimental conditions, potentially harming subjects or leading to incorrect conclusions. Furthermore, the adaptive nature of the experiments, relying on sequential sampling and treatment assignment based on past observations, may raise data privacy and consent issues, especially in sensitive fields like healthcare. In conclusion, while our work aims to advance the field of machine learning by providing a novel framework for efficient ATE estimation through adaptive experiments, it is crucial to consider the ethical implications and societal impacts of its application. Researchers and practitioners should be mindful of the potential risks and strive to apply these methods in a responsible and equitable manner.
\bibliographystyle{icml2024}
\bibliography{AEE.bbl}