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.
70,413 characters · 16 sections · 57 citation commands
Prediction-Guided Active Experiments
\RUNAUTHOR{Ao, Chen and Simchi-Levi}
\RUNTITLE{Prediction-Guided Active Experiments}
\TITLE{Prediction-Guided Active Experiments}
\ARTICLEAUTHORS{ \AUTHOR{Ruicheng Ao\textsuperscript{1} \quad\quad Hongyu Chen\textsuperscript{1} \quad\quad David Simchi-Levi\textsuperscript{1,2,3} } \AFF{\textsuperscript{1}Institute for Data, Systems, and Society, Massachusetts Institute of Technology, Cambridge, MA 02139\texorpdfstring{\\}\textsuperscript{2}Department of Civil and Environmental Engineering, Massachusetts Institute of Technology, Cambridge, MA 02139\texorpdfstring{\\}\textsuperscript{3}Operations Research Center, Massachusetts Institute of Technology, Cambridge, MA 02139\texorpdfstring{\\}\EMAIL{\{aorc, chenhy, dslevi, fengzhu\}@mit.edu}} }
\ABSTRACT{ In this work, we introduce a new framework for active experimentation, the Prediction-Guided Active Experiment (PGAE), which leverages predictions from an existing machine learning model to guide sampling and experimentation. Specifically, at each time step, an experimental unit is sampled according to a designated sampling distribution, and the actual outcome is observed based on an experimental probability. Otherwise, only a prediction for the outcome is available. We begin by analyzing the non-adaptive case, where full information on joint distribution of the predictor and the actual outcome is assumed. For this scenario, we derive an optimal experimentation strategy by minimizing the semi-parametric efficiency bound for the class of regular estimators. We then introduce an estimator that meets this efficiency bound, achieving asymptotic optimality. Next, we move to the adaptive case, where the predictor is continuously updated with newly sampled data. We show that the adaptive version of the estimator remains efficient and attains the same semi-parametric bound under certain regularity assumptions. Finally, we validate PGAE's performance through simulations and a semi-synthetic experiment using data from the US Census Bureau. The results underscore the PGAE framework’s effectiveness and superiority compared to other existing methods. }
Experimentation is fundamental to scientific research across diverse fields, from online platforms to the pharmaceutical industry, due to its capability to evaluate the performance of new treatments bakshy2014designing, cerqueira2020adaptive, li2022interference, chen2023efficient. A central question in experimental design lies in optimizing the data collection scheme. Collecting individual data points can be costly, especially for high-quality labels, which often restricts experimenters from gathering sufficient data and can compromise statistical efficiency. At the same time, modern machine learning techniques frequently provide experimenters with a predictor of outcomes. This predictor could be a traditional machine learning model, such as Random Forests, Boosting Trees, or Neural Networks trained on historical data, or even Large Language Models (LLMs) that serve as simulators for human agents li2022interference, gui2023challenge, xie2024can.
This study tackles the question of how existing predictors can be leveraged to create a more efficient data collection strategy. Specifically, we consider a setting in which an experimenter seeks to evaluate a new algorithm's performance with a focus on estimating the mean outcome of a known population under this new approach. This objective contrasts with traditional experimental designs, which often aim to measure treatment effects—that is, the difference in outcomes between the new algorithm (treatment) and an existing one (control) deaton2018understanding, kato24active. Our framework assumes the experimenter possesses sufficient knowledge of control arm and is solely interested in assessing the treatment arm, which is related to clinical trials using external control data (schmidli2020beyond, rahman2021leveraging). This scenario also parallels survey sampling sarndal1980pi, where the primary goal is to estimate the mean outcome for a specific population (e.g., income, education levels, or political views).
Our approach to experimental design considers two factors: the sampling distribution and the experimental probability. Specifically, we assume the experimenter first samples according to a designated sampling distribution, then conducts experiments and collects the outcome data based on the experimental probability; otherwise, only the predicted outcome is available. In other words, not all sampled units are tested subjected to the new algorithm. Some of the outcome of the sampled units replaced by the prediction. In this way, we can potentially minimize the number of experimental units needed and achieves a higher precision given the same budget.
Thus, the target is to optimize both the sampling distribution and experimental probability to minimize the asymptotic variance of the final estimator. This study poses two main questions: (1) How should we optimally allocate the sampling budget with known population density and an existing predictor? (2) How should we conduct inference after data collection? To address the first question, we introduce an experimental design framework called “Prediction-Guided Active Experiments,” which builds on a simple intuition: in a heterogeneous population, more samples should be collected from groups with either higher outcome variance or lower prediction accuracy. For the second question, we propose an estimator that achieves the lowest asymptotic variance for a given sampling scheme. This estimator is built on the recent Prediction-Powered Inference (PPI) framework (angelopoulos2023prediction), which has valid statistical guarantees without knowledge on the precision of the predictor.
We begin by analyzing the non-adaptive setting, where the joint distribution of individual covariates, predictions, and target outcomes is known in advance. In this scenario, we derive efficient sampling densities and experimental probabilities by minimizing the semiparametric efficiency lower bound, which is a lower bound for the asymptotic variance of any regular estimator. The resulting solution is intuitively linked to the conditional variance decomposition of the target outcome with respect to the predictor. We then consider the adaptive setting, where the relationship between predictions and actual targets is initially unknown. In this setting, we propose an adaptive scheme that aligns with the non-adaptive solution and an efficient estimator that achieves the variance lower bound in both settings. Finally, we validate our approach through simulations and semi-synthetic experiments using U.S. Census data, demonstrating the method's efficiency.
Our work is related to three streams of literature: adaptive experimentation, prediction-powered inference, and active learning.
Adaptive Experimentation. Adaptive experimentation involves dynamically adjusting the allocation of treatment and control groups across different populations to optimize the estimation accuracy of treatment effects. The concept of adaptivity was first introduced by neyman1934two and has been extensively developed since then (robbins1952some; solomon1970optimal; rosenberger2001optimal; hahn2011adaptive). A central principle here is that if one group (either treatment or control) demonstrates lower outcome variance, then fewer samples are required for precise estimation. With the growth of online platforms, adaptive experimentation has gained traction due to its practical applicability, as seen in recent studies (tabord2023stratification; zhao2023adaptive; simchi2023multi; ao2024online; kato24active). Unlike these works, our research focuses on enhancing inference by leveraging the predictive power of auxiliary models, offering a novel perspective on adaptive experimentation.
Active Learning. Active learning frameworks enable machine learning algorithms to selectively label specific unlabeled data points for model training, typically to maximize predictive accuracy (settles.tr09; ren2021survey). Sampling is often optimized based on uncertainty measures, with the primary goal of improving model performance (balcan2006agnostic; settles2011theories; sener2018active; sinha2019variational; zhang2023active). Our work aligns with active learning in that both approaches aim to optimize sampling distribution to achieve greater efficiency. However, our focus diverges in that we also incorporate the information from a predictor and the goal is to make inference on the mean outcome, while active learning predominantly aims to enhance model accuracy.
Prediction-Powered Inference (PPI). Our work is also inspired by The literature on prediction-powered inference (PPI) (angelopoulos2023prediction; angelopoulos2023ppi++; zrnic2024cross). PPI assumes access to true labels for only a subset of data, with predictions from an external model for the remainder. This framework employs established semi-parametric methods to conduct valid inference with the partially labeled data (rubin1976inference; robins1994estimation; robins1995semiparametric; chernozhukov2018double). Recent studies, such as zrnic2024active, explore actively collecting labels to optimize statistical efficiency within the PPI framework. Our work extends this by jointly optimizing both the sampling distribution and experimental probabilities, demonstrating higher efficiency gains, as evidenced by our numerical results.
We consider a population, where each individual is characterized by a covariate $X \in \mathcal{X} \subset \mathbb{R}^d$ that represents the demographic (age, sex, occupation, etc.). The distribution of $X$ is denoted by $\mathbb{P}_X$, with a density function $q(x)$, which is assumed to be known to the experimenter. Suppose there is a new algorithm, and the goal of the experimenter is to evaluate the performance of that new algorithm on this population. When the policy is applied to an individual with covariate $X$, the resulting outcome is denoted by $Y \in \mathbb{R}$ with the conditional distribution $\mathbb{P}_{Y \mid X}$. The conditional expectation of the outcome is given by $\mathbb{E}[Y \mid X] = \mu(X)$. The primary objective is to estimate the mean outcome, denoted by $\theta$, defined as:
Now we assume that the experimenter can gather data not only by directly applying the algorithm to a sampled individual and observing the outcome $Y$, but can also by utilizing a predictor $f$ that takes $X$ and some other information $W\in\mathcal{W}\subset\mathbb{R}^{m}$ as input and provides predicted outcomes $F = f(X, W)\in\mathbb{R}$. Here, $W$ can be other characteristics of the user that is not part of the demographic, or it is some outside information like the environment when the experiments are conducted, or it can be the inherent randomness from the prediction model itself. We don't assume we know the distribution of $W$ a priori. The predictor $f$ here can be any machine learning model, such as a random forest, boosting trees, or a large language models (LLMs).
Importantly, the accuracy of $f$ is not assumed a priori; instead, the proposed method will adapt to the accuracy of $f$. When the predictions are accurate, the method will produce an estimator with significantly reduced error compared to traditional approaches. Even if the predictions are not highly accurate, the proposed method will still provide valid uncertainty estimates through confidence intervals.
To start with, let us first examine a naive approach. Let us say the experimenter would just sample $n$ individuals $X_1, X_2, \dots, X_n$ independently from the population distribution $\mathbb{P}_X$ and observe their outcomes $Y_1, \dots, Y_n$. Then, a simple average $\hat{\theta} = \frac{1}{n} \sum_{i=1}^n Y_i$ would yield a consistent estimator for $\theta$. However, this approach can be inefficient in two ways:
Observing the inefficiency in the above naive experimentation scheme, we introduce a framework for Prediction-Guided Active Experiments (PGAE), aimed at improving experimental efficiency in both aspects described above, for which we describle in the following section.
In this active experimentation framework, the experimenter utilizes the known covariate distribution and the prediction tool onhand to strategically select which individuals to sample and whether to apply the algorithm and observe outcomes or rely on the predictor. Specifically, at each time step $t$, the experimenter determines two key quantities: the sampling distribution $p_t(x)$ and the experiment probability $\pi_t(x)$, and the sequence of actions at each time step $t$ is as follows:
After $T$ time steps, the experimenter collects a dataset of the form $\{X_t, F_t, \Delta_t, \Delta_t Y_t\}_{t=1}^T$, where $\Delta_t \in \{0, 1\}$ indicates whether the experiment was conducted ($\Delta_t = 1$) or not ($\Delta_t = 0$). Then an estimation method is used at the final period to estimate the mean outcome. Importantly, compared to the naive method, the number of people that are actually subjected to the new algorithm is $\sum_{t=1}^T\Delta_t$, which is much smaller than $T$ depending on the experimental probability $\pi_t$.
The goal of this adaptive approach is to ensure that the sampling and experimentation are focused where they are most needed—sampling more from high-variance subpopulations and conducting experiments where the predictor is less reliable. This strategy helps to minimize costs while maintaining statistical validity and improving the efficiency of the estimation process. In the following section, we will derive the optimal choice of the $p_t$ and $\pi_t$ for the experimenter along with an estimation scheme that can ensure the estimator achieve the smallest asymptotic variance.
In this section, we focus on the non-adaptive setting where the experimenter chooses the sampling distribution $p(x)$ and the experimental probability $\pi(x)$ at the beginning of the experimental period. Then the experimenter sample $T$ people according to $p(x)$ and conduct actual experiments for people with covariate $x$ with probability $\pi(x)$. Under the non-adaptive setup, the final data $\{X_t, F_t, \Delta_t, \Delta_tY_t\}_{t=1}^T$ are i.i.d. for each $t$, which enables us to conduct further statistical analysis. In particular, we will first derive a variance lower bound for every regular estimator given the $T$ i.i.d. data. Then we obtain the optimal choice of $p(x)$ and $\pi(x)$ by minimizing this lower bound.
To derive the optimal choice of $p(x)$ and $\pi(x)$, we first focus on the semi-parametric efficiency bound for estimating the estimand $\theta$ given the i.i.d. data $Z_t=(X_t, F_t, \Delta_t, \Delta_tY_t)$. The efficiency bound gives the lowest asymptotic variance for a class of regular estimators (See Chapter 23.5 of van2000asymptotic for a formal definition of regular estimators). In short, these estimators are $\sqrt{T}$-consistent estimators and remains so under any perturbation of size $1/\sqrt{T}$ of the data generating process. To start with, we first calculate the efficient influence function in below
The efficient influence function acts as the derivative of the final estimand with respect to the data distribution. By using the efficient influence function in Lemma (ref), we are able to calculate the variance lower bound for the regular estimators, which is simply the variance of the influence function (See Chapter 23 of van2000asymptotic).
Theorem (ref) provides the variance lower bound for the class of regular estimators. In particular, there are two parts in the expression, which corresponds to the variance decomposition of $\operatorname{Var}(Y\mid X)$ with respect to the predictor $F$. The first part involves $\mathbb{E}[\operatorname{Var}(Y\mid X, F)\mid X]$, which is the the part of conditional variance $\operatorname{Var}(Y\mid X)$ that cannot be explained by the predictor $F$. The second part relates to $\operatorname{Var}(\mathbb{E}[Y\mid X, F]\mid X)$, which is the the part that can be explained by the predictor $F$. Here, the experimental probability $\pi(x)$ enters the expression only in the first term, meaning that experimentation will contribute to reducing the variance of the estimator by reducing the unpredictable part of the variance.
To present some more intuition, if $p(x)=q(x)$ and $\pi(x)=1$, which is the naive case where the experimenter samples according to the original distribution and conduct experiments on every sample, the lower bound becomes $\mathbb{E}[\operatorname{Var}(Y\mid X)]$. In this case, the lower bound is attainable by first estimating the conditional mean $\mathbb{E}[Y\mid X]$ and then aggregating them according to the true $q(x)$. Note that for the naive estimator, the asymptotic variance is actually $\operatorname{Var}(Y)$. The difference between $\mathbb{E}[\operatorname{Var}(Y\mid X)]$ and $\operatorname{Var}(Y)$ is because we already know the true covariate distribution, which facilitates us in estimation of the population mean.
Given the semi-parametric lower bound $V(p, \pi)$ for any sampling distribution $p$ and experimental probability $\pi$, we should select those $p$ and $\pi$ that minimize the lower bound. It is straightforward to see that the lower bound in (ref) decreases with $\pi(x)$, which is intuitive since with more actual experiments, we should be able to do estimation at a higher accuracy. Thus, we assume the experimenter faces a budget constraint $\gamma\in[0,1]$, which is the proportion of samples that are subjected to actual experimentation. In fact, for an experimental design with sampling probability $\pi(X)$, its portion of the experimental units can be calculated as $\mathbb{E}_{X\sim p(x)}[\pi(X)]$. Thus, given a pre-assigned experimental ratio $\gamma$, the optimization problem becomes
The solution of the optimization problem (ref) can be characterized in the following theorem.
In Theorem (ref), we give the relative magnitude of the absolute solution $\pi^*(x)$ and $p^*(x)$. One can obtain the actual value of $p^*$ and $\pi^*$ normalizing them according to $\int_{\mathcal{X}} p(x)dx=1$ and $\mathbb{E}_{X\sim p(x)}[\pi(X)]=\int_{\mathcal{X}}p(x)\pi(x)=\gamma$. If we regard $\operatorname{Var}(\mathbb{E}[Y\mid X, F] \mid X)$ as the predictable variance and the $\mathbb{E}[\operatorname{Var}(Y\mid X, F)\mid X]$ as the unpredictable variance. Then the sampling distribution $p^*$ is proportional to the actual density $q$ times the predictable part of the variance. As for the experimental probability $\pi$, it is proportional to the ratio between two parts --- higher unpredictable variances compared to predictable variance will lead to more experimentation in that particular group. The expression relates to the optimal sampling distribution without predictions, which is $q(x)\sqrt{\operatorname{Var}(Y\mid X)}$ (kato24active). The difference is that with a predictor, the conditional variance term becomes the unpredictable component of the conditional variance. This reveals the role of experimentation in the presence of a predictor --- experimentation helps to reduce the part of variance that is unpredictable by the predictor.
In this section, we derive an efficient estimator $\hat{\theta}^*$ that actually achieves the semi-parametric efficiency lower bound $V(p, \pi)$ for any given experimental design rule $p$ and $\pi$. The derivation here follows the standard one-step correction estimator in semi-parametric statistics using sample-splitting techniques (e.g. chernozhukov2018double, kennedy2022semiparametric). Denote $Z_t = (X_t,F_t,\Delta_t,\Delta_tY_t)$ as the $t$-th sample (here we use $\Delta_tY_t$ to emphasize that we can only observe $Y_t$ while $\Delta_t=1$). Let $\hat\mu(X)$ and $\hat\tau(X,F)$ be any estimators of $\mathbb{E}_{}\left[Y\vert X\right]$ and $\mathbb{E}_{}\left[Y\vert X,F\right]$, we define the estimated influence function given $\hat\mu, \hat\tau$ for $Z_t$ as
To reduce the dependence between $\hat\mu,\hat\tau$ and $\{Z_t\}_{t=1}^T$, we apply techniques of sample-splitting and cross-fitting to get efficient estimators. For samples $(Z_1,\dots,Z_T)$, we split them into $K$ disjoint folds randomly. Let $F_t\in[K]$ denote the indicator of the fold for $Z_t$ with probability $\mathbb{P} (F_t=k)=1/K.$ Now we define $\hat \mu^{-k}$ and $\hat\tau^{-k}$ as the estimator of $\mathbb{E}[Y\mid X]$ and $\mathbb{E}[Y\mid X, F]$ trained on observations excluding fold $k$. Our new estimator for $\theta$ is constructed as follows:
The estimator (ref) consists of two parts. The first part is the standard estimator of the mean $\mathbb{E}_{X\sim q(x)}\left[\hat\mu^{-k}(X)\right]$, which is the population mean after obtaining the estimation of the conditional mean $\hat\mu^{-k}(X)$. Note that here, since the covariate distribution $q(x)$ is known, one can calculate the expectation here at high precision through Monte Carlo simulation. The second part is the term involving the influence function. This is the first-order correction term that makes the estimator efficient. Below, we prove that this estimator $\hat\theta_T$ is efficient under some regularity assumptions.
Theorem (ref) holds for every sampling distribution $p$ and experimental probability $\pi$. Specifically, it holds for $p^*$ and $\pi^*$ proposed in Theorem (ref). Combining Theorem (ref), (ref), and (ref), we know that as long as the experimenter follows the optimal $p^*$ and $\pi^*$, and uses the estimator (ref), he will obtain the lowest possible asymptotic variance for estimating the population mean, which solves the experimental design problem in the non-adaptive setting.
In this section, we consider the adaptive setting where the experimenter samples one single or a batch of data points at a time and updates the sampling and experimental probability $p_t, \pi_t$ in a sequential manner, which is the Prediction-Guided Active Experiments (PGAE) framework introduced in Section (ref). We discuss both its implementation detail and asymptotic efficiency.
We first discuss how to implement PGAE in the adaptive setting, where $p_t, \pi_t$ and $f_t$ are adaptively estimated. The key is to use sampled data in two ways: (1) to estimate $\operatorname{Var}(\mathbb{E}[Y\mid X, F]\mid X)$ and $\mathbb{E}[\operatorname{Var}(Y\mid X, F)\mid X]$ consistently so that the sampling rule $p_t$ and $\pi_t$ converge to $p^*$ and $\pi^*$. (2) to update $f_{t}$ so that the prediction achieves a higher precision. To estimate \(\alpha(x) := \text{Var}(\mathbb{E}[Y \mid X, F] \mid X)\) and \(\beta(x) := \mathbb{E}[\text{Var}(Y \mid X, F) \mid X]\), we do the following steps:
After obtaining the estimator $\hat{\alpha}_t(x)$ and $\hat\beta_t(x)$, one can obtain $\hat{p}_t$ and $\hat\pi_t$ as \[\hat p_t(x) = \frac{q(x)\sqrt{\hat{\alpha}_t(x)}}{\mathbb{E}_{X\sim q}[\sqrt{\hat{\alpha}_t(X)}]}, \quad \hat\pi_t(x) = \min\left\{\frac{\gamma\sqrt{\hat{\beta}_t(x)}}{\mathbb{E}_{X\sim \hat p_t}\left[\sqrt{\hat{\beta}_t(X)} \right] },1\right \} \] Here, we explicitly write out the normalization term for both of the estimators. Note that in Step 3, we also truncate the estimated variance $\hat\alpha(x)$ to $[\sigma_{\min}, \sigma_{\max} ]$ to ensure numerical stability. Thus, with the estimated $\hat{p}_t$ and $\hat\pi_t$, one can proceed the sampling phase. One can do the update in a batch manner, where the estimation is conducted after a batch of experimental units is sampled. Next, we discuss how to perform estimation of the population mean $\theta$ in the adaptive setting.
In the adaptive setting, all the adaptive estimators $\hat{\mu}^{(1)}_t, \hat{\mu}^{(2)}_t, \hat\tau^{(1)}_t, \hat\tau^{(2)}_t$ and $\hat\alpha_t, \hat\beta_t$ are trained use only the data in the first $t-1$ time steps, which are $\mathcal{F}_{t-1}$-measurable and are independent of the $t$-th observation. This property facilitate us to construct an efficient estimator as in van2008construction and luedtke2016statistical. We now define the adaptive one-step correction estimator as follows.
where $\widehat{\dot{\Psi}}(Z_t;\hat\mu_t^{(1)}, \hat\tau_t^{(1)})$ is given by
To derive the asymptotic result for estimator $\hat{\theta}_T^{Adap}$, we will assume the adaptive estimators and the predictor converge in an appropriate sense. Moreover, we will also assume the actual distribution of $Y$ and $X$ are sub-Gaussian. Under these two regularity assumptions, we state the efficiency result for $\hat{\theta}_T^{Adap}$ as below.
To conclude this section, we summarize the complete algorithm in Algorithm (ref).
In this section, we conduct both simulation and semi-synthetic experiment to test the performance of the proposed design. In particular, we focus on the sample efficiency, meaning that given the number of units that actually received treatment, the mean squared error of different experimental scheme.
To start with, we first conduct a numerical simulation to validate the performance of the PGAE framework. We assume $X$ and $W$ are both real numbers sampled uniformly from $[-1,1]$, and the target $Y$ satisfies the following relationship \[ Y = 2W + X + XW + \epsilon, \] where $\epsilon$ is a normal random noise satisfies the following condition \[\mathbb{E}[\epsilon\mid W, X]= 0, \qquad \operatorname{Var}[\epsilon\mid W, X] = 4\sin\left(\frac{3\pi X}{2}\right)^2. \] We do not assume the predictor is perfectly accurate, instead, we only assume that the oracle estimator follows from a linear regression \[ F(X, W) = 2W + X.\] And in the PGAE framework, we update and fit $F$ based on linear regression of $Y$ on $W$ and $X$. Under this setup, we have $ \operatorname{Var}(\mathbb{E}[Y\mid X, F]\mid X)\propto (2+X)^2$ and $ \mathbb{E}[\operatorname{Var}(Y\mid X, F)] = \sin(\pi X/2)^2$, which means the optimal experimental design is given by \[ p^*(x)\propto 2+x, \qquad \pi^*(x)\propto \left| \frac{\sin(\pi x/2)}{2+x} \right| \] We plot this setup in Figure (ref). Here, the optimal experimental design will sample more data with larger $X$, which is those data points with higher variances. The experimental probability $\pi^*$ puts more emphasis on data points with low prediction power of $F$, i.e., those points with larger $\mathbb{E}[\operatorname{Var}(Y\mid X, F)\mid X]$.
We test in total 5 policies: (1) PGAE: the algorithm introduced in Algorithm (ref). (2) PGAE-Oracle: Algorithm (ref) with information on the actual distribution of $(X, F, Y)$, hence knows the oracle $p^*$ and $\pi^*$. (3) Naive: Naive policy where data is sampled according to the true distribution $q(x)$ and estimator is constructed by sample average. (4) Opt-Sample: only do optimal sampling but do not use prediction in estimator. (5) PPI-Oracle: method proposed in zrnic2024active, which samples according to true distribution $q$ but conduct experimentation with different probability. This method constructs estimator using the predictions according to Prediction-Powered Inference framework (angelopoulos2023prediction). We conduct the experiments across different experimental proportion $\gamma\in\{0.2, 0.4, 0.6, 0.8, 1.0\}$. In every experiment, the number of people that receive treatment is set to be 3000. Meaning that with predictor that uses prediction, there will be approximately $3000/\gamma$ samples. We repeat for 2000 times and plot the result in Figure (ref).
From the results in Figure (ref), we make the following observation: (1) By optimizing the sampling distribution and experimental probability, the proposed PGAE framework does achieve lower MSE across all scenarios. (2) By comparing with the Opt-Sample and PPI-Oracle estimator, we know that both the sampling density and experimental probability contributes to reduce the MSE. (3) By comparing PGAE with PGAE-Oracle, we observe that the estimation step does not reduce the efficiency of our estimator. On the contrary, it could potentially improve the accuracy. This phenomenon is seen in many semi-parametric estimators (hitomi2008puzzling). (4) Even if the predictor here is not perfect and the model is actually wrongly specified (we do not include $XW$ term in the regression), the estimator is still unbiased and performs relatively well. This indicates how we can harness the power of prediction in the experimentation and estimation.
We also test the algorithm on real world data from the US Census Bureau. In particular, we study the American Community Survey (ACS) Public Use Microdata Sample (PUMS), which is an annual survey that asks about education, income, citizenship, etc. We use the interface in ding2021retiring for access of the dataset, which is also the data source in zrnic2024active.
Here, the research question we care about is the average income of the US family in the year 2019. The target $Y$ here is the income, the covariate $X$ we choose is age and sex, other predictive variable $W$ consists of 15 features includes education, citizenship, race, etc. We conduct semi-synthetic experiments by doing a bootstrap sampling of the survey population. In particular, for each experiment, we will sample $n/10$ data points where $n=221045$ is the total number of data points. We compare the PGAE framework with three other algorithms: (1) Naive: naively sample according to the true distribution and output the sample average. (2) PGAE-No-Pred: sampling according to the optimal distribution derived in PGAE framework but only use actual experimental data for estimation. (3) PPI: sampling according to the true distribution, doing experiments with fixed probability $\gamma$, and construct estimator with the PPI framework (angelopoulos2023prediction). For a fair comparison, we fix the number of actual experimental units to be the same across different methods. That is, if the experimental portion is $\gamma$, then Naive and PGAE-No-Pred will sample $\gamma n/10$ data but PGAE and \textbf{PPI} will still sample $n/10$ data points with average experimental probability of $\gamma$. Here, comparison with \textbf{PGAE-No-Pred} is trying to capture the effect of using prediction in the estimator and \textbf{PPI} is trying to capture the effect of adaptive sampling.
We start with a non-adaptive approach where we use a CatBoost model trained by data from the year before (2018) as the predictor the income. Then we use that model to generate prediction for the income for each individual of the year 2019 and do not update that model. We present the MSE and width and coverage of the 95% confidence interval in Figure (ref). Here the confidence interval is constructed according to the asymptotic variance $V^*=V(p^*, \pi^*)$, where the $(1-\alpha)$-level confidence interval is given by $\hat{\theta} \pm z_{\alpha/2} \cdot \sqrt{{V^*}/{n}}$.
Then we test the case where we continuously update both the estimator $f_t$ and the sampling probability $p_t$ and $\pi_t$. In this case, we do not use any data from previous years and only fit the model using the current sampled data. we update the prediction model in a batch manner where we retrain the model once for every $1000$ samples added. Here, the PPI-Adaptive label means that the predictor is also adaptively estimated using the data of 2019. The result is in Figure (ref).
From the result of both the adaptive and non-adaptive settings, we make the following observations: (1) The PGAE method outperforms all other methods in terms of the average MSE and the width of the confidence interval. The advantage is more obvious when $\gamma$ is small, i.e., when we only have limited access to actual experimental data. (2) Both the optimal sampling and optimal experimental probability contributes to the efficieny of PGAE as it outperforms both PGAE-No-Pred and PPI methods. (3) The PGAE method performs very well even if there's no pre-trained prediction model, which highlights the effectiveness when using the prediction in inference. This can be seen when we compare the performance between the two settings.
In this work, we propose a novel data collection scheme called Prediction-Guided Active Experiment, which leverages an existing predictor to guide the sampling process. This scheme prioritizes sampling data with higher outcome variance and conducts more experiments to obtain actual labels for data where the predictor underperforms. We begin by addressing the non-adaptive case, deriving a semi-parametric efficiency lower bound for any regular estimators, and identifying the optimal sampling distribution and experimental probability by minimizing this bound. We then construct an estimator whose asymptotic variance achieves the theoretical lower bound. Subsequently, we extend our approach to the adaptive case, proposing an estimator that achieves the lower bound even under adaptive sampling. Finally, we validate the performance of our proposed method through simulations and semi-synthetic experiments. We believe our method has strong potential for wide application in active experimentation, given the accessibility and ease of obtaining predictors today.
\setcounter{page}{1}