Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
51,150 characters · 21 sections · 83 citation commands
Active Adaptive Experimental Design for Treatment Effect Estimation with Covariate Choice
\twocolumn[ \icmltitle{Active Adaptive Experimental Design for Treatment Effect Estimation\\ with Covariate Choice}
\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
Experimental approaches play a pivotal role in uncovering causality in various fields of science and industrial applications, such as epidemiology and economics ChowChangPong,Hahn2011,ChowChang201112,fda. This study focuses on cases with binary treatments. Typically, one of these treatments corresponds to the treatment, and the other to the control imbens_rubin_2015. Among the quantities representing causal effects, this study focuses on an average treatment effect (ATE), which defines a causal effect as the difference between the expected outcomes of binary treatments.
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 imbens_rubin_2015. However, due to the counterfactual property, individual treatment effects cannot be observed. Therefore, we use ATEs to capture causality.
A fundamental experimental approach for estimating ATEs is a 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 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, adaptive experimental designs have gained increasing attention across various fields Laan2008TheCA,Hahn2011,ChowChang201112,villar,fda.
For estimating ATEs efficiently (with smaller asymptotic variances), Laan2008TheCA and 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. Meehan2022, Kato2020adaptive, and 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 Augmented Inverse Probability Weighting (AIPW) estimator with covariate shift adaptation using Importance Weighting (AIPWIW). We refer to the setting of an adaptive experiment as active adaptive experimental design with covariate choice and our designed experiment as the 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. Laan2008TheCA, Hahn2011, Meehan2022, and 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, 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 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 mcbook. In causal inference under a covariate shift, Uehara2020 presents an efficient estimator of policy values, a generalization of ATEs, using double machine learning 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) in Section (ref). 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 active adaptive experiment for ATE estimation.
\paragraph{Contributions} Compared to existing work, our contributions are summarized as follows:
This section provides our problem formulation. We define potential outcomes and observations separately by following the Neyman–Rubin causal model Neyman1923,rubin1974. Then, we discuss estimating the ATE from the observations gathered from an adaptive experiment.
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).
We make the following regularity assumption.
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
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), 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.
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 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:
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 active adaptive experiment with covariate choice.
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 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
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 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).
Theorem (ref) 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), 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$.
This section investigates $w$ and $p$ that minimizes $\tau(w, p)$ by considering the following functional optimization problem:
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 efficient covariate density and propensity score since the semiparametric efficiency bound is minimized under them.
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
This efficient propensity score is referred to as the Neyman allocation, which has been traditionally employed in designing experiments with binary treatments without covariate choice Neyman1934OnTT,Laan2008TheCA,Hahn2011,Kato2020adaptive,adusumilli2022minimax,Dai2023.
Next, we find an efficient covariate density that minimizes the semiparametric efficiency bound for fixed $w$.
By substituting $w^*$ obtained in Section (ref), 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).
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:
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
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
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
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
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}$.
We design an adaptive experiment with covariate choice, referred to as the 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) and then show the asymptotic property in Section (ref).
This section provides the definition of our experiment, which consists of the sampling phase and 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).
\paragraph{Sampling phase} Define the following estimators:
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 yang2002 and qian2016kernel. In our simulation studies, we employ the Nadaraya-Watson kernel regression Nadaraya1964,Watson1964, whose consistency has been proved by qian2016kernel in an adaptive experiment when the propensity score is lower bounded by a universal constant.
Then, we estimate the efficient propensity score as
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$
\paragraph{Estimation phase} In the estimation phase, we estimate the ATE $\theta_0$ by using the AIPWIW estimator, defined as
where $\Psi_t(Y_t, A_t, X_t; \widehat{w}_t, \widehat{p}_t) \coloneqq$
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 Adaptive AIPW (A2IPW) estimator proposed in Kato2020adaptive with importance weighting 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 Laan2008TheCA, frequently utilized in semiparametric analysis and recently refined as DML ChernozhukovVictor2018Dmlf.
This section provides an asymptotic distribution of $\widehat{\theta}_T$. First, we make the following assumptions.
This assumption also implies $\widehat{p}_t(x) \xrightarrow{\mathrm{a.s.}} p^*(x)$.
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 klaassen1987, Laan2008TheCA, Luedtke2016, ChernozhukovVictor2018Dmlf, and Kato2021adr. With this technique, we can derive the asymptotic distribution only with the convergence rate condition in Assumption (ref), without assuming the Donsker condition.
We also note that $\widehat{\theta}_T$ is consistent for $\theta_0$, even without assuming Assumption (ref).
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) 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). For example, if $d$ is significantly large, we might use arguments about sparsity or benign overfitting to guarantee that Assumption (ref) holds. Once Assumption (ref) holds, the asymptotic variance of $\widehat{\theta}_T$ does not depend on $d$.
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.
\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
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:
For $x\in[-10, 10]$, we plot $\sigma^2_0(a)(x)$ in Figure (ref).
\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 Laan2008TheCA, Hahn2011, and 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.
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)--(ref), 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.
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). We show the results in Figure (ref).
\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). We show the results of simulation studies in Figure (ref).
\paragraph{Results} From Figures (ref) and (ref), 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 Henmi2004paradox,Hitomi_Nishiyama_Okui_2008.
In this section, we discuss related topics.
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 Wald1945,Robbins1952,Solomon1970, Rosenberger2001,bai2024primer. Pioneering contributions for ATE estimation by Laan2008TheCA and Hahn2011 lay the foundation for this field.
Building upon Laan2008TheCA and Hahn2011, Kato2020adaptive and 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 cook2023semiparametric, who proves the asymptotic normality of the A2IPW estimator under weaker assumptions than those of Kato2020adaptive. They also propose an adaptive experiment that ensures finite-sample stability by truncating the propensity score Waudby-Smith2024. Furthermore, they derive tighter concentration inequalities using results from Howard2020TimeuniformNN, Waudby-Smith2023, and Waudby-Smith2023_nonpara. gupta2021efficient designs an adaptive experiment involving instrumental variables. There are other studies discussing related topics Blackwell2022, Westort2021.
Recent work by 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. 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 Kato2020adaptive based on the arguments in Laan2008TheCA and ChernozhukovVictor2018Dmlf. Kato2021adr and li2023double propose DML for adaptive experiments. hadad2019 and Zhan2021 also focus on this topic, presenting different estimators under a different setting. There are also several approaches for statistical inference in adaptive experiments Zhang2020, Bibaut2021.
In the context of active learning, 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 Sugiyama2009, which is relevant to our context. 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, zrnic2024active explores active learning for causal inference.
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 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 Thompson1933, Lai1985, aimed at identifying the best treatment (arm) Audibert2010, Kaufman2016complexity, Russo2016. BAI has two settings: fixed-confidence Garivier2016 and fixed-budget. The fixed-confidence setting is related to sequential testing bechhofer1968sequential, Garivier2016, while the fixed-budget setting is BAI with a fixed sample size Bubeck2009, Bubeck2011, degenne2023existence. 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, russac2021abn, kato2021role, and simchi2024experimentation investigate the BAI with covariates, which is referred to as contextual information. In the fixed-budget setting, kato2024adaptivepolicylearning designs an adaptive experiment for policy learning.
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 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 Kato2020adaptive.
We can consider several extensions of our designed experiment. In Appendix (ref), 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.
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.
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.