Near-Exponential Savings for Population Mean Estimation with Active Learning
Authors: Julian Morimoto
JACOB GOLDIN
Daniel E. Ho
Venue: NeurIPS 2025 poster ()
Keywords: Adaptive Sampling
Variance-UCB
Disagreement Based Active Learning
Primary area: general_machine_learning

Abstract: We study the problem of efficiently estimating the mean of a $k$-class random variable, $Y$, using a limited number of labels, $N$, in settings where the analyst has access to auxiliary information (i.e.: covariates) $X$ that may be informative about $Y$. We propose an active learning algorithm ("PartiBandits") to estimate $\mathbb{E}[Y]$. The algorithm  yields an estimate, $\widehat{\mu}_{\text{PB}}$, such that $\left( \widehat{\mu}_{\text{PB}} - \mathbb{E}[Y]\right)^2$ is $\tilde{\mathcal{O}}\left( \frac{\nu + \exp(c \cdot (-N/\log(N))) }{N} \right)$, where $c > 0$ is a constant and $\nu$ is the risk of the Bayes-optimal classifier. PartiBandits is essentially a two-stage algorithm. In the first stage, it learns a partition of the unlabeled data that shrinks the average conditional variance of $Y$. In the second stage it uses a UCB-style subroutine ("WarmStart-UCB") to request labels from each stratum round-by-round. Both the main algorithm's and the subroutine's convergence rates are minimax optimal in classical settings. PartiBandits bridges the UCB and disagreement-based approaches to active learning despite these two approaches being designed to tackle very different tasks. We illustrate our methods through simulation using nationwide electronic health records. Our methods can be implemented using the \textbf{PartiBandits} package in R.

Near-Exponential Savings for Mean Estimation with Active Learning
Julian M. Morimoto" Jimorinoto@berkell.. edu
Jacob Goldin jagoldin@uchicago.o edu
Daniel E. Ho dho@law. www.w..taafeee stanford.ce edu
Abstract
We study the problem of efficiently estimating the mean of k-class random variable, Y. using a limited number of labels, N. in settings where the analyst has access auxiliary information (i.e.: covariates) X that may be informa- tive about Y. We propose an active learning algorithm ("PartiBandits") to esti- mate E[Y] The algorithm yields an estimate, (ApB, such that (Ørn E[Y])* is Ô N
constant and v the risk the Bayes-optimal classifier. PartiBandits is essentially a two-stage algorithm. In the first stage, it learns partition of the unlabeled data that shrinks the average conditional variance of Y. In the second stage it uses UCB-style subroutine ("WarmStart-UCh UCCB") to request labels from each stratum cound-by-roum Both the main algorithm's 's and the subroutine's convergence rates minimax optimal in classical settings. PartiBandis bridges the UCB and disagreement-basee proaches active learning despite these two approaches being designed to tackle very different tasks. We illustrate our methods through simulation using nationwide electronic health records. Our methods can be implemented using the PartiBandits package in R.
where >
ap-
Introduction
Estimating the mean of a k-class random variable, Y, with limited data from a subset the population of interest is pervasive problem in statistics and machine learning. A classical solution to this problem to draw a simple random sample (SRS) N independent and identically distributed (IID) labels and compute the resulting sample mean. However, this may be inefficient use of the label budget if one has information X (i.e., covariates) that may be related to Y. In such cases, one approach is to leverage X to get better estimate of E with fewer labels, perhaps through stratified random sampling (StRS) over X and allocating the label budget across strata in proportion to how frequently each stratum occurs the population. But in practice, there are many ways to define strata, and choosing a poor definition can result in minimal gains, or even worse performance than SRS. In general, analysts who use X poorly, through stratification or otherwise, may over-sample some subpopulations and neglect others, resulting in biased or sub-optimally noisy estimates (see, e.g., Aznag et al. (2023): : Henderson al. (2022)) ). This challenge has motivated the development of different adaptive sampling techniques for mean estimation (sec, c.g., Seber and Mohammad Salchi (2015); Thompson (1911) )) these approaches focus on asymptotic performance and do not address whether fast rates of convergence can be achieved in finite samples. In parallel, the active learning literature has developed strategies for learning with limited labels. While classical active learning
DDepartment of Statistics, University of California, Berkeley; Regulation, Evaluation, and Governance Lab, Stanford Law School; World Bank Group University of Chicago: American Bar Foundation Stanford University
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

results primarily focus on classification (see, e.g., Puchkin and Zhivotovskiy (2022) Hanneke and Yang (2014); [Hanneke (20TT), recent work uses active learning to efficienlly estimate subgroup (i.e. within-strata) means in settings where strata are predefined (Aznag et al.) 2023). However, there has been no thorough exploration of active learning methods for population mean estimation when the researcher does not know an optimal stratification scheme. In this paper, we carry out such an exploration by developing an active learning framework for population mean estimation of k-class random variables, its convergence guarantees, and to what extent fast rates of convergence are achievable.
Our main problem setup revolves around estimating the mean of a k-class random variable Y. where the analyst has access to auxiliary information X that may be informative about Y. The analyst may adaptively choose which instances to query for their corresponding labels, Y. round by round, with a budget of N label requests. This setup parallels the pool-based active learning setup, where the analyst observes large collection of IID unlabeled instances X1, X2,... and sequentially selects which ones to label, ultimately giving the analyst the labeled dataset, (X1, Y1), ..... (XN, Yv). The hope is to obtain an estimate of E[{ that is closer to [[[] than the SRS strategies with high probability,
where the latter convergence rates are on the order of O Var(()
There are two important reasons why it is hard efficiently estimate population means in this problem setup. An ideal strategy would first partition the data into strata that minimize the average within-stratum variance of Y, then allocate the label budget across these strata according to the Neyman allocation to minimize the variance of the mean estimate that aggregates the subgroup mean estimates (see Jo et al. 2025); Bosch al. (003)). But in most settings, this optimal stratification is not known ahead of time. Moreover, variance within each stratum is not observed directly and must be estimated from noisy samples, so an allocation strategy that may seem optimal early on na--baed on preliminary variance estimates may prove suboptimal as more data is collected. Thus, the analyst must (1) learn a good stratification from unlabeled data, and (2) decide how to allocate labels across strata adaptively in way that reflects estimated (as opposed to oracle) variances.
1.1 Summary of Contributions
Our contributions are five-fold. First, we develop an active learning algorithm ("PartiBandits") for efficiently estimating the mean of a k-class variable Y. This algorithm yields an estimate, PPp, such that (7pB- ) is Ö 1+texp(e=[--M (rrxx((((-NNooo(NN))) where e > O is a constant and v is the risk of the
N
Bayes-optimal classifier (Theorem 3 Figure I.. It performs least as well as SRS in N, and almost exponentially better when X is predictive of Y (i.c.. when // is small). It also closely resembles the exponential savings observed disagreement-
active learning for classification (Puchkin and Zhivotovskiy 2022, Hanneke and Yang, 2014 Hanneke| [)11)) even though such results do not help with the task of mean estimation Y when X does not perfectly predict Y (see Dong et al. (2025)). ). Second, we show that if X can be stratified in advance using a stratification scheme G, the PartiBandits subroutine ("WamsStartUUCB" achieves error Ô
where (g) is the average -uthin-group variance of Y (Theorem 1 Figure 1. Third, we show that both convergence rates are minimax optimal in classical settings (Theorems 2 and 4p. Fourth, we bridge gap between Upper Confidence Bound (UCB) algorithms and disagreemment-base approaches in the active learning literature despite these two approaches being developed for very different tasks (Section 4.2). Finally, we conduct simulation studies using real-world data from over 6 million electronic health records and find that the gains predicted by our theory for population mean estimation can be achieved even in realistic small- sample regimes (Section 5p..
(g) N
2 Related Work
Our work builds on, and bridges, two different strands of prior work in active learning. The first is disagreement-basee theory, as developed and refined by Hanneke designed for classification where labeled data are costly but unlabeled data are abundant. In this setting, the analyst queries labels for instances drawn from large pool, concentrating effort on regions of the input space where candidate hypotheses disagree. defining feature of this approach is its potential for "cxponential savings", which refers to the convergence rate of excess classification error
This theory was originally

---

0.5
044
Estimetion Enor
Estimetion Error
PartBandits (Main Algorthm) High
WarmStart-UCe (Subroutine) High I:(o)
03
0.
X,(G)
0.1
Number Labels (N)
Number Labels (N)
This plot compares the performance of PantiBandits and WarmStart-UCH UCB, to SRS in different problem settings. The left panel compares SRS to PartiBandits for label budgets from 10 100. Here, X ~ Unif(o, 1] and Y { 2 0.5). with a fixed fraction Y's (between 0% and 10%) randomly flipped to introduce noise. The proportion of flipped labels is equal to v by definition. For each label budget, we generate 500 hypothetical datasets in this way. apply SRS and PartiBandits each, and compute the resulting error rates. We then take the 90th percentile of these
Figure 1:
rates obtain classical 90% high- probability/coniidnee bound. PartiBandits eventually outperforms SRS with relatively fewer samples, with performance gains becoming more pronounced when better predicts Y and y decreases. The right panel compares SRS to WarmStart- from 50 to 200. In this panel, X ~ Unif 1] and Y [X 2 D.5}. with 5% of the labels randomly flipped to introduce noise. We examine the effect of specifying different stratification schemes beforchand that reduce the within-group variance of Y to varying degrees, where lower values En (g) indicate better average within-groo group variance reduction. Each scheme defines strata by applying a threshold between 0.3 and 0.5 and grouping observations based on whether X falls to the left or right of the threshold. We run the same simulation procedure as above to obtain the 90% confidence bounds. WarmStart-UCH consistently outperforms SRS, and the gap grows when stratification reduces variance more effectively (i.c., when Σ, (9) shrinks).
for label budgets
(relative to the risk of the Bayes optimal classifier) shrinking at roughly Ô (exp(c- (-N/log(NV) (N)()))))) far faster than the (VVr((Y)/N) rate typical in passive learning. While disagreement-based learning has been extensively studied in classification, has not, problem of estimating population means. Our work is the first to show that the core insights of this framework can be used construct stratification schemes that substantially reduce estimation variance in the mean estimation setting.
knowledge, been applied to the
The second strand that our work connects to is UCB-style active learning. In particular, our proposed WarmStart-UCH subroutine is closely related to the work of Azzag et al. (2023), which developed a Variance- UCB algorithm for estimating the means of predefined subgroups using fixed label budget, using upper confidence bounds on within-grous group variance to guide sampling. Our subroutine uses a similar approach to estimate the overall population mean using the strata selected by the first stage of our algorithm. Additionally, we build on the Aznag et al.) (2023) results by showing in Theorem [ that the rate of subroutine for estimating population means from pre-defined strata explicitly quantifies the effect of how "informative" the subgroups are for estimating the quantity of interest, something that cannot be obtained through direct application of the main Aznag et al. results alone. Our rate also has improved dependence on key parameters such as the number of strata and amin the smallest conditional variance of Y over all strata.
3 Notation and Problem Setups
3.1
Main Setup
Our main problem setup is that of estimating the population mean of a k-class random variable, Y, whose realizations come from the set (0....,k) ...,, using limited label budget N. Where appropriate, the realizations of Y may also be any set with A distinct elements in JR. The analyst has abundant access unlabeled information X E (ex: covariates), which may be informative about Y, and can choose which examples to label in order to estimate μ E The analyst uses algorithm to
3

---

construct an estimator ( of the population mean that uses auxiliary information X' and only N labels. The goal is to minimize the squared error, (i(NV- =n)
We also have the following terms and notations. A hypothesis class C is any set of measurable classifiers h : X' + R C E where [R] is finite. As we will show in Section[4.2 PartiBandits parallels classical disagreement- based active learning algorithms in that requires as an input. For any measurable h X + {O....,k) k). we define the squared loss of h as er(h) (X( - Y) Let inc er(h), the infimum loss of C.
As discussed in the Appendix, our main results still hold for other loss functions, including asym- metric misclassification costs. Such alternatives can produce more informative Bayes classifiers in applications where the ordinary squared-loss version fails to identify the threshold (such as when Pr(Y- = 1 X) < 1/2 for all X).
Our main assumption is the following:
Assumption 1 (Exponential Savings in Classification). We assume that the joint of distribution (X, Y) and the hypothesis class are such that an active learning algorithm, S. can be used to learn a classifier h such that with high probability, B[[M(X)- - Y) - 2 exp (c- log(N)) where N is the label budget, and e > 0 is some N-independent constant.
There are many problem setups in which this assumption is satisfied, as we discuss in Corollaries
3.2 Setup for a PartiBandits Subroutine, WarmStart-UCB
Additionally, PartiBandits contains a subroutine that depends on the following problem setup that builds on the one discussed above. The following notation and definitions are drawn from Aznag et al.) (2023). We assume that can partition X' using stratification scheme G = {AI.... Aa where A, S Ed are disjoint. Let Pa = P(X- E A₂) and M4 E[y X A, E A,). so the population mean is μ Zg=1 Mg. We define of == Var(P, Y | X Ag), which equals p2 oa2, where o² :- Var(Y X E A,) is the unweighted conditional variance of Y given X E Ay. The distribution of X is assumed to be virtually known (as is the case in classical active learning setups), so Pg is also known to the analyst. (g is the average within-group variance of Y, The analyst wishes to compute an unbiased estimate of the population mean with N label requests, sampling only one group from {Ai...., Act } at time. The set of feasible policies for estimating is defined as II := { {TTe [re | ne E G²-1 x R'-1 + (G)), Vt E [N]], where AG) is the set of measures supported on [G]. For some policy æ E II, let (g. denote the number of collected samples from group Ay after N label requests by way of policy æ e II, and let (T)) be the weighted sample mean estimator of (ig for TT))(T) collected samples, that is: (TT)(T)) == (G. Et:TEA, Ye Pg. Once all data have been collected using the full label budget, N, the analyst will compute the population mean by aggregating the subgroup mean estimates obtained from the policy π E II: fTTT, N) = Eg (T)). Going forward, we drop the explicit dependence on π and N in the notation when the policy is clear from context. As long as at least one sample is collected from each group, each Ng.N is an unbiased estimator of Pg TE(Y X E Ag] Pg- and hence the aggregated mean estimator ju is an unbiased estimator of the population mean j.. The analyst wishes to obtain a high probability bound on the variance of fi. The analyst does not know the true standard deviation vector o := (oI....,,o OG), which can be used to obtain an upper bound on the variance of få via Var(ā) = Var [Egn =1 (g.N) S Eg o²
Σye(a o P.
Therefore the analyst must learn through their decisions. We define the regret of a policy Regret (=) (Ā(T, N) 10=-)) n)²
"g.N"
4 Our Algorithms and Performance Guarantees
We now discuss our algorithms and their performance guarantees in turn. The proofs are in the Appendix. PartiBandits is our main algorithm, but since incorporates UCB-style subroutine, WarmStart-UCB, we first analyze the subroutine.

---

4.1 PartiBandits Subroutine: WarmStart-UCB
The first Algorithm is similar to the Variance-UCB algorithm of Azzag et al. (2023) except that we include an initial "warm-start" step (Step 1). We estimate Og via the sample standard deviation,
002 = / ng. Eosux,= (PyY. - pg,,)² We can then define UCB
0gg + Cx(8)
Vagut
where (6) := / 20g log (3) log (7) + 2 log()) +1 ()))
N In ON 0i and C2
(1-61/21og( log(()))
constant upper bounds the sub-gaussiaa parameters of Y (which exist since Y is k-class), and 8 E (0, 1) is parameter representing the confidence level for obtaining a high probability bound. WarmStart-UCB selects each round the group with the largest upper confidence bound its variance estimate, but begins with a "warm-start" phase that allocates fixed fraction of the label budget, 7, evenly across all groups (initiated by Step 1).
Algorithm 1 WarmStart-UCB
Require: Label budget N, stratification scheme G. confidence level 6, buffer fraction 1: Initialize 11.00 = 0, and 0g = too for all y E [G] and / S N
2: Compute (5)
3: for t = .,, 1 do
4: Compute (o,) 0g.t + Cx.
Vg E [G]
Vng.s
(ag)
5: Select group X,+1 arg maxy
ng.t
6: Observe feedback Y
7: Update the number of samples: Ng,t+1 ng,t + 1 1
A [G] PaYs, [
8: Update the mean estimates, Ag,t+ Mgget1
Ex=1 1x 1x,=g
9: Update the standard deviation estimates,
gg.t+1 M1.41-1 (Pg #get1)3 e
10: end for Output: 7ws. = -Eg
The following an upper bound on the performance WarmStart- UCB.
Theorem 1. [uws-cc || = Ō (2)10)
(g(
Theorem shows that when stratification scheme given a priori, WarmStart-UCB efficiently estimates the population mean with error scaling Ô (g) where Mn (9) captures how informa-
tive the grouping is. The more informative the grouping, the smaller (g) is, and the faster the rate of convergence. By the law of total variance, Σ, (9) S Var(Y), so this rate always at least as fast as that obtained with SRS. The proof is relatively straightforward. In Section A.5 of their work, [Aznag] al. (023) showed that (n) (nn)) (1/N). We do not define R₁ (n) and R3 Rn(n)) explicitly (n)))
here as this would involve significant technical detail. However, we show in the Appendix that R₂ (n) equivalent the variance of fwws.UB. while Rí (n*) corresponds Σ, (9)/N. This identification allows us to directly leverage this bound from Section A.5 of Aznag et al. (023). We then calculate how large N must be in order for this quotient to be bounded from above by some constant (though this threshold is quite large, as it is inversely proportional to min)) and we get a bound on the variance of e1 for sufficiently large N. For all other N, we use the fact that minimum fraction of the label budget is allocated to each group and obtain a similar high probability bound using classical Hoeffding arguments. This is why and the WarmStart step are important, as they safeguard the Variance- UCB procedure by ensuring that part of the label budget is allocated to StRS (every group gets minimum number of samples), and this allows for nice convergence guarantees even when the proper rate of Azzag et al. (2023) does not hold. This allows us to obtain an analogous rate for label budgets that do not meet the threshold, thereby eliminating the counterintuitive dependence on amin that is typical in the active learning literature (Aznag et al. 2023 Carpentier et al. 2015). This ensures that our rate holds uniformly over label budgets and constitutes a proper non-asympootic, high-pro high-probabilltyy bound.
5

---

We note in the Appendix that when the dependence on the constants and G is made explicit, the rate is /AWS-LCB - = G.E, constant, and do the same for т. As allude to above and show the Appendix, the dependence on T vanishes for large N relative to min, and for all other N, the rate still holds with slightly larger constants (including constant factor of 11). We also discuss in the Appendix the special cases when T E {0,
N-T however, we follow Aznag et al. (2023) in treating as
While Aznag et al. )202)) established a rate of ((--2) for the task of multi-group mean estimation (which can extended to the task of population mean estimation) for a particular regret definition, the upper bound for WarmStart-UCh both (I) explicitly accounts for the signal of Y in X through Σ, (9). and (2) exhibits tighter dependence on the number of groups, G, and eliminates the dependence on the smallest conditional variance of Y across all subgroups, min- The latter result in particular addresses an open problem in the active learning literature on mean estimation by demonstrating that not all active learning mean estimation frameworks result in the counterintuitive inverse dependence on min (sec, e.g., Aznag et al. (2023): Carpentier et al. 20155: Ganti and Gray (2133). While these improvements to the rate of convergence come the cost of slower dependence the label budget
N (from N 2-1), this is expected as the ONN ((N-2) of Aznag al. (0033) is for different definition of regret than the one we are interested in here, (îi p))
The following provides a matching lower bound.
Theorem (Lower Bound for WarmStart- UCBB. Let X Unif(0, 1] and Y = 1 2 t for some t E [0, 1]. Assume that Pr-fraction of labels of is flipped at random over X < t. and analogously with p> for X and that ps, P> 1/4. 4. The stratification scheme G partitions the covariate space at the true threshold (i.e., groups X < and X 2 t). Then,
7 - E(Y] 2 (g)
N
for some constant Cj > O and all estimators It of E[Y] in this setup.
Since this lower bound matches the upper bound of Theorem [ we have that the rate of Theorem iss minimax optimal in this classical setting. This lower bound is based on the classical threshold example where the stratification scheme is such that the strata are chosen according to the decision boundary, and represents a favorable case where X highly predictive of Y and the analyst has knowledge of how to group observations to reduce within-stratuu variance exactly the kind of setting any analyst would hope to operate in. The main point to note about this lower bound that when the stratification scheme is well-chosen, the dependence in N is still on the order of 1 1/N. This will be important in the discussion of the lower bound for the main PartiBandits algorithm (Theorem
4.2 PartiBandits
We now present our main algorithm, PartiBandits (Algorithm) 2
Algorithm 2 ParttiBandits
Require: hypothesis class C, active learning classification algorithm S, label budget N, confidence
level 8, buffer fraction т.
1: Stage 1: Learn stratification using S
2: Run S with hypothesis class C, label budget [N/2], and confidence level 5 to obtain classifier h
3: Construct stratification scheme G by defining A; 1i-1(i) for all E Im(h) and setting
g {A;};
4: Stage 2: Apply Stratified Sampling Subroutine (WarmStart-UCB) to estimate means over G
5: Run WarmStart-UCB with label budget N - [N/2], stratification scheme & and buffer fraction т Output: Tip Eg Ng.N.
essentially two-stage algorithm. In the first stage, runs a disagreement- algorithm, S, that the analyst chooses. S helps identify partition of the unlabeled data that shrinks the average

---

conditional variance of Y. In the second stage, it runs the WarmStart-UCE UCBB subroutine on that learned stratification. Examples of handle the case when Y binary (A+ 1) include the A² algorithm of [Balcan et al. (2006) and Algorithm 1 of Puchkin and Zhivotovskiy (2022). For the multiclass setting (A= >> 1), one may instead use algorithms such as Algorithm of Agarwal (2013) to learn a partition of the unlabeled data reduces the mean conditional variance of Y. In Algorithm 2 we present PartiBandits with a general choice of S, and show in Theorem 3 that it can achieve near-exponentin savings whenever Assumption is satisfied given the gatta-enerating process, hypothesis class, and the choice of S. We then illustrate in Corollaries [ how different choices of S can accommodate different gett--enerating processes (e.g., binary vs. multiclass Y) and assumptions about the problem setup (such as the hard margin condition or the assumption that the Bayes optimal classifier is in the hypothesis class). The main theorem and its corollaries show that PartiBandits allows efficient mean estimation for multiclass outcomes across a wide range of structural assumptions and problem settings.
(2013)
The following an upper bound on the performance of Algorithm 2
Theorem 3. For any joint distribution of (X, Y), hypothesis class C, and S such that Assumption 7 holds, we have
UPB - -E[Y] (2 (-N/log(N)))
N
where O is constans.
We note in the Appendix that when the dependence on T and G is made explicit, the rate is
1PPB - B[y] ò g) [=texplee-NN/ log( (og(N))) where [g] the number of strata. Theorem N-7
shows that PartiBandits efficiently estimates [E by learning a stratification scheme G of Y that yields an average within- -tratum variance, (g)),, that is bounded from above by 1/ exp(ec-p---- log ))).
Asymptotically, this rate is faster than-or least as fast as the 1/N decay achieved by classical adaptive sampling methods (Félix-Medina 2003 Thompson 1991p. since our bound decays at the rate of roughly exp(-eXX log N)/N when y is small.
Result Intuition. Disagreement-baseed active learning algorithms effectively learn a stratification scheme where within-group variance is reduced, since the labels in each stratum (i.e., strata induced by the inverse mapping of the classifier's prediction function) will tend to concentrate around a single class. Furthermore, they can do this with very few labels. Because active learning algorithms can identify these low-variance strata rapidly, we can then perform an adaptive stratified sampling procedure (Algorithm IT] using the learned stratification to estimate the population mean. Since estimation error depends primarily on the average within-group variance, reducing that variance quickly leads to a correspondingly fast decline in estimation error. We discuss in Corollary 4 how PartiBandits can further decompose relatively homogenous strata into sub-strata with higher and lower conditional variances, which allows the algorithm to allocate more samples to more heterogeneous sub-strata, yielding even better estimates of the population mean.
Proof Sketch. If Assumption is satisfied, then there is active learning algorithm, S, such that when used in Step 2 of Pantis, we obtain a classifier, hi whose excess risk decays at an exponential rate, (X) - 2 - v S exp(c (-N/ log N)) for some constant c > 0. We can then show that the variance of the mean estimate, Çinn, is bounded from above by E (ht ( - Y) In particular, we first use the law of total expectation show that B[([((X)) - -Y) ΣSE, (((j-Y)²
R(X)= = j] Pr(((x)= j). where J is the image of h. Then use the Bias Variance decomposition to show that the latter quantity is equal to Eee (Var(Y h = 5) + (j E h = )1)²) Pr = j). Then easily follows that this quantity is an upper bound on the average within- group variance of Y using the stratification induced by h, and therefore an upper bound on the variance (and therefore the estimation error) of fipg-
The constant e depends on C's VC-dimension and disagreement coefficient (Hanneke 2011) Part- iBandits yields better mean estimates with smaller label budgets if C is well constructed and relatively small (as the case when the analyst has good prior knowlede knowledge about possible ways in which X may be related to Y), since less of the label budget is needed to eliminate incorrect hypotheses. It is typical for disagreement-bant
active learning algorithms to exhibit this dependence on the hypothesis class C (Puchkin and Zhivotovskiy 2022 Hanneke and Yang 2014).
7

---

With different choices of S. we can obtain the following corollaries that allow for different assump- tions and problem setups regarding the joint distribution of (X; Y) and the hypothesis class C. All proofs may be found in the Appendix.
Corollary 1 (Classical Binary case with low noise). Suppose is binary, and the joint distribution of (X, Y) and hypothesis class C are such that there exists μt < co such that for all > O diampe; C) S ue, where diam(E;C) C) is the diameter of the E-minimal set of (this is the "hard margin" condition. For further details, see Section 2 and Theorem 4 of Hanneke (2011)). If we set S to be the A algorithm of Balcan et al. (2006). then, given a label budget of N and E (0, 1/2), we have with
probability at least 8 that 1iPB E[ =
['texple--- log N1)
N
Corollary 2 (Binary, weaker structural conditions on C). Assume Y E {0,1] and that the joint distribution of (X, Y) and hypothesis class C are such that Massart's noise condition is satisfied (Assumption in Puchkin and Zhivotovskiy (2022)). without requiring that contain the Bayes optimal classifier. Surposefurther that the joint distribution and hypothesis class are such that the star number s and the (combinatorial) diameter of C are finite (see Section 2 and Theorem 4.1 of Puchkin and Zhivotovskiy (2022)) If we set S to be Algorithm 4.2 of Puchkin and Zhivotovskiy (2022) then, given a label budget of N and 8 E (0, 1/2), we have with probability at least 8 that 1ip8 B[y]]= =0 Ö "-toxppee( N/ log NY)
NV
Corollary 3 (Multiclass). Suppose Y is k-class (k 2) and that the joint distribution of (X, Y) and hypothesis class C satisfy Assumptions 1-3 and the multiclass Tsybakov noise condition (Assumption 4) of [ggawall (2013). If we set S to be Algorithm 1 of [Agaraa (2013), then, given label budget of N and 8 E (0, (0,) we have with probability at least that 1ūrg E[
(+exp(e-(NN// (NN/ log N))
Corollary 3 efffectively allows PartiBandits to also handle real-valued outcomes (ex: Y ~ Unif(O, 1) if the analyst first discretizes Y into bins, effectively turning the problem setup into that of Corollary
3
Up to this point, we have focused on active learning algorithms S that guarantee exponential savings by grouping together instances that are likely to have similar labels. However, we may also consider S that not only identify homogeneous regions but also heterogeneous regions. Such S would be helpful for identifying strata for a distribution where labels are assigned by a simple threshold rule that outputs 0 if æ S 1/2 and otherwise, except that in the regions x E 11/4,1/2 and E (1/2,3/4) the label is flipped with probability 0.1. The optimal stratification scheme here splits the domain into four intervals 0, 1/4], (1/4, 1/2]. (1/2, 3/4], (3/4, 1). and allocate more of the Stage-2 samples to the middle two strata where the labels are noisier. In Corollary (below, we introduce an example of an S that helps identify such scheme. The proof is in the Appendix.
Corollary 4 (Heterogeneity-Awaan S). Assume the setup of Corvllary 2 Define S in the following way:
Algorithm 3 Heterogeneity-Aware Active Learning Algorithm
Require: hypothesis class C, label budget N. confidence level 6.
1: Run Algorithm 4.2 of Puchhin and Zhivotovskiy (2022) with a given label budget N' and hypothesis class C to obtain a classifier, hi
2: Let X. denote the abstention region obtained in Step 2 of Algorithm 4.2 of Puchkin and Zhivo tovskiy (2022).
3: Define h² (x) T((x)) + =(N) if x E X., and h (x)))) h(()) otherwise. e(N) is an arbitrarily small number relative to N (we may choose e e(N) exp( N/ log N)).
Output: h (xx
Then we have that given a label budget of N and 8 E (0, we have with probability at least 8.
Pirz E = δ (+texplci-- N/ log N))
What this S does capture heterogeneity via the abstention region X. produced within Algorithm 4.2 of Puchkin and Zhivotovskiy (2022): the region where labels are more likely to be ambiguous. It
8

---

converts the final classifier returned by Algorithm 4.2 of Puchkin and Zhivotovskiy (2022) into new classifier hi* that takes values in f,, 1, + s), splitting the space into four strata that isolate both homogeneous and heterogeneous regions. The result intuitive because the difference between h from Corollary 2 and h* is small, so Assumption is satisfied for h' if it is satisfied for that h.
The following yields matching lower bound.
Theorem 4 (Lower Bound for PartiBBandits) Consider the satta-generating process where X Uniff0, 1] and Y = 1 2 for some E [0, 1), with a ps- and ps- fractionn of labels Y flipped at random over X S and X > t, respectively, and Ps:P> - 1/4. Let C {1 {()) t) : E [0, 1]}. Then we have that for sufficiently large N,
v + exple (-N/ log(N)))
7 -
N
for constants e, 02 > and all estimators To BY) in this setup.
Since this lower bound matches the upper bound in Theorem 3 we have that the rate of Theorem is minimax optimal for this classical setting. This lower bound is based on a simple threshold setup with segmented label noise. As will show in Section [] this setup reflects common situation where subset of the unlabeled data is highly predictive of Y The proof of Theorem 4 starts by noting that the minimimum of 2 (9) among all possible stratification schemes, G. is precisely because of the way Y is generated. We can then use the lower bounds in Hanneke and Yang (2014) and Hanneke (201T to show that the exponential decay of the excess risk the optimal rate in this setup, so no algorithm can cause Σ₁ (g) to converge to its optimal value faster than this exponential rate.
5 Empirical Illustration
We empirically evaluate the performance of main algorithm, PartiBandits, and comparing it to SRS. We also test the WarmStart-UGe subroutine on the analogous mean estimation task when X can be stratified according to some stratification scheme priori. In most 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 present comparisons to other baselines, as well as analyses with other data generating processes, in the Appendix. We use Monte Carlo simulations and simulations involving nationwide electronic health records, with further details in the Appendix.
5.1 Simulations for Theorems 1 andd] 3
The error of the PartiBandits mean estimate is Ô [utoxp(e((-//oog())
as shown in Theorem 3
.V Hence the critical parameter that affects this rate is 1/, which is closely related to X's relationship with Figure[sshwss shows how PartiBandits performs as the strength of the relationship between X and Y varies. We see that generally, PartiBandits eventually outperforms SRS with relatively fewer samples, and this performance gap increases as the relationship between X and Y strengthens. We set S = 222 for our runs of PartiBandits. The right panel shows the performance of the War subroutine, which estimates the mean of Y when X can be stratified in advance according some stratification scheme G. We see that WarmStart-UCB consistently outperforms SRS when the stratification scheme closely aligns with the underlying decision boundary, effectively grouping observations with similar values of Y. We do not compare PartiBandits to A², A 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, and attempting to do can yield biased results Dong et al. 2025).
Y. The left panel of Figure
5.2 Simulations for Theorem 3 Using Health Records Data
To illustrate the gains of PartiBandits in a real-world setting, we leverage access to the American Family Cohort (AFC) dataset, which contains patient-level level data from over 1,000 practices participating

---

in the American Board of Family Medicines PRIME Registry. AFC contains fine-grained longitudinal records of patient race and clinical diagnoses.
Our outcome of interest a binary random variable, Y = HB, where H indicates the presence of hypertension and B is an indicator for whether the patient is Black. The estimand is B[Y], which corresponds to the fraction of patients who (1) are Black and (2) have hypertension. X is the probability that an individual is Black based on their zip code. Z, which is available in the data. This setup reflects common scenario in which researchers are studying demographic prevalences but lack access direct demographic labels (Andrus et al. 202T U.S. E.O. 202I), which motivates the use of proxy information like geolocation to guide sampling decisions. Though X is defined as a probability mapping. can also be viewed as an upper bound on the probability that Y 1 for an individual from a certain zip code since Pr Pr = 1 Z) < (Pr(H 1|2) A Pr(B B = 1 1 Z)). and therefore also a bound on the variance of Y. If H is known but is not (analogous to the problem setting of Elzayn et al. (2025)), we can obtain even more sampling efficiency gains by forcibly setting X = 0 for all H 0, since the variance of Y 0 for all such H. As a result, PartiBandits will not request labels B from individuals such that H = O. We compare the PartiBandits and SRS estimates of [[[]] under different label budgets. Because obtaining diagnosis and race data is often costly, this setting is such that labels are expensive, making it well suited to illustrate the benefits of using PartiBandits.
We focus on individuals whose derived probabilities of being Black, X, fall in the top and bottom 5th percentiles of the distribution. Although we restrict to these tails, this experimental setup does not simplify the problem. We are interested in the interaction variable HB rather than a single class label, which makes identifying any separation more difficult. Even when B = 1, not all such individuals have hypertension, so H B not always 1 and can even be O more often than when B O. Since, 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, H B may even be 0 more often than for those classified as B = O, so the class separation not immediate. This setup reflects the reality of many datasets that are not IID samples from the general population.
Our focus on the tails also shows that X does not need be highly predictive of across the entire population. We show that the mean within these tails where X is highly predictive can be estimated accurately with much fewer labels than SRS requires, freeing up the remainder of the label budget for regions where X is less informative of Y. This experimental setup thus illustrates how PartiBandits still offers a way of efficiently estimating the population mean.
To run this simulation, we draw, for each label budget, 500 random subsets of 10,000 patients each from the full AFC dataset of million patients. Within each subset, we restrict attention individuals whose geocoding-derived probabilities of being Black, X, fall in the top or bottom 5th percentile, and estimate the mean of Y for this subpopulation using both PartiBandits and SRS. We compute the 90th percentile of the resulting estimation eтrors to obtain Our choice of S is the classical 2 algorithm of Balcan et al. (2006) (thus putting us the regime of Corollary D. Figure shows the results and confirms that PartiBandits eventually outperforms SRS with relatively fewer samples. For smaller label budgets, SRS fares better but only by a universal constant factor.
high-probability bound for each method.
Performance SRS and PartiBandits (AFC Data)
Estimation Enrrr
SRS
0.050
Part Bandis
100
125
Number of Labels
Figure 2: Comparison of estimation error for different label budgets using the AFC data.
10

---

References
Agarwal, A. (2013). Selective sampling algorithms for cost-sensitive multiclass prediction. In Proceedings of the 30th International Conference Machine Learning, pages 1220-1228. PMLR.
ISSN: 1938-7228.
Andrus, M., Spitzer, E., Brown, J., and Xiang, A. (2021). "What We Can't Measure, We Can't Under- stand": Challenges to Demographic Data Procurement in the Pursuit of Fairnsss, arXiv:201 [cs].
Aznag, A., Cummings, R., and Elmachtoub, A. N. (2023). An active learning framework for multi- group mean estimation. Advances in Neural Information Processing Systems, 36:32602-32633
Balcan, M.-F., Beygelzimer, A., and Langford, J. (2006). Agnostic active learning. In Proceedings of the 23rd international conference on Machine learning, ICML '06, pages 65-72, New York, NY, USA. Association for Computing Machinery.
Bosch, V., and Wildner, R. (2003). Optimum Allocation of Stratified Random Samples Designed for Multiple Mean Estimates and Multiple Observed Variables. Communications in Statistics Theory and Methods, 32(10):1897-1909 Publisher: Taylor & Francis _eprint: https://doi.org/100 0.000..00111 1/STA- 120023258.
Burnashev, M. and Zigangirov, K. (1974). An Interval Estimation Problem for Controlled Observa- tions. Probl. Peredachi Inf., 10(3):51-61.
Carpentier, A., Lazaric, A., Ghavamzadeh, M., Munos, R., Auer, P., and Antos, A. (2015). Upper- Confidence-Bound Algorithms for Active Learning in Multi-Armed Bandits. 1.Giv:1507.05522 [es].
Castro, R. M. and Nowak, R. D. (2008). Minimax Bounds for Active Learning. IEEE Transactions on Information Theory, 54(5):2339-2333.
Cohn, D., Atlas, L., and Ladner, R. (1994), Improving generalization with active learning, Machine Learning, 1522)2201-221.
Dong, E., Schein, A., Wang, Y., and Garg, N. (2025). Addressing discretization-indueee bias in demographic prediction. PNAS Nexus, 4(2):pgaf027.
Elzayn, H., Smith, E., Hertz, T., Guage, C., Ramesh, A., Fisher, R., Ho, D. E., and Goldin, J. (2025).
Measuring and Mitigating Racial Disparities in Tax Audits*. The Quarterly Journal of Economics, 140(1):113-163.
Félix-Medina, M. H. (2003). Asymptotics in adaptive cluster sampling. Environmental and Ecological Statistics, 10(1):61 (011)..1-2..
Ganti, and Gray, A. G. (2013). Building Bridges: Viewing Active Learning from the Multi-Armed Bandit Lens. arXiv: 130906830000 [es].
Hanneke, S. (2011). Rates of Convergence in Active Learning. The Annals of Statistics, 39(1)3333- 361. Publisher: Institute of Mathematical Statisties.
Hanneke, S. and Yang, L. (2010). Negative Results for Active Learning with Convex Losses. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pages 321-325. JMLR Workshop and Conference Proceedings. ISSN: 1938-7228.
Hanneke, and Yang, L. (2014). Minimax Analysis of Active Learning, arXiv:14410000 [cs].
Henderson, P., Chugg, B., Anderson, B., and Ho, D. (2022). Beyond Ads: Sequential Decision- Making Algorithms in Law and Public Policy. In Proceedings of the 2022 Symposium Computer Science and Law, pages 87-100, Washington DC USA. ACM.
Jo, N., Andrea, V., Derek, O., and Ho, D. E. (2025). Not (Officially) in My Backyard: Characterizing Informal Accessory Dwelling Units and Informing Housing Policy With Remote Sensing. Journal of the American Planning Association, 91(1):30-45. Publisher: Routledge _eprint: http://wo..o////...0000044400000

---

Puchkin, N. and Zhivotovskiy, N. (2022). Exponential Savings in Agnostic Active Learning through Abstention, IEEE Transactions on Information Theory, 68(7):455-4465. arXiv:2102.00451 [es. math, stat].
Seber. G. A. F. and Mohammad Salehi, M. (2015). Adaptive Sampling. In Wiley
StatsRef: Statistics Reference Online, pages 1-11. John Wiley & Sons, Ltd. _eprint:
111997811184444 14:445522.stt005222.ppp222
Shin, J., Ramdas, A., and Rinaldo, A. (2019). Are sample means in multi-armer bandits positively negatively biased? arXiv:1905. 1197 [math].
Thompson, S. K. (1991), Stratified Adaptive Cluster Sampling. Biometrika, 78(2):389-377, Publisher: [Oxford University Press, Biometrika Trust].
U.S. E.O., (2021). Advancing Racial Equity and Support for Underserved Communities Through the Federal Government.
12

---

Appendix
5.3 Proofs
Before proving Theorem | we need few auxiliary lemmas. Lemma 1. Let (0,1) and
o²
R1(aN) = R1(n)= =
g
Ng.N
61
where ny = n (M1.,,...,, .G..). Then we have that with probability at least 6,
Tww 112 sc (MN)))) log
2 8
for some absolute constant C
Proof.
We begin by recalling that for any policy we have:
PgNNN
Σ Yg Pg.
Ng.N E.X,AA
and
AWS-UCB
Σ Po.N.
Since fws.ucB is unbiased estimator for μ and Var(Pg,N) we have:
- "a.N
Exx,es, Var(Y,P,) o "g..
Pr ((TWS-UCC - s) Pr
[
(Ag.N
El
28
(Lincarity)
s2
expp
(Hoeffding)
>>
0²
1
1g..N
Through the classical exercise of setting the left-hand-side to 8 and writing s in terms of 6, we have:
2
2
2
log
(R(N))
The next Lemma is simple but is important for linking the results of Aznag et al. 2023 to our work here. Lemma 2. Let
(n(n)) = (Egela) og) 2
13

---

Then,
R1(n") S E1(9)
N
where (g) is the expected conditional variance of given the stratification G:
Ei(9) - Σ
o0 Pg
gE(G]
Proof. This follows immediately from the fact that oz = P2 and from the norm equivalence property t and Le norms on RR
g oz and Pg (0,1) for all g E [
Proof of Theorem
We have the following directly from Section A.5 of Aznag et al. 2023
R1 (n) R2(n")
Gl n
7(3) min
n1 g.N
n n 1
(nn)
N ming n
N3
where:
min minge[[] Tg.
Σ1 Ege(G] Tg.
N Σ
Na N,
and n ( n.
This simplifies to:
Rí (n) - (n() (nnn))
GlIn n 11 N Wiing ng.N
63
|/n 12" 1
+
max
(2)
7g..N
Na
(II)
By Lemmas and it is sufficient to show that for sufficiently large N, the right hand side of Equation 2 upper bounded by some constant. We first show that term (I) upper bounded by constant and then we show that the same goes for term (IT). For all N that are not sufficiently large, we will perform classical Hoeffding analysis.
Bounding Term (I), We have that
2GON ΣS
||n n= 5 3G
V
mmn I
(3)
by Equation 23 of Aznag et al. (2023). and
Mmin N Σ
min N
by the analysis in Lemma of Aznag et al. 2023 What we want to do first is combine the summation on the right hand side of into one term. As long as
Σ
N >
2Σ
(Condition 1)
Tmin
(E.E.)))
14

---

then:
min RN V Σ
Owin N
1
(by Condition 1)
which implies that:
2GON Σ
min nin 4G 4G24224 2033, 3G,
(ON 2 by construction)
and so we can write:
8GON E1
Il -
mva-n'lle n0 5
min n
Thus, if we assume Condition 1 and
64G
N N
(Condition 2)
then we have that
2
GlIn - n 1130
G
8GCx 21
ming n
<
N Wming N
N miny ng.N
(1)
640 E ming ny
G
ming ni.N
C4G
Σ{N
< 1.
/632? [|n-n"
Bounding Term (II). We rewrite this term as the product of two terms,
G2
Ns
6) )
N and focus on each term in the product separately.
maxx
ng.N
6332 n 11
We first consider the term
We observe that if we assume Condition and
No
2/3
200 C
N
(Condition 3)
Σ
11/2 1/2 07 min
Cs(G,Ey.
then
15

---

No 2002 GC N3/2
(by Condition 3)
// min
1/2
32256 GOON N3/2
1,11/2 O
63 512
G N3/2
1 //1/33
1/2
min
632
512G M
min Σn
N
3/2
o2
63E]
8GCN Σ,
Omin Σg
3
N
min
63
|| 72 11:0
(by Condition
which implies
6352]
|| - n 1=
< 1.
o min
Na
Next, consider the term max,
Pa ng.. N
We have by the proof of Theorem I in Aznag
et al. (023) that if N is sufficiently large such that In-a" 1loe < 1, then
ming 15
12 9.N
max Ng.NN
22-n leee ming "7
Hence, we have that as long as Condition satisfied and
N (16GCN)²
(Condition 4)
Miomin
then, dividing both sides of that inequality by 2/N, we have
8GCN
Eiimin N
SGCy Σ,
Omin N
N
In n loc
min n
and therefore we have:
n
max g
9.N
<
< 64.
Ng.N
1 1/2
This gives that:
16

---

635
Ng.N
1|
n1" n='1.
max
- 1 64
o min
Ng.NN
N/
if Conditions 1, 3, and hold.
Hence, if we let:
C(G, E1, min) == max
C2(G, E1, miin), Cs E1, mmin), Ca Σ1, m
then we have that as long as N 2 C(G, E1, mmin), then
R1 (n) R1 (n*)
I 1 64
R1 (1(1))
x 65,
which means that
(i(n) S -
for some absolute constant C² and so:
2 3
(πws-ucB - -1} -12s C Ri (1x) log
2
$ C" (n(n")
Eg(
$ C' C'
log 8)
G
(Lemmal
N
Now for the analysis when N C(G, En, minn). By construction of Algorithm all E [G]:
have that for
T N. Ng.N 2
G
So by we have:
0² Rg.N
2 llg
(Ōws - 1u)2 <
(4)
8
9
G
2
2
C
log
(5)
N
9=1
C
log
(6)
N
We follow Aznag et al. 2023) in treating G as constant, and do the same with to obtain:
(TwS. - -1²²== (((g)
N
17

---

Remark (Dependence on т and G.). We note that by 6 and G made explicit, the rate is
when the dependence on the constants
TWS.UCC - = (
(g) .T
We discuss the dependencies and G further below.
Proof of Theorem The proof follows from the fact that the Neyman allocation corresponding to this 9 yields the smallest value of Σ, (9') over all possible stratification schemes G' and classical results about the estimation the mean of IID Bernoulli random variables (ex: Lemma from Hanneke and Yang (2010)). The idea is that first, (g) minimizes Σ1 (9') among all possible stratification schemes 9 (precisely because 9 aligns with the way the data are generated) and the fastest possible rate for estimating the conditional mean of Y on each stratum of G is simply the conditional variance divided by the number of labels requested on that stratum (ex: Lemma from Hanneke and Yang (010)). Combining the results yields the bound.
The proof essentially amounts to showing that the value of Σ (9) produced by Algorithm 2 is bounded from above by + exp(c (.N/oog(N)) To do this, first show that:
Proof of Theorem 3
(h(() Y) C C.E)))
We first use the law of total expectation obtain that:
BI(((x) =
Σ
BIG - Yy 1 Pr((() j)
(7)
jEJ
where is the image of h. We will now proceed with a birs-variance decomposition. We have that for all j E J,
E ((x) = 5] = ( Y)² [((x) = 5]
Define
M4 == T((x) = j].
Then we have that
B[G -)² T((X) = = -4 1 j].
Expanding using the identity
(a (a-c+c-b)+ 6)² (a (a
(c- 6) +
c) b),
we get:
BIU M4 + Y)² 17(x) = -
- Eluu; Y)2 Tx(x) = + + m4) Bl(ug - YY) [(x)) = j].
Note that the last term, 2(j - -41)= B[(Hy - Y) T((x)= = j]. is 0 by linearity expectation and the definition of μyz
So ultimately we have that:
[(((X) Y)² -12 h(x) = =- = 115) B((ns Y)2 ((x)) = 5].
18

---

Now using 7 we can rewrite:
- Σ
11)) 11)2 1 1j-
B(( - v)2 1ī(X) 5]] Pr(((X)= j)-
(8)
jEJ
we let of Var(Y Aj) for each j E J, then by definition of jaj we have that:
((((X)-X)) -
Σ of Pr(m((x=)) =
jEJ
(g).
Now by Assumption
we have that:
E[(h(X- r2 S exxppee (xp((e((-N/ogg(M) (-N/log(A))),
so we have that:
(9) + exple log()))))))
for some constant > 0. Applying Theorem with our choice of and the remaining label budget of N/2 we obtain that:
(APB - E[Y)) = Ō V (++explc (-N/log(N))
N
Remark (More on the Dependence on T and G (19)). We note that by Remark is straightforward to show that when the dependence on and G is made explicit, the rate is where [g] is the number strata.
[Pre - E[y]|2 = (191- ( (Nesplee-N/oog())) N-T
Remark (Different Loss Functions), We note that Theorem 3 still holds if one is interested in asymmetric misclassification cost, a 1{h(X) = 0, Y = 1} 1(h(X) = 1, = 0} instead of the ordinary squared loss. We assume here that Y is binary.
Proof. To see why we first start by doing an analogous decomposition to
E 1fh(x) 0, Y = = b = y =0]
= E 1(Y = 0} [((x) = 1] p()(x)== = 1) + + a ( =0] p(N(x) = -),
and this precisely equal to:
--B[(((x)- ((() ( -0)o())))))--B((//00 (1) -)- 0).
Therefore the rest the proof for this alternative misclassification cost is analogous to the proof of Theorem just with adjustments for the constants and b.
[i
This result is helpful in a setting where X Unif(0, and
Bern(O).
X < 0.5,
Y
B0.n5..25),
X > 0.5.
Here, an alternative loss would be helpful because the typical Bayes optimal classifier based on the squared loss would not distinguish between X 0.5 and 2 0.5, but the Bayes optimal classifier resulting from an asymmetric misclassification loss would.
19

---

Proof of Corollary 7
A² learns classifier with exponential PARTIBANDITS. Hence Assumption T is satisfied, and Theorem follows.
By Theorem 5 of Hanneke (2011). S ==
savings in Stage
Proof of Corollary
2 By Theorem 4.1 Puchkin and Zhivotovskiy (2022). Algorithm 4.2 of Puchkin
and Zhivotovskiy
2022 learns a classifier with exponential savings in Stage I of PARTIBANDITS. is satisfied, and Theorem follows.
Hence Assumption
Proof of Corollary 5 As demonstrated by Corollary of Agaawaal 2013), Algorithm AgaAgal
(2013) learns classifier with exponential savings in Stage 1 of PARTIBANDITS, Hence Assumption iis satisfied, and Theorem folloow..
-N (TR((T) )
Proof of Corollary 4 We assume without loss of generality that exp
S 1 (the result can be
easily adjusted when this assumption does not hold).
By Theorem
3 it is sufficient to show that Assumption I is satisfied:
(x) (X) 2) x exp (c- log(N))
for some constant > For this, sufficient show that the excess risk of and he differ by at most exp showing that:
-N Ing(N)). ). since i already satisfies assumption (Corollary 2 Formally, this means
B((h(x) - r) -)(x)((( - y]9] $ exp (c.
-N
(og((N)
We have that:
B[(h(x) - Y - EX) Y)
B[(h(x) ry² (fx(() -1)
(linearity)
E [(x)))2 - 2b(X)Y +y² (x)xx) 2(x((x)) y2
E [(x)2 i (xx) 2Y 22--22222 -ixx) h h-
S E M(x)2 hm (x)² . R
+2 2 (f(x)) ix)x) h + 7-
(triangle inequality)
We note that E (()xx))-<(0 7 7 [T| tion of h Furthermore, we have using the difference of squares decomposition that:
< c₁ exp
N Tog(NN
for some e₁ 2 0 by construc-
b :
- +x -xx)) : T
= (xxx)))xxxxx(((xxx
S 2B ( hx(x)) 7
where the inequality follows from Jensen's inequality, definition of hr, and the assumption that
NV [og(NN
< 1. Again by definition of Ti*, we have that 2E
(f(()) - h (x)) h Tr
exp
N log(M)
for some C2 0. Thus, have that E [(x)² - hx)2 h
C2 exp
C2 exp
[og(NV)
.
and therefore that:
B[((((x) - r) -) -)(((((())-))3
r)2a S exp (c- (og(t)
20

---

Proof of Theorem 4 As shown in Theorem 2 the best possible performance achievable by any algorithm is vN for some constant C > O that may depend on other parameters of the problem. So determining the lower bound for any algorithm returns an estimate of j becomes a task of determining the lower bound of C O for any algorithm. We were able deduce the lower bound of C' in Theorem 2 by restricting to a particular situation (that which the analyst basically is able to leverage X perfectly in a sense). We consider here more general problem setting where the analyst limited to estimators f that compute weighted aggregates of stratum means over partitions consisting of exactly Two strata, So and Si, which make up some stratification scheme G. The analyst learns So and S₁ through some procedure after N/2 label requests, and uses the rest of the label budget to estimate the means within each stratum and then aggregating them the end. This is a best case scenario because this is aligns with the data generation process and helps approximate the optimal value of C > (one only needs to learn the best choice of Sg and S1 to arrive at the optimal C, as shown in Theorem 2). Hence, the lower bound on (īi E( is simply the lower bound on Eg) so we focus on the latter quantity in this proof.
Assume without loss of generality that μ40 - E|y X E Sn] S 1/2 and B|y X E Si] 2 1/2. There exists classifier such that h- h0) So and h h-1(1) St. Recall from 8 that we have the following for any h:
/ [[(h(X) - 2) = ((1 ([1--1)² E[(u Y)² ((X) = 11) P(M(X) 1)
+ (18 + Bi(po Y)² |(X)) = 0]) P(h(X)
Since Mo0 S 1/2 and μr 2 1/2, we have that (1 - n1)² s (1 11)p1 and HB S Ho(1 Mu). Thus,
we have that:
(((())-))) -
((1 (11--1)s E(411 V)² ((X) 11) P(M(X) 1)
+ (Ho(1 - mO) E[( - Y)² I(() 0) P(h(X) 0)
Since Y is Bernoulli, we have by the corresponding variance formula that:
[[(((X) - YY (E(( Y)² 1)²1 1] P(I(X) 1)
+ 2E Y)² A(X) 0) P(h(X) = 0)
$ 2 (G).
What this shows is that if we can lower bound [[(h(X) y) then we can lower bound Σ, (G). In their discussion of Theorem 4, Hanneke [201 I noted that the rate:
E ( y) - v exp(c Caappe-(-N/loonn (-N/log(NN)
minimax optimal in precisely this situation where the strata map on to threshold classifiers (see also Castro and Nowak (2008): (Cohn et al. (1994); Burnashev and Zigangirov (1974)), so we have that there a 1) such that with probability least
Eol Y) 2 ( exp(c-
for some absolute constant K 0, and therefore that:
2 (9) K (u + exple (-N)))(NN))),
for (X,Y) D.
21

---

Additional Details Regarding Simulations
Details Regarding Simulations for Figure
In both simulations, each dataset contains 10,000 unlabeled instances. For each label budget from 80 to 140, we run 500 simulations. For each run, we compute the absolute difference between the mean estimated by the algorithm and the true mean; we then report the 90th percentile of these 500 errors.
Details Regarding AFC Simulations
In each simulation, we construct the mapping from ZIP code to the probability of being Black using the 10,000-patient sample drawn for that run. This means that the probability mapping X is re-estimated in every simulation based the ZIP code-race distribution within the sampled subset.
Additional Notes
Note (Notes about Theorem I. When the dependence on the constants T and G is made explicit, the rate is
( (9 N
Pws-uce -[Y]
Following Aznag et al.) (2023), we treat G as a constant, and do the same for T. As discussed above and shown here, the dependence on vanishes for large N relative amin and for all other N, the rate still holds with slightly larger constants (including a linear factor of 1 We also discuss the special cases when {0, 1}.
Setting T = 0 allocates all labels to Variance-UCB achieving optimal rates when both min and the label budget are sufficiently large. For sufficiently large N relative m is longer relevant to the problem and the ordinary rate holds thanks to the main result of Aznzg et al. (2023). However, a positive ensures that least some samples are drawn uniformly from every group, preserving the guarantee of Theorem 7 with only slightly larger constants. For large label budgets, the effect of 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 present the simplified rate Ô En (x(
in Theorem 7
When T is small (close to 0), we rely more heavily on the Variance-UCB adaptive sampling strategy. which is reasonable if the label budget is large and/or mii is not too small. When large (close to 1), 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. the final bounds, essentially appears as a constant in the denominator, but for sufficiently large label budgets, its effect becomes negligible as the guarantees of the Variance-UCE UCC procedure dominate.
The WarmStart in Algorithm I serves as a buffer, ensuring that Variance-UGG remains effective when the overall label budget is small and the algorithm might otherwise undersample low-variance groups. As As(Azag et al. (2023) discuss, the regret of Variance-UCT can scale inversely with a the smallest group variance. By assigning each group minimum number of samples, the warmstarl helps guarantee the algorithm's 's fast-rate properties. For small N, the risk bound behaves ((1/TN) where т is the warmstart proportion. As N grows and the fast-rate conditions hold, the dependence on vanishes and the sharper rates of Aznag et al. (2023) apply. We follow followAznag et al. (2023) in focusing only on the dependence on the label budget in the bound.
Note (Note about Empirical Illustration.). Here present additional experiments using alternative data- generating processes, varying degrees of class separation, and asymmetric class distributions. We also consider the comparison of PartiBandits to alternative baselines.
22

---

02
SRS
Estimatton Enror
PartiBandits (Main Algorithm)
Lown
50
75
100 Number of Labels (N)
125
150
Figure This plot compares the performance PartiBandits SRS when the labels are generated according the following logistic data generating process: X Unif(o, 1] and Bernoulli [+expi-(8++a X where 3 11vv and 31 2/v. This corresponds to a Logit-type DGP, with 1/v governing the steepness of the logistic curve. For each label budget. we generate 500 hypothetical datasets in this way, apply SRS and PartiBandits each, and compute the resulting error rates. We then take the 90th percentile of these error rates obtain classical 90% ceatt-oobbbilitycondiiii relatively fewer samples, with performance gains becoming more pronounced when better predicts Y and decreases.
bound. PartiBandits eventually outperforms SRS with
0.
04
Estmalion Error
PartiBandits (Main Algorithm) Hight
0.3
0.2
Low
0.1
50
75
100
125
150
Number of Labels (N)
Figure This plot compares performance PartiBandits to SRS when labels are generated according to the following asymmetric probit data generating process: X ~ Unif[-5,5 5] and Y Bernouli ( ((1/v) (X 0.25))) ) where 2(-) denotes the standard normal CDE This corresponds to Probit- type DGP, with 1/v controlling the steepness of the probability curve and X 0.25 marking the midpoint threshold. For each label budget, we generate each, and compute the resulting error rates. We then take 90th percentile these error rates to obtain classical 90% high-probabilityconnieee samples, with performance gains becoming more pronounced when X better predicts Y and v decreases.
hypothetical datasets in this way, apply SRS PartiBandits
bound. PartiBandits eventually outperforms SRS with relatively fewer
23

---

0.5
SRS
Estimetion Erro
PartiBandits (Main Algorithm) High
02
0.1
Lowi
0.0
Number Labels (N)
Figure 5: This plot compares the performance of PartiBandits SRS and Thompson sampling and SRS for label budgets from 0% and 10%) randomly flipped to introduce noise. The proportion of flipped labels is equal to v by definition. For each label budget, we generate 500 hypothetical datasets in this way, apply SRS, Thompson sampling, and PartiBandits to each, and compute the resulting error rates. To execute the Thompson sampling, we use the standard Beta-Ba Bernoulli Thompson Sampling algorithm with uninformative prior Beta( L, 1). At each round, algorithm samples a success probability from each arm's posterior, selects the arm with highest draw, observes Bernoulli reward,
100. Here, X Unif(O, 1] and Y = (X 0.5). with fixed fraction of Y"s (between
updates the corresponding posterior. In our setup, we ran rounds with K = 3 arms (true = 10.1, 0.5, 0.8)) for the prototype and K = 5 hins over [0, 1] with threshold of 0.5 for the binned variant. We then take the 90th percentile of these error rates to obtain classical 90% hath---obbbbbillycooideeeee bound. PartiBandits eventually outperforms SRS and Thompson sampling with relatively fewer samples, with performance gains becoming more pronounced when better predicts Y and decreases. We also observe that, over time, Thompson sampling ceases to yield better mean estimates, consistent with theoretical results suggesting that this procedure can yield biased mean estimates in common settings (Shin ct 2019).
3000
NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility. transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. [NA]
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist one of the factors in their evaluation. While "[Yes] " is generally preferable "[No] " is perfectly acceptable answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write justification elaborate. All supporting evidence can appear either in the

---

main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point the section(s) where related material for the question can be found.
IMPORTANT, please:
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims in the abstract and the introduction are directly supported by the Theorems and the empirical results in our paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss the assumptions required for our theoretical results to hold. We also state the parameters and contours of the simulations conducted.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
25

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The paper does provide the full set of assumptions for each theoretical results. The complete and correct proofs are available in the supplementary material.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
The paper provides the details behind the simulations we conducted, the specific data generating processes, etc. The code for the monte carlo simulations is available in the supplementary material. The data for the AFC simulations cannot be provided as it contains highly sensitive personal identifying information and would present significant ethics concerns if released.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
26

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: As noted above, we will the data and code for the monte carlo simulations will made available in the supplemental material. However, this cannot done for the AFC data for the reasons discussed above.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cr cc/
public//giidee/CodeSuuiiii www.sooobssssiiooooooooooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips https///cccpllicgggeeeeeeeeeeeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Please see the answers above regarding data and code and the details for conducting our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
27

---

Justification:
The simulations do not report "error bars", however, the main metric in the simulations is the 90% confidence bound on the estimation enror, which functions as the upper end of a confidence interval on the estimation error. Thus, the experimental results already incorporate measurement of uncertainty.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We describe the procedure for the monte carlo simulations which are quite simple and do not take long to run.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
We confirm that this submission complies with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
28

---

10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: This work contributes positively by reducing the number of data intrusions required obtain accurate estimates, which can minimize burden, cost, and potential harm in sensitive applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
We do not use pretrained language models, image generators, or scraped datasets.
Justification:
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
29

---

Justification: All datasets and code packages used in this paper are properly credited, and their licenses (including CC-BY 4.0 and equivalent open- licenses) have been reviewed and respected.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g.. website), the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, www.ssppppsssaaeee thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We introduce a new R package, PartiBande ts, and provide documentation, usage examples, and licensing information alongside the code.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmdel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not invovle crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
30

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non non-standard components.
Please refer to our LLM policy //npur:////uuraie cc. Cc//oonferencee/22255///// for what should or should not be described.
31


Reviews:
[{"review_id": "fVpU3Rw5vp", "reviewer": "Reviewer_ZBCc", "summary": "The authors consider the problem of estimating the mean of a binary random variable $Y$ using limited number of labelled instances. They assume that they have an access to covariate $X$ that may be informative about $Y$. The authors suggest a PartiBandits algorithm, which takes a hypothesis class $\\mathcal C$, label budget $N$, a confidence level $\\delta$, and a buffer fraction $\\tau$ as an input. The algorithm performs in two steps. First, using the $A^2$ (agnostic active) algorithm as a subroutine, PartiBandits stratifies the covariate space into two sets. After that, it estimates stratified mean using Variance-UCB algorithm [Aznag et al., 2023] with a ''warm-start'' step (WarmStart-UCB). Assuming the Tsybakov noise condition and that the Bayes optimal classifier belongs to the concept class $\\mathcal C$, the authors prove minimax optimal upper and lower bounds on the accuracy of mean estimation. They illustrate advantage of their procedure over stratified random sampling using Health Record Data.", "questions": "1. How the result of Theorem 5 will change if $Y$ is a (non-Bernoulli) random variable taking its values in $[0, 1]$?\n\n2. Is it possible to further improve the algorithm by increasing the number of strata?\n\n3. In [Puchkin and Zhivotovskiy, 2022], the authors show that exponential improvements in agnostic active learning are possible even if the Bayes optimal classifier does not belong to the concept class $\\mathcal C$. Instead, they require $\\mathcal C$ to have a bounded combinatorial diameter and a finite star number. Is it possible to replace Assumption 2 by similar requirements in your setup?\n\n4. Is stratified random sampling the only existing competitor? Can you compare PartiBandits with other algorithms?", "limitations": "Yes.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strenghts.\n\n1. According to the numerical experiments, the PartiBandits algorithm significantly outperforms stratified random sampling. \n\n2. The authors prove rigorous guarantees on the PartiBandits performance. They also prove a matching lower bound claiming optimality of their procedure.\n\n\nWeaknesses.\n\n1. The setting is limited to binary outcomes.\n\n2. The assumption that the Bayes classifier must belong to the concept class $\\mathcal C$ (Assumption 2) is rarely satisfied in practice.\n\n3. The number of strata is limited to 2.\n\n4. The proofs of the upper and lower bounds (Theorems 5 and 6) significantly rely on the well-known results from the literature on active learning (for example, Theorem 4 from [Hanneke, 2011]).", "was_revised": true, "final_justification": "The authors have thoroughly addressed my main concerns, namely Weaknesses 1-3 in my initial review. The authors extended the setup for multiclass and multistrata scenarios (see Corollary 2 in their rebuttal). Besides, relying on findings of Puchkin and Zhivotovskiy (IEEE Trans. Inf. Theory, 2022), they showed (Corollary 3) that nearly exponential savings are still possible even if the Bayes optimal classifier is not in the concept class.", "rebuttal": "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 \\textbf{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.\n\n\\textbf{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.}\n\n\n\\textbf{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.}\n\n\n\\textit{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.\n\n$$***$$\n\n\\textbf{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: \n\n    $$\n    \\left| \\widehat{\\mu}_{\\text{PB}} - \\mathbb{E}[Y] \\right|^2 = \\tilde{\\mathcal{O}}\\left( \\frac{\\nu + \\exp(c \\cdot (-N/\\log(N))) }{N} \\right),\n    $$\n\n    where $c >0$ is a constant. \n    \n$$***$$\n\n\\textit{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.\n\nWe 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.\n\nCorollary 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).\n\nCorollary 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). \n\nCorollary 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). \n\n\\subsection{Strengths and Weaknesses}\n\n\\textbf{W1:} \\textbf{We have added Corollary 3 to address this point (see Point 0).}\n\n\\textbf{W2}: \\textbf{We have added Corollary 2 to address this point.}\n\nAs 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.\n\n\\textbf{W3: We have added Corollary 2 to address this point.}. \n\nWe discuss this point more in point W1 in our response to Reviwer ehtb.\n\n\\textbf{W4}: \\textbf{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.}\n\n\\subsection{Questions}\n\n\\textbf{Q1}: \\textbf{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.}\n\nThe 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]$.\n\n\\textbf{Q2}: \\textbf{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.}\n\nPlease see our response to point W3 and point W1 in our response to Reviwer ehtb for further details.\n\n\\textbf{Q3}: \\textbf{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).}\n\nSee also our response to point W2 above.\n\n\\textbf{Q4}: \\textbf{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.}\n\nWe address this in section W4 in our response to Reviewer MZfC."}, {"review_id": "AA49jrO3il", "reviewer": "Reviewer_ehtb", "summary": "The paper introduces a novel (meta) algorithm for the estimation of the mean of a binary random variable Y that depends on some (informative) X in a scenario where labels are expensive.  The proposed algorithm consists of two parts: first, an active-learning approach learns a partition G of the space X so that the partitions $A_1,\\dots, A_G$ with probability $P_k=P(X\\in A_k)$ minimize Var(Y|A_k). This is done using an online classifier that partitions based on its predictions. Then in a second step a bandit algorithm is applied which selects the partition A_k so that a bound on the weighted variance $P_k^2 Var(Y|A_k)$ is minimized. The first part of the algorithm is a direct application of the existing algorithm $A^2$, while the second part is a slight variation of Variance-UCB.The idea is that the initial split reduces the variance based on how informative X is, while the second part then focuses the majority of the remaining sample budget on the high variance partitions. The work establishes matching lower and upper bounds for each algorithm, where the lower bound is given by example. Finally, two small experiments illustrate the work.", "questions": "1. Can you provide experimental results with more difficult settings, e..g, harder to separate classes?\n2. Can you provide scaling results for Algorithm 2 as g increases?\n3. Can you sopecify what the significance of tau is in Algorithm 1, and how your results change when tau is set to 0?", "limitations": "no. A conclusion and discussion is missing. I mentioned above the missing discussion of the scaling behavior in several cases.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "I was not able to conduct a proper review since the work is very far outside my core domain and i only realized very late that the total review load is around 150 pages and since I scheduled this paper last due to its complexity, I was not able to read up on related work. I still do a best-effort review based on the main article without appendix. I am choosing low confidence and opt for weak accept.\n\nThe strength of the paper is a clear theoretical analysis of the problem including matching lower and upper bounds. This is based on the reliance on strong results from prior work, but is still an achievement. \n\nThe main weaknesses i would like to point out is the overall minimalistic approach of the paper\n1. Algorithm 2 seems underexplored and minimalistic.\n2. Not all sources of uncertainty seem to be carried over to the analysis \n3. the experimental part is very short and especially experiment 2 is not understandable to me. The experiment do not provider enough information to establish the theoretical rates and are too minimalistic to properly explore strengths and weaknesses.\n4. In general, bounds and their intuitive meaning are underexplored.\n\nDetails:\n1. From my understanding, Algorithm 2 partitions the input space in just two parts: A_0 and A_1 based on the classification of the hypothesis produced by the A2 algorithm. As a result, the gain of Algorithm 2 is just based on how well the classifier predicts the right class but does not make use of variability of regions within a class. E.g., assuming the true data is generated by logistic regression, we would ideally want to partition in three classes, where the middle class is close to the boundary.\n2. Algorithm 1 assumes that $P_k=P(x\\in A_k)$ is known. This either means that the assumption is that the initial unlabeled dataset is the full distribution, or the values are given by an oracle. I do not think that matters for small number of partitions since we can assume that the amount of unlabeled data is always sufficiently large compared to the amount of labeled data that the additional error can be included in some of the factors, but formally it is not discussed in the theorems.\n3. Experiment 2: Do I understand it correctly, that in each trial from the 10000 drawn people only the ones are selected which fall in the 5% lower and upper intervals, that is we see the results based on easily separable classes? In that case, I think the experiment is almost a \"strawman\" ( for a lack of better word) since in this example it is also incredibly easy to find the classifier that is separating them with just a few samples. In this case, a maximum margin classifier with just 1 sample from each class would have a high probability to be bayes optimal (when both labels are correct). So this experiment shows the extreme case of where the algorithm must perform best. I still do not understand the shown figure, because it appears the result is almost constant, even though i would expect to see 1/N convergence in both algorithms. A larger sample range, especially in the small sample regime would be interesting. Similarly, for experiment 1 only the case of symmetric noise and two same size classes is investigated. I do not find an explanation in the paper why this special case should be sufficient, since in this case both partitions should have identical variance.\n4. I would have loved to read more about the intuitive meaning of the bounds. For Algorithm 2 for example, I do not understand why it should not depend on the number of partitions. I would intuitively expect the bound to become worse as the number of partitions grows. Similarly, for Algorithm 1 i do not understand why the warmstart is necessary, since it does not appear in the bound.\n\n\ntechnicalities:\n1. I was not able to find the warmstart component in Algorithm 1. I think the initialisation of the std-dev in line 1 should do that, but then i would expect some dependency of t on g(not G) to ensure that each g has an area where its variance is maximized.\n2. A number of quantities appear to be undefined in the main article. For me as a non-expert, especially considering the density or the article, it made me search quite a few times because i was not sure whether i missed something. It might be that some are standard notation, but i would like the authors to check the following variables that i mark with the first occurrence i found:\n- $\\tilde{\\mathcal{O}}$ line 65 (is this: $\\mathcal{O}$ ignoring log-terms?\n- $R_1$, line 201\n- $c_1,c_2$ line167. \n- $ \\Sigma_1(\\mathcal{G})$ line 178\n3. Theorems do not clearly state assumptions. I know that you stated assumptions globally, but I think it is crucial to describe concrete assumptions of each analysis as minimal as possible, e.g., Theorem 3 should not depend on Assumptions 1 and 2 since the algorithm assumes $\\mathcal{G}$ is given. At the same time, considering the choice of $c_1,c_2$ in Algorithm 2, it appears that Theorem 3 holds for a much broader set of random variables than binary Y.", "was_revised": true, "final_justification": "The authors answered all points I raised, and also the ones raised by other reviewers. Additional experiments were added and the discussion did not uncover any fundamental flaws in the maths or experiments. To the contrary, the authors expanded their analysis to cover the issues raised in form of additional corollaries and thus greatly reduced my feeling of the work being \"minimalistic\". I think this strengthens the paper significantly.\n\nWhile I am not an expert in the area, I see that the other reviewers express happiness with the response, and indicate a possible raising of their scores. This was the best paper in my batch, and in light of all the above, I feel safe to propose acceptance of this paper at NeurIPS.\n\nFinally, I would like to thank the authors for a very honest rebuttal and discussion, with clear indication of weaknesses and mistakes.", "rebuttal": "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 \\textbf{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.\n\n\\subsection{Strengths, Weaknesses, and Technicalities}\n\n\\textbf{W1}: \\textbf{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.}\n\nCorollary 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.\n\nThe 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.\n\n\\textbf{W2}: \\textbf{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.}\n\n\n\\textbf{W3}: \\textbf{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.}\n\nAdditionally, 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.\n\n\\textbf{W4}: \\textbf{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.}\n\nYou 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.\n\nThe 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.\n\n\\textbf{T1}: \\textbf{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.}\n\n\\textbf{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.}\n\n1. Yes, $\\widetilde{\\mathcal{O}}$ ignores logarithmic factors and fixed constants.\n2. 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.\n3. $c_1$ and $c_2$ are simply the subgaussianity constants (see Aznag et. al (2023)), but we will clarify them in the revised paper.\n4. $\\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. \n\n\\textbf{T3}: \\textbf{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.}\n\n\\subsection{Questions}\n\n\\textbf{Q1: Please see our response to point W3 above.}\n\n\\textbf{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).}\n\n\\textbf{Q3}: \\textbf{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.}"}, {"review_id": "QBZDhuuAtX", "reviewer": "Reviewer_Wv21", "summary": "This paper proposes a PartiBandits algorithm for estimating the mean of a binary random variable with auxiliary information, which includes two stages: (1) Using a Agnostic Active algorithm to learn a partition of the data to shrink the average conditional variance of the variable; (2) Using a UCB-style subroutine to actively request labels for the learned strata. The paper proves that the algorithm could achieve a near-exponential improvement in estimation accuracy in comparison to Simple Random Sampling, and the convergence rate is minimax optimal. This paper also provides empirical verification of the algorithm's performance using real-world data.", "questions": "1. Is it difficult to extend the algorithms and theoretical results to non-binary cases?\n\n2. How is the computational cost of the PartiBandits? Could it be applied to large datasets in practice?", "limitations": "The Bayes Optimal Classifier is in the Class assumption (Assumption 2) could be too strong for practical settings, despite it is common in theoretical analysis.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength: This paper shows strong results and proposes algorithms that achieve theoretically-optimal convergence rates. It also bridges a gap between UCB algorithms and disagreement-based methods. This paper provides both theoretical and empirical evidence to support the claims. It is also well-structured and the presentation of the main results are clear.", "was_revised": false, "final_justification": "", "rebuttal": "\\textbf{Q1}: \\textbf{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.}\n\n\\textbf{Q2}: \\textbf{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.}"}, {"review_id": "iNWvwYn3ka", "reviewer": "Reviewer_MZfC", "summary": "This paper proposes PartiBandits, an active learning algorithm for estimating the population mean of a binary label with covariates and limited labels. The method combines stratified sampling with an adaptive allocation policy, and achieves a near-exponential convergence in Massart low-noise condition. Matching upper and lower bounds are provided and experiments on semi-synthetic data validate the gains under ideal stratification.", "questions": "What is the convergence rate if the assumption is not Massart low noise condition (i.e. hard-margin) but rather more general Tsybakov condtion?\n\nThe authors eliminate the dependence on \\sigma_{min} but introduce another buffer \\tau, what is the specific dependence on \\tau in the final bounds? Is there a trade-off between big \\tau and small \\tau?\n\nWhy is A^2 not used in the evaluation, is there any practical concern, e.g. efficiency?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nTackles an under-explored but important problem: active learning for population mean estimation.\n\nBoth matching lower bounds and upper bounds are given, making the algorithms minimax optimal.\n\nGiven a pre-defined stratification G, WS-UCB can achieve nearly exponential convergence rate when \\mu is small.\n\nEmpirical validation, though limited, shows clear improvements over SRS.\n\nWeaknesses:\n\nThe main result only considers the hard-margin condition, which is rarely the case in practice.\n\nIt is not clear what is the role of buffer \\tau in the final bound of theorem 3\n\nThe stratification process is under-specified: A^2 is mentioned in theory but not used in experiments. \n\nExperiments have only SRS baseline.\n\nMinor presentation issues: e.g. what are \\rho_{\\leq} and \\rho_{>} in the statements of Theorem 4? Some notations, like R_1(n),  lack definition in the main text.", "was_revised": true, "final_justification": "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.", "rebuttal": "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 \\textbf{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.\n\n\\subsection{Strengths and Weaknesses}\n\n\\textbf{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.}\n\n\\textbf{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.}\n\n\\textbf{W3}: \\textbf{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.}\n\n\\textbf{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.}\n\nWhile 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.\n\n\\textbf{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.}\n\n\\subsection{Questions}\n\n\\textbf{Q1: Thank you for raising this question. We address this in point W1.}\n\n\\textbf{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.}\n\n\\textbf{Q3: Thank you for raising this question. We believe we have addressed this point in our response to W3 above.}"}]

Author rebuttal:

Meta-review: This article explores the disagreement-based active learning framework to improve the sample efficiency for the population mean estimation of a target variable $Y$, enabled by the information of covariates $X$ which are predictive of $Y$. In the proposed algorithm "PartiBandits", a predictor of $Y$ is first obtained by a disagreement-based active learning method, allowing the construction of homogenous subgroups w.r.t. to the value of $Y$. Then the population mean estimation is carried out from these subgroups with a UCB (upper confidence bounds)-style procedure. The minimax optimality of "PartiBandits" is demonstrated with matching upper and lower bounds. It is also tested on a real-world data set against the standard baseline of random sampling.

The main issues raised by Reviewers concern the strong assumptions for the theoretical results and the restrictive setting for the deployment of the algorithm. The authors effectively addressed these limitations, many of which are simply removed by using more advanced alternatives for disagreement-based active learning. Reviewers were satisfied with the rebuttal, and recognized the interest and the theoretical strength of the proposed method.
Decision: Accept (poster)