Near-Exponential Savings for Population Mean Estimation with Active Learning

Thank you for your thoughtful and detailed review. Your feedback highlights several important issues, and we have used your comments to guide our planned revisions and clarifications to the paper. Below, we address each of your points, with the main responses highlighted in bold. After reviewing these clarifications and planned updates, we hope you will consider whether a re-evaluation of the score would be appropriate.

\subsection{Strengths and Weaknesses}

W1: Thank you for pointing this out. We recognize that the hard-margin condition is restrictive in practice. Accordingly, we have extended the main theorem to cover more general cases, and Corollary 2 provides an example where the hard-margin condition is not required. Please see Point 0 and point W2, and point Q3 of our response to Reviewer ZBCc for details.

W2: Thank you for bringing this to our attention. Please see points W4 and Q3 in our response to Reviewer ehtb for our response and planned revisions in response to this feedback.

W3: Thank you for your question about the stratification process and the use of $A^2$. You are correct—our experiments do not include a line corresponding solely to $A^2$, since we use it only as a subroutine within the PartiBandits algorithm, not as an independent baseline. $A^2$ is an active learning algorithm for classification, not mean estimation, so comparing its output directly to mean estimates from PartiBandits or SRS might not be meaningful (see lines 69–71 of the manuscript). We use $A^2$ to identify covariate regions likely to have homogeneous outcomes, which then define strata for adaptive sampling in mean estimation. Although one could construct a mean estimate directly from the classifier’s output, this approach generally yields biased results and is not recommended (see lines 69–71). Based on this feedback, we will better explain this point in the paper.

W4: Thank you for this helpful feedback. We will add Thompson sampling as an additional baseline in our experiments (see lines 242-244), addressing this point directly.

While our experiments primarily compare PartiBandits to simple random sampling (SRS), we agree that it is valuable to consider additional baselines. In most practical settings, SRS and stratified random sampling (StRS) are the standard approaches—and often the only realistic choices available—since the effectiveness of more sophisticated methods is highly domain- and application-specific. In practice, whether alternative sampling algorithms outperform SRS or StRS depends crucially on the relationship between observed covariates and the outcome of interest, which may not always be known or exploitable. That said, we have discussed adaptive sampling methods, such as Thompson sampling and stratified adaptive cluster sampling, as possible additional baselines for future comparison. We will be conducting additional experiments involving Thompson sampling so that it can serve as an additional baseline for comparison. It should be appropriate for use with the AFC data and can also be naturally incorporated into the Monte Carlo experiments we have already conducted.

**W5: Thank you for pointing out these presentation issues. We agree that the notation for \rho_{\leq** and $\rho_{>}$ in Theorem 4 could be confusing. These symbols were not meant to denote special quantities in themselves, but rather to communicate that less than a quarter of labels were flipped on each side of the partition; we will revise the statement to clarify this intent and remove any ambiguity. Regarding $R_1$, it is not explicitly defined in the main text and was included only to help readers locate the corresponding quantity in Aznag et al. (2023) for comparison. As it is not essential for our main discussion, we will either clarify the reference or remove it to avoid confusion.}

\subsection{Questions}

Q1: Thank you for raising this question. We address this in point W1.

**Q2: Thank you for this question. Please see also points W4 and Q3 in our response to Reviewer ehtb. The dependence on $\tau$ is as follows: when $\tau$ is small, we rely more heavily on the Variance-UCB adaptive sampling strategy, which is reasonable if the label budget is large and/or \sigma_{\min** is not too small. When $\tau$ is large, the procedure defaults to simple stratified random sampling, ensuring the stated dependence on the label budget in the theorem but resulting in slightly less optimal constants. In the final bounds, $\tau$ essentially appears as a constant in the denominator, but for sufficiently large label budgets, its effect becomes negligible as the guarantees of the Variance-UCB procedure dominate. In response to this feedback, we will better explain this dependence and trade-off more explicitly in the revised text.}

Q3: Thank you for raising this question. We believe we have addressed this point in our response to W3 above.

Thank you for the detailed rebuttal! The authors managed to make a significant addition to more general cases on the top of hard margin condition and binary setup, explain the role of \tau more explicitly, add the Thompson sampling baseline, and revise the confusing notations. This is impressive and I tend to raise my score to accept.

Q1: Thank you for raising this important question. Extending the near-exponential savings result to non-binary cases would greatly strengthen our contribution. To achieve this, we revised the main theorem and added Corollary 3, as detailed in Point 0 of our response to Reviewer ZBCc.

Q2: Thank you for this question. Active learning algorithms like $A^2$ can be computationally intensive because of their sequential nature. To improve scalability, we revised the PartiBandits pseudocode and main theorem (see Point 0 in our response to Reviewer ZBCc) to support alternative, more computationally efficient active learning subroutines for classification. This flexibility allows us to readily incorporate faster algorithms (whether now or in the future) into the PartiBandits framework for large-scale applications.

Thank you very much for your thoughtful review. Your feedback raises a number of valuable points and questions, and these have helped shape our plans to revise the paper with important clarifications and additions. Below, we provide responses to each point, with the main responses in bold. After reviewing the clarifications and planned revisions below, we hope there will be consideration of whether a re-evaluation of the score is warranted.

\subsection{Strengths, Weaknesses, and Technicalities}

W1: Thank you for highlighting this point. Your feedback led us to write Corollary 2 (see Point 0 in our response to Reviewer ZBCc), which shows how to construct partitions with more than two strata and leverage additional structure in the covariate space.

Corollary 2 covers cases where we can leverage multiple strata, even with binary random variables. Specifically, Algorithm 4.2 of Puchkin and Zhivotovskiy (2022), which forms the basis for Corollary 2, combines standard binary classification with an abstention mechanism. This approach lets the analyst partition the covariate space into three regions: where the learned classifier predicts 1, where it predicts 0, and where it abstains (predicts $(*)$). These correspond to three strata: likely 1’s, likely 0’s, and ambiguous cases (see Point 0 in our response to Reviewer ZBCc). Our framework and Corollary 2 therefore support the kind of multi-strata setup you describe, even in the binary case.

The main guarantee of exponential savings (i.e., the rate in Theorem 5) remains intact, and this approach enables the algorithm to leverage three strata as you suggest. While further adjustments to the main results from Puchkin and Zhivotovskiy may be needed to optimize the final rate, experiments based on this framework are likely to exhibit improved performance, as you note. We will include additional details in both the main text and Appendix to address your feedback on the number of strata and elaborate on these points.

W2: Thank you for raising this point about the assumption that the $P_k$'s are known. Our analysis assumes we know the $P_k$'s and that we have enough unlabeled data to estimate these probabilities accurately. This assumption is standard in the active learning literature (see, for example, Puchkin and Zhivotovskiy (2022) and Hanneke (2011)), which typically presumes access to a large pool of unlabeled data. Your observations about the impact of few versus many partitions are correct. We have revised the paper to clarify this point in response to your feedback.

W3: Thank you for your thoughtful comments on Experiment 2 and the experimental design. You are right that focusing only on the most easily separable classes would create an overly favorable scenario, allowing a classifier to achieve near-perfect separation with few samples. However, Experiment 2 takes a different approach: our outcome is the interaction variable $HB$, not just a class label, which adds complexity. For example, even if a patient is very likely to be Black ($B=1$) based on covariates, they may not have hypertension ($H=1$), so $HB$ is not always $1$ in the upper tail. Because the AFC data are not IID draws from the general population, the upper tail may, for example, include individuals from predominantly Black but affluent neighborhoods who are less likely to have hypertension. Thus, $HB$ may even be $0$ more often than for those classified as $B=0$, so the class separation is not immediate. It would be, however, if the outcome of interest was just $B$, since the covariate represents $\Pr(B = 1)$ so it would be precisely the upper and lower part of the distribution of the outcome, unlike the case when the outcome is $HB$ instead. This setup reflects the reality of many datasets that are not IID samples from the general population. We will clarify this distinction in the paper in response to your feedback.

Additionally, your observations about the figure prompted us to review our code. We found that we had applied the quantile operation in reverse—calculating lower rather than upper quantiles—and had not correctly implemented random sampling for SRS. After correcting these issues, we reran Experiment 2 with the proper sampling and a larger label budget range. The updated figures now reflect the trends you suggested. For Experiment 1, we also ran additional experiments with different noise levels, data generating processes (logit and probit), asymmetric class distributions, and degrees of class separation. Across these expanded results, PartiBandits continues to show consistent performance. We are ready to incorporate these findings into the revised paper.

W4: Thank you for your comments on the intuitive meaning of the bounds, especially regarding the number of partitions in Algorithm 2 and the role of the warmstart parameter in Algorithm 1. We agree these points are important and have added further clarification to the text.

You are correct that the performance of stratified sampling algorithms generally depends on the number of strata. We have updated the PartiBandits pseudocode and main theorem to explicitly accommodate more strata and to show near-exponential rates (see Point 0 in our response to Reviewer ZBCc). For the revised theorem statement, we will clarify the dependence on the number of strata in the main text.

The warmstart in Algorithm 1 serves as a "buffer," ensuring that Variance-UCB remains effective when the overall label budget is small and the algorithm might otherwise undersample low-variance groups. As Aznag et al. (2023) discuss, the regret of Variance-UCB can scale inversely with $\sigma_{\min}$, the smallest group variance. By assigning each group a minimum number of samples, the warmstart helps guarantee the algorithm’s fast-rate properties. For small $N$, the risk bound behaves as $\widetilde{\mathcal{O}}( 1/\tau N)$, where $\tau$ is the warmstart proportion. As $N$ grows and the fast-rate conditions hold, the dependence on $\tau$ vanishes and the sharper rates of Aznag et al. apply. We followed Aznag et al. in focusing only on the dependence on the label budget in the bound, but we will clarify this dependence and its implications in the text of the Theorem.

T1: Thank you for your careful reading of Algorithm 1 and for highlighting the role of the warmstart component. You are correct that the initialization of the standard deviation in line 1 serves as the warmstart. Our procedure does not include an explicit dependence of $t$ on $g$; instead, the warmstart ensures that when the label budget is too small for Variance-UCB’s rate guarantees, we fall back to uniform (StRS) sampling (see Point W4 above). We have revised the text to clarify this point.

T2: Thank you for drawing attention to these points regarding notation. Most of the notation is standard, but we have clarified some of it in the revisd text.

1. Yes, $\widetilde{\mathcal{O}}$ ignores logarithmic factors and fixed constants.
2. We did not explicitly define $R_1$ in the main text; we referenced it only to help readers compare with Aznag et al. (2023). Since $R_1$ is not essential to our main discussion, we will either clarify its reference or remove it to prevent confusion.
3. $c_1$ and $c_2$ are simply the subgaussianity constants (see Aznag et. al (2023)), but we will clarify them in the revised paper.
4. $\Sigma_1(\mathcal{G})$ is described in English in lines 72–73 (as the average within-group variance of the outcome given stratification $\mathcal{G}$), and is defined formally in the Appendix. We agree it would be helpful to include the formal definition directly in the main text, and we will do so.

T3: Thank you for raising this important point about the clarity of assumptions in our theorems. You are correct that Theorem 3 does not rely on Assumptions 1 and 2 and holds for a much broader class of random variables than just binary $Y$. We will revise the paper to state the assumptions for each theorem explicitly and clarify that Theorem 3 applies in this greater generality.

\subsection{Questions}

Q1: Please see our response to point W3 above.

Q2: Thank you for this question about how Algorithm 2 scales as the number of groups $g$ increases. We have addressed the theory behind this question in Corollary 3 (see Point 0 in our response to Reviewer ZBCc).

Q3: **Thank you for highlighting the significance of $\tau$ in Algorithm 1. The parameter $\tau$ safeguards the Variance-UCB procedure by ensuring that part of the label budget is allocated to simple stratified random sampling (StRS), which maintains theoretical guarantees even when the label budget or \sigma_{\min** (the smallest within-group variance) is small. Setting $\tau = 0$ allocates all labels via Variance-UCB, achieving optimal rates when both $\sigma_{\min}$ and the label budget are sufficiently large. However, a positive $\tau$ ensures we sample from every group, preserving the guarantee of Theorem 3 with only slightly less optimal constants. For large label budgets, the effect of $\tau$ disappears and the optimal rates of Aznag et al. (2023) apply. Since the rate’s dependence on the label budget remains unchanged (only the constants differ), we did not state this dependence explicitly in the theorems. We will clarify these points and add more detail to the theorems in response to your question.}

We thank you very much for your careful review. It raises many good points and questions that have informed key revisions we have made to the paper. In what follows, we provide responses to each of the issues, with the main points in bold. However, there is one point we would like to discuss first as it addresses many of the issues raised. Once the clarifications and plans for revision below have been reviewed, we kindly ask the reviewer to consider if a re-evaluation of the score may be possible.

Point 0: In response to feedback about the paper’s restrictive assumptions (e.g., Tsybakov noise, Bayes classifier in class) and binary problem setup, we revised Step 2 of the PartiBandits algorithm and revised our main result of near exponential savings in Theorem 5. These limitations arose simply because we used the classical $A^2$ algorithm in Step 2 for clarity and intuition. As noted in lines 234–245, one can substitute less restrictive subroutines as needed. To make this explicit, we updated the PartiBandits pseudocode and reformulated the main theorem to accommodate broader assumptions and setups. Only minor adjustments to the proof are required, as detailed below.

The revised theorem allows us to incorporate any algorithm for classification into the PartiBandits framework through Step 2, provided it achieves exponential savings. PartiBandits directly inherits the assumptions and problem setups of any such algorithm, ensuring it remains effective for efficient mean estimation as active learning methods continue to advance.

What we changed. In Step 2 of PartiBandits, we now allow any algorithm, $\mathcal{S}$, that learns a binary or multiclass classifier $\widehat{f}$ satisfying, with high probability, $\mathbb{E}[(\widehat{f} - y)^2] \lesssim \nu + \exp\left(c \cdot \left(-\frac{N}{\log(N)}\right)\right)$, where $\nu$ is the infimum risk over the hypothesis class, $c$ is a constant, and $N$ is the label budget ("Condition 1"). With this modification, along with minor updates to the rest of the PartiBandits pseudocode, we reformulated the main theorem to allow more flexibility in assumptions and problem setups. The revised theorem is as follows.

$***$

Theorem 5: Suppose that for some joint distribution of $(X,Y)$ where $Y$ is a $k$-class random variable, there is a subroutine, $\mathcal{S}$, that learns a $k$-valued classifier $\widehat{f}$ such that with high probability, $\mathbb{E}[(\widehat{f} - y)^2] \lesssim \nu + \exp\left(c \cdot \left(-\frac{N}{\log(N)}\right)\right)$ where $\nu$ is the infimum risk of all possible classifiers in the hypothesis class, $c$ is a constant and $N$ is the label budget. Then, if $\mathcal{S}$ is used as the subroutine in Step 2 of PartiBandits (Algoirthm 2), we have:

$$
\left| \widehat{\mu}_{\text{PB}} - \mathbb{E}[Y] \right|^2 = \tilde{\mathcal{O}}\left( \frac{\nu + \exp(c \cdot (-N/\log(N))) }{N} \right),
$$

where $c >0$ is a constant. 

$***$

Why this change is possible. We can make this change because the only requirement for Step 2 is to identify strata that are likely to be homogeneous with respect to the outcome, a task that any well-performing classification algorithm can achieve. In the original version of PartiBandits, we used the $A^2$ subroutine in Step 2 to identify homogeneous regions of the covariate space, but any suitable classification algorithm can fill this role. Intuitively, a classifier with low excess risk tends to group together samples with similar labels, which helps us define more homogeneous strata. This structure enables us to obtain stable mean estimates within each stratum and efficiently aggregate them to efficiently estimate the overall population mean.

We can formalize this argument as follows. For binary $y$, Theorem 5’s rate holds as long as the classifier used in Step 2 of PartiBandits satisfies Condition 1 above. We show in the proof of Theorem 5 (see Appendix) that the average within-stratum variance is bounded above by $\mathbb{E}[(f(x) - y)^2]$ (where we change notation so that $f$ in this paragraph is now the learned classifier). To see this, we first apply the law of total expectation to write $\mathbb{E}[(f - y)^2] = \mathbb{E}[(f - y)^2 \mid f = 1] P(f = 1) + \mathbb{E}[(f - y)^2 \mid f = 0] P(f = 0)$. Applying the bias-variance decomposition, we obtain $\mathbb{E}[(f(X)-Y)^2] \geq \Var(Y \mid f=1) \Pr(f = 1) + \Var(Y \mid f=0) \Pr(f = 0)$. This right-hand side is exactly the average within-group variance that determines the high-probability bound on the estimation error for the PartiBandits population mean. This derivation is purely algebraic and does not rely on any assumptions specific to the active learning algorithm used in Step 2. The same argument extends directly to multiclass $y$. We now outline a few corollaries that illustrate how PartiBandits adapts to different assumptions and setups through the choice of Step 2 subroutine.

Corollary 1: Suppose $Y$ is binary and the Bayes optimal classifier is in the class and the hard margin condition is satisfied. Then the rate of Theorem 5 holds with $\mathcal{S}$ being $A^2$ (this was the original result of our main Theorem).

Corollary 2: Suppose $Y$ is binary and the hypothesis class has bounded combinatorial diameter and a finite star number. Then the rate of Theorem 5 holds with $\mathcal{S}$ being Algorithm 4.2 of Puchkin and Zhivotovskiy (2022).

Corollary 3: Suppose $Y$ is a $k$-class random variable and the covariate space is such that the conditions necessary for Corollary 2 in Agarwal (2013) hold. Then, an analog of the near exponential savings rate of Theorem 5 holds with $\mathcal{S}$ being Algorithm 1 of Agarwal (2013).

\subsection{Strengths and Weaknesses}

W1: We have added Corollary 3 to address this point (see Point 0).

W2: We have added Corollary 2 to address this point.

As you correctly note in Question 3, the algorithm of Puchkin and Zhivotovskiy (2022) achieves exponential savings under weaker assumptions. As we explain in Point 0, using this algorithm as the subroutine $\mathcal{S}$ in Step 2 of PartiBandits allows PartiBandits to inherit these more flexible assumptions, so an analogous near-exponential savings result follows.

W3: We have added Corollary 2 to address this point..

We discuss this point more in point W1 in our response to Reviwer ehtb.

W4: It is true that our proofs rely on existing active learning results, but this is a feature, not a bug. We designed the proof of Theorem 5 to apply to any active learning classification algorithm used in Step 2 of PartiBandits, as long as it reduces excess risk with exponential savings (Condition 1). This is thanks to an original decomposition we derive that bridges classification tasks and mean estimation tasks, as explained above. As active learning is a rapidly developing field with more advanced algorithms being developed every year, we wanted to ensure all such algorithms for classification, either now or in the future, can be incorporated into the PartiBandits framework (by way of Step 2) to help with efficient mean estimation. In response to this feedback, we have revised the Pseudocode of PartiBandits and the statement of our main Theorem proving the near exponential rates to make this contribution clearer.

\subsection{Questions}

Q1: Thank you for raising this question. We can extend the PartiBandits framework to general real-valued outcomes by discretizing $Y$ and applying Corollary 3, which addresses the feedback about the non-binary case. Achieving similar results without discretization is substantially more difficult and would likely require new advances in active learning for "classifying" continuous outcomes. To our knowledge, no one has yet shown exponential savings in the continuous outcome setup without discretization, for classification, estimation, or otherwise. Success on that front would itself merit a separate paper. We will revise the manuscript to explain these points.

The idea behind the discreitzation approach is to partition the range of $Y$ into $k$ intervals and treat each as a separate class; for example, mapping $Y \in [0,1]$ into $k=10$ bins such as $[0,0.1), [0.1,0.2), \ldots, [0.9,1]$.

Q2: Thank you for raising this question. Corollary 2 (see Point 0) has been written to directly address this, demonstrating that such an improvement is indeed possible.

Please see our response to point W3 and point W1 in our response to Reviwer ehtb for further details.

Q3: Thank you for raising this question. This understanding of the work of Puchkin and Zhivotovskiy (2022) is correct, and yes, Assumption 2 may be replaced by similar requirements in their setup. We have written Corollary 2 to address this feedback (see Point 0).

See also our response to point W2 above.

Q4: Thank you for this helpful suggestion—we will add Thompson sampling as an additional baseline in our experiments (see lines 242-244), addressing this point directly.

We address this in section W4 in our response to Reviewer MZfC.