Near-Exponential Savings for Mean Estimation with Active Learning

2025-11-07

AI Summary

This paper addresses the problem of estimating the mean of a $k$-class random variable $Y$ given high-dimensional covariates $X$, using only a small number $N$ of labeled samples. To this end, we propose PartiBandits, a novel active learning algorithm that unifies the Upper Confidence Bound (UCB) framework with disagreement-based active learning. Our method employs a two-stage strategy: first partitioning the sample space based on $X$, then applying a WarmStart-UCB subroutine within each partition to construct a low-variance stratified sampling scheme. We theoretically establish that the estimation error decays as $ ilde{mathcal{O}}ig( frac{ u + exp(-c N/log N)}{N} ig)$, achieving near-exponential convergence—matching the minimax optimal rate in classical settings. Empirical evaluation on real-world electronic health record data demonstrates substantial improvements in both estimation accuracy and sample efficiency over baseline methods.

Technology Category

• Machine Learning: Active Learning
• Reasoning under Uncertainty: Stochastic Optimization
• Constraint Satisfaction and Optimization: Constraint Learning and Acquisition

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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}_{ ext{PB}}$, such that $left( widehat{mu}_{ ext{PB}} - mathbb{E}[Y] ight)^2$ is $ ilde{mathcal{O}}left( frac{ u + exp(c cdot (-N/log(N))) }{N} ight)$, where $c>0$ is a constant and $ u$ 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 PartiBandits package in R.

Problem

• Estimating mean of k-class random variable with limited labels
• Using auxiliary covariates to improve label efficiency
• Developing active learning algorithm for minimax optimal convergence

Innovation

• Active learning algorithm with two-stage design
• Partitions data to reduce conditional variance
• UCB-style subroutine for label selection