Best-Arm Identification with Generative Proxy
July 08, 2026
1 Introduction↩︎
Best-arm identification is a canonical formulation for decision making: among a finite set of alternatives, the decision maker seeks to identify the best one using as few samples as possible facing observation uncertainty. In many applications that motivate best-arm identification, these samples are online and expensive, e.g., a reward observation may require exposing a user to a recommendation, running a pricing experiment, or deploying a candidate policy in a real environment. Fortunately, cheap proxy evaluations for decision units can be produced with modern predictive and generative models. Such proxies are not reliable enough to replace the real reward, but they can be informative about its variation. Following the same philosophy of prediction-powered inference [1], the cheap proxy prediction can explain part of the uncertainty in an expensive target quantity and thereby reduce the number of costly observations needed for inference. In this paper, we study how to use such proxy scores as control variates in fixed-confidence best-arm identification. These proxy scores finds its place in many real world best-arm identification problems:
Recommendation Systems: In personalized recommendation, item cold-start is a tough problem. When a set of items (or themes) is newly launched to a user, the recsys needs to find which item (or theme) is the favorite of the user by click-through rate or watching time, facing uncertainty in user behavior due to many aspects. However, online samples are expensive as the user needs to be exposed to the items and usually only a few interactions are allowed. With the help of large language model (LLM), we are able to prompt the LLM with the item information, user interaction history, and all context information including session state, time and device, other recommended items, etc., to receive a proxy score, e.g., a score from 0 to 3 where 0 corresponds to no interest and 3 corresponds to very interested. During the preparation for the item launch, LLM can be called for many times for different user and different possible contexts, and we can acquire an accurate offline proxy score for each item, but the correlation between this proxy score and the real CTR or watching time remains unknown. Only after the items (or themes) are launched, the correlation can be estimated at the same time recsys explores the user’s favorite. Due to the requirement of high concurrency and low latency, only a few or only one LLM call is available after each user interaction.
Pricing Policy Selection: A firm choosing among a set of candidate pricing rules, each mapping a customer to an offered price, need to identify the rule that yields the largest expected revenue. Here, online samples are expensive because each one requires offering a price to a real customer. A predictive model or an LLM-based scorer can instead produce a cheap proxy score for any customer-policy pair. Before the pricing experiment is deployed, this scorer can be queried offline on a large pool of historical contexts, so that the average proxy score of each policy is estimated to high accuracy. The correlation between the proxy score and the realized revenue, however, is not available offline and can only be estimated once the experiment is running. Moreover, because pricing decisions must be made in real time, only a few proxy evaluations, can be afforded for each realized customer interaction.
These two examples share the same structure. The proxy is correlated with the real reward and is therefore useful for variance reduction, serving to identify the best arm with fewer costly observations. Underlying both examples is an asymmetry between two kinds of proxy information: proxy-only information is cheap and abundant before the experiment begins, whereas reward-paired proxy information is scarce and must be acquired online. On the offline side, proxy-only scores can be generated in large quantities before the experiment by repeatedly querying the generator or predictive model at essentially no cost, so the marginal proxy mean can be estimated to arbitrary accuracy and treated as known. This does not require the proxy to be calibrated to the true reward mean. On the online side, the proxy information that actually reduces variance is the part paired with the same unit that generates the realized reward, and such paired evaluations are far more costly than offline ones, limited by latency and scalability.
This asymmetry is the source of both the opportunity and the difficulty in our setting. The opportunity is clear: had the reward-proxy correlation been known in advance, a classical control-variate adjustment by [1] would remove the predictable part of the reward noise, leaving each arm with a smaller residual variance and saving a large number of online samples in proxy-aided best-arm identification. The difficulty is that this correlation, unlike the proxy mean, must be estimated from the same costly reward-proxy paired samples that the identification is trying to economize. This naturally suggests learning the correlation as a first step and reducing the residual noise accordingly. However, aiming at the correlation is myopic. Estimating the correlation and proceeding with the estimated result may trust the observed reward more than warranted and commit to a decision before the arms are reliably separated, which jeopardizes correctness and not merely efficiency. In fact, simply being cautious about the correlation still does not repair this, because a guarantee on the correlation translates poorly into a guarantee on the variance reduction, and the translation degrades exactly when the proxy is most informative. In this paper, we show that such a cheap but unreliable proxy can be turned into a provable reduction in the number of costly online samples needed for best-arm identification.
1.1 Main Contribution↩︎
In this work, we formulate proxy-augmented fixed-confidence best-arm identification and propose PROBE, PRoxy OLS for Best-arm Exploration, an algorithm that turns an unreliable but cheap proxy into provable
savings in costly online samples with reward-proxy correlation. Our main contributions are summarized as follows.
In Section 2, we formalize a best-arm identification problem in which each costly reward pull is paired with a cheap proxy score whose marginal mean is known but reward-proxy correlation is not. We show that a control-variate adjustment turns this model into a heteroscedastic BAI problem, in which the proxy reduces the effective noise variance of each arm from one to the residual variance. This yields a clean oracle benchmark: were the correlation \(\rho\) known, the sample complexity would improve by the arm-wise factor \(1-\rho^2\). We further show that this single-proxy model is without loss of generality, as a batch of repeated proxy queries at one pull reduces exactly to a single averaged proxy.
In Section 3, we design
PROBE(PRoxy OLS for Best-arm Exploration), a phase-elimination algorithm that learns the variance reduction while performing identification. The central difficulty is that the correlation, and hence the variance reduction it implies, is unknown and must be learned from the same costly samples; a plug-in estimate of the reduced variance is anti-conservative and can compromise correctness.PROBEinstead maintains a one-sided residual-variance upper certificate obtained from an ordinary least squares fit, whose exact \(\chi^2\) law makes the certificate valid regardless of the unknown correlation. A calibration stage and monotone updates keep the certificate within a factor \(1+\kappa\) of the true residual variance throughout, and a deliberate one-round lag decouples mean estimation from variance calibration so that the concentration of mean estimates remain Gaussian.In Section 4, we prove that
PROBEis \(\delta\)-PAC and, with high probability, attains the oracle gap- and variance-dependent sample complexity up to a multiplicative factor \(1+\kappa\) and an additive calibration cost of order \(\widetilde{O}(K/\kappa^2)\) that is independent of the gaps and correlations. The parameter \(\kappa\) admits a clean interpretation as the price of not knowing the correlation, trading the tightness of the variance certificate against the cost of establishing it. Fixing \(\kappa\) to a small constant recovers the known-correlation oracle rate up to constants. We also show thatPROBEextends to the \((\epsilon,\delta)\)-PAC setting with a minimal change to the stopping rule, preserving the same improvement.In Sections 5 and 6, we validate
PROBEon synthetic instances and on an auto-loan pricing replay environment. A paired Gaussian benchmark shows thatPROBEclosely tracks the known-correlation oracle across the whole range of correlations, and a nonlinear environment with proxies learned by standard regression models confirms that the sample saving is governed by the induced reward-proxy correlation rather than by how the proxy is produced. In the replay study, proxy scores generated by large language models and by a structured tabular predictor reduce the fixed-confidence stopping time by up to \(61\%\) relative to a reward-only baseline, with the savings tracking the strength of the same-unit reward-proxy correlation.
1.2 Related Work↩︎
1.2.0.1 Fixed-confidence identification and pure-exploration objectives.
Our work builds on fixed-confidence best-arm identification (BAI), where an algorithm adaptively samples arms and stops with a recommendation that is correct with probability at least (1-\(\delta\)). This objective is related to ranking and selection in stochastic simulation [2]–[6] and to classical BAI algorithms based on elimination, confidence bounds, and information-theoretic sampling [7]–[14]. Related pure-exploration objectives include top-(m) identification, thresholding, PAC subset selection, and all-(\(\epsilon\))-good arms [15]–[21]. These works study identification from reward observations alone; we study how paired proxy observations reduce reward-sample complexity.
1.2.0.2 Model-generated feedback and LLM-assisted decisions.
Our motivation comes from decision problems where model-generated signals are informative but unreliable. [22] study regret minimization with correlated surrogate rewards, and related work uses biased offline or model-assisted data to improve bandit learning [23], [24]. In experimentation, [25] study BAI with LLM judges audited by human feedback, and [26] use LLMs to cold-start content experiments. Marketing and economics work similarly treats LLM responses as synthetic data and documents both useful preference signals and systematic deviations from humans [27]–[33]. Different from the regret-minimization setting of [22], we study fixed-confidence identification, where the proxy enters through within-arm residualization and fresh-batch OLS mean estimation is separated from \(\chi^2\)-based residual-variance certification.
1.2.0.3 Structured information in BAI.
Several pure-exploration models exploit additional structure, including arm-level variances [34], linear or transductive features [35]–[40], and delayed, correlated, or covariance-aware feedback [41]–[43]. These works change the cross-arm sampling structure. Our proxy is paired with the reward from the same pull, so the useful dependence is within an arm and within a sampled unit.
1.2.0.4 Control variates and prediction-assisted inference.
Our variance reduction is a control-variate mechanism [44], related to stochastic bandits with correlated side observations for regret
minimization [45], [46]. It also connects to prediction-powered
inference and proxy-data transfer, which use predictions to improve efficiency while correcting bias [1], [47], [48]. In fixed-confidence BAI, however, the learned residual variance determines future
sampling effort. PROBE therefore uses one-sided OLS variance certificates rather than a plug-in estimate of the reward-proxy correlation.
2 Problem Formulation↩︎
In this section, we introduce the fixed-confidence best-arm identification problem and the proxy-augmented observation model studied. We first recall the classical benchmark without proxies, then discuss the oracle benchmark in which the reward-proxy correlation is known in advance.
2.1 BAI Baseline↩︎
We consider a fixed-confidence best-arm identification (BAI) problem in finite-armed bandit with \(K\) arms indexed by \(i\in[K]=\{1,\dots,K\}\). At each time step \(t=1,2,\dots\), the learner pulls an arm \(i\) for the \(l\)-th time and receives a reward \(X_{i,l}\). For each arm \(i\), the sequence \(\{X_{i,l}\}_{l\ge 1}\) are i.i.d. samples from normal distribution \(\mathcal{N}(p_i,1)\) with unknown mean \(p_i = \mathbb{E}[X_{i,1}]\in[0,1]\), and reward sequences are independent across arms with unit variance. Denote the unique best arm as \(i^*=\arg\max_{1\le i\le K}p_i\), the reward gap can be defined for each suboptimal arm \(i\ne i^*\) as \(\Delta_i=p_{i^*}-p_i\), and we further let \(\Delta_{i^*}=\min_{i\neq i^*}\Delta_i\) for notation consistency. Given \(\delta\in(0,1)\) The objective of BAI is to find the best arm \(i^*\) with probability at least \(1-\delta\). To be specific, with single input \(\delta\), we need to find an algorithm that adaptively samples arms, stops at time \(T\), and returns a recommendation \(\hat{i}\), and such algorithm is called \(\delta\)-PAC (probably approximately correct) if \(\text{Pr}(\hat{i} = i^*)\geq1-\delta\).
The sample complexity of such algorithms is typically governed by the gap-dependent quantity \(\sum_{i=1}^K \Delta_i^{-2}\). A standard fixed-confidence BAI baseline to obtain such gap-dependent guarantees is by phase elimination with exponential-gap [8], [11]. The algorithm maintains an active set of arms and runs in rounds \(r=1,2,\ldots\), with a target accuracy \(\varepsilon_r\) that decreases geometrically, for example \(\varepsilon_r=2^{-r}\). In each round, every active arm is sampled enough times, e.g., \(\tilde{O}(\varepsilon_r^{-2})\), to estimate its mean up to order \(\varepsilon_r\). The algorithm then selects a reference arm, either by a robust subroutine such as median elimination [11] or simply by the empirical maximum among active arms, and removes arms whose empirical means are worse than the reference value by a constant multiple of \(\varepsilon_r\). Intuitively, a suboptimal arm with gap \(\Delta_i\) survives only until the first round in which \(\varepsilon_r\) is smaller than a constant multiple of \(\Delta_i\). Then, such algorithm identifies the best arm with sample complexity of order \[\sum_{i=1}^K \frac{1}{\Delta_i^2} \log\left(\frac{1}{\delta}\log\left(\frac{1}{\Delta_i}\right)\right), \label{eq:PE-bound-base}\tag{1}\] up to universal constants and lower-order logarithmic terms.
2.2 BAI with proxy scores↩︎
We now augment each reward observation with a proxy score. At each time arm \(i\) is pulled (for the \(l\)-th time), the learner observes a paired reward \((X_{i,l},Y_{i,l})\), where \(X_{i,l}\) is the noisy version of the unknown real reward \(p_i\) and \(Y_{i,l}\) is the proxy score (or reward) on the the same arm \(i\). The proxy score may be produced by a predictive model, a simulator, or an AI evaluator, e.g., a pre-trained or fine-tuned AI models like large language models (LLMs). For each arm \(i\), the pairs \(\{(X_{i,s},Y_{i,s})\}_{s\ge 1}\) are i.i.d., and the sequences are independent across arms.
The proxy score is informative through its statistical dependence with the true reward: it can help estimate \(p_i\) by its correlation with \(X_{i,l}\), even if its own mean is not calibrated to the reward mean. Formally, for each arm \(i\), the reward and proxy score as a bivariate Gaussian pair [22] as \[\begin{pmatrix} X_{i,l}\\ Y_{i,l} \end{pmatrix} \sim \mathcal{N}\!\left( \begin{pmatrix} p_i\\ m_i \end{pmatrix}, \begin{pmatrix} 1 & \rho_i\\ \rho_i & 1 \end{pmatrix} \right), \qquad \rho_i\in(-1,1). \label{eq:proxy-model}\tag{2}\] Here, \(p_i=\mathbb{E}[X_{i,1}]\) is the unknown reward mean to be learned, while \(m_i=\mathbb{E}[Y_{i,1}]\) is known and used only to center the proxy score, e.g., by define \(W_{i,l} = Y_{i,l}-m_i\) we have \(W_{i,l}\) with zero mean. The correlation coefficient \(\rho_i=\text{Cov}(X_{i,l},Y_{i,l})\) captures the quality of the proxy information, where larger \(|\rho_i|\) means that the centered proxy \(W_{i,l}\) carries more information about the reward noise. Importantly, we do not require \(m_i=p_i\), or even that \(m_i\) be close to \(p_i\). The proxy may therefore be biased as a direct surrogate for reward, as long as its centered variation is informative about \(X_{i,l}\).
Remark 1 (Proxy asymmetry between offline and online). The formulation in 2 reflects the offline-online asymmetry in proxy scores highlighted in Section 1, which makes the proxy mean \(m_i\) treated as known while leaving the correlation \(\rho_i\) to be learned.
Known proxy mean. Before the online phase, proxy calls are cheap, so the generator or predictive model can be queried repeatedly on each arm to estimate the marginal mean \(m_i\) to arbitrary accuracy, and we therefore treat \(m_i\) as known. The correlation \(\rho_i\), by contrast, is a joint property of the reward and the proxy on the same online unit due to shared underlying features or explicit contexts, and can only be estimated from the costly paired samples drawn during the online phase.
Single proxy per pull. During the online phase, latency and concurrency constraints permit only a restricted number of proxy calls per interaction. Assuming a single proxy per pull is nonetheless without loss of generality: when the proxy is repeatedly called independently, they can be averaged into one control variate, as we formalize in Proposition 3.
In bandit problems, the estimation of mean rewards plays an important role. With the help of proxy scores, the variance of estimation can be reduced through a control-variate adjustment [44]. To be specific, for each arm \(i\) with unknown mean \(p_i\), we can define the adjusted observation \[X_{i,l}-\beta(Y_{i,l}-m_i) \label{eq:control95variate}\tag{3}\] based on observed arm reward \(X_{i,l}\) and proxy score \(Y_{i,l}\) for any coefficient \(\beta\in\mathbb{R}\).
Proposition 1. For any coefficient \(\beta\in\mathbb{R}\), the adjusted observation with proxy score is unbiased with variance \[\text{Var}(X_{i,l}-\beta(Y_{i,l}-m_i)) = 1+\beta^2-2\rho_i\beta.\] The variance is minimized at \(\beta=\rho_i\), in which case \[\text{Var}(X_{i,l}-\rho_i(Y_{i,l}-m_i)) = 1 - \rho_i^2=:v_i.\]
Note that in current model 2 and 3 , we have assumed a single proxy per pull. Although in real world application, proxy calls can be repeated more than once, if the samples share the same underlying feature or explicit context, more proxy scores are redundant without additional correlation information. Otherwise, the averaged proxy is a sufficient summary of the whole proxy vector, discussed in detail in Section 10 in the appendix.
If \(\rho_i\) is known before sampling, then the control-variate adjusted random variable \(X_{i,l}-\rho_i(Y_{i,l}-m_i)\) is a Gaussian observation with mean \(p_i\) and residual variance \(v_i=1-\rho_i^2\). Consequently, the fixed-confidence best-arm identification problem with known-correlation proxy reduces to a heteroscedastic BAI problem in which arm \(i\) has noise variance \(v_i\) in place of one. To be specific, the variance of each arm \(i\) is reduced from \(1\) to \(1-\rho_i^2\). Rerunning the phase-elimination argument behind 1 with each arm’s sampling times scaling by its own residual variance \(v_i\), the baseline complexity with known proxy becomes \[\sum_{i=1}^K \frac{1-\rho_i^2}{\Delta_i^2} \log\left(\frac{1}{\delta}\log\left(\frac{1}{\Delta_i}\right)\right), \label{eq:PE-bound-known}\tag{4}\] up to universal constants and lower-order logarithmic terms, and serves as the oracle benchmark. This shows the informative value of a proxy: when a proxy can explain part of the reward fluctuation for the same sampled unit, the variance of remaining noise reduces from one to \(1-\rho_i^2\) after subtracting this predictable component. As how much fluctuation is explained can be measured by the correlation \(\rho_i\), a highly correlated proxy improves BAI by reducing the residual uncertainty in the estimation of mean reward.
Despite the \(1-\rho_i^2\) reduction in sample complexity given in 4 , the main challenge is that the correlation \(\rho_i\), and hence the residual variance \(1-\rho_i^2\) that the oracle schedule relies on, is not known in advance. As we will discuss in Remark 2, estimating \(\rho_i\) and plugging it into \(1-\hat{\rho}_i^2\) is not enough as \(|\rho_i|\) might be overestimated and the sampling is not enough before an elimination decision. A conservative bound on \(\rho_i\) cannot yields a reliable certificate for the residual variance either especially when \(\rho_i\) is large. Our algorithm instead estimates and certifies the residual variance directly through a linear-regression formulation, which we develop next.
3 PROBE: Proxy OLS for Best-arm Exploration↩︎
In this section, we propose our algorithm PROBE that learns the control-variate coefficient while maintaining conservative residual-variance upper certificates. The algorithm uses ordinary least squares (OLS) to estimate the reward mean and
the reward-proxy correlation. The residual-variance estimate is then converted into a residual-variance upper certificate that determines future batch sizes.
3.1 OLS Preliminaries↩︎
Under 2 , and plugging into \(W_{i,l} = Y_{i,l}-m_i\), we note that control-variate adjusted observation in 3 of each pull of arm \(i\) can be written as the regression model \[X_{i,l}=p_i+\rho_i W_{i,l}+\varepsilon_{i,l}, \qquad W_{i,l}\sim \mathcal{N}(0,1),\qquad \varepsilon_{i,l}\sim \mathcal{N}(0,v_i),\] where \(v_i=1-\rho_i^2\) and \(W_{i,l}\) is independent of \(\varepsilon_{i,l}\). Here, \(\rho_i\) is the slope, and \(p_i\) is the intercept in OLS. Fix arm \(i\) and a fresh batch of size \(s\ge 3\), with observations \(\{(X_{i,l},W_{i,l})\}_{l=1}^s\), the OLS estimator is \[(\hat{p}_i,\hat{\rho}_i)=\arg\min_{a,b}\sum_{l=1}^s (X_{i,l}-a-bW_{i,l})^2. \label{eq:ols}\tag{5}\] Equivalently, with \(\bar X_i=s^{-1}\sum_{l=1}^s X_{i,l}\) and \(\bar W_i=s^{-1}\sum_{l=1}^s W_{i,l}\), \[\hat{\rho}_i= \frac{\sum_{l=1}^s(X_{i,l}-\bar X_i)(W_{i,l}-\bar W_i)}{\sum_{l=1}^s(W_{i,l}-\bar W_i)^2}, \qquad \hat{p}_i=\bar X_i-\hat{\rho}_i \bar W_i.\] Let \[h_i=\frac{1}{s}+\frac{\bar W_i^2}{\sum_{l=1}^s(W_{i,l}-\bar W_i)^2},\] and define the unbiased OLS residual-variance estimator for arm \(i\) by \[\hat{v}_i=\frac{1}{s-2}\sum_{l=1}^s(X_{i,l}-\hat{p}_i-\hat{\rho}_i W_{i,l})^2. \label{eq:ols-var}\tag{6}\] Then, the OLS estimator \(\hat{p}_i\) and \(\hat{v}_i\) satisfy the following statistical backbone of the algorithm.
Proposition 2. For any fixed arm \(i\), conditional on \(W_{i,1},\ldots,W_{i,s}\), the OLS intercept satisfies \[\hat{p}_i-p_i\sim \mathcal{N}(0,v_i h_i).\] Moreover, \[\frac{(s-2)\hat{v}_i}{v_i}\sim \chi^2_{s-2}.\]
Remark 2 (Estimate \(\hat{v}_i\) instead of \(\hat{\rho}_i\)). The \(\chi^2\)-distribution of \(\hat{v}_i/v_i\) plays an important role in our algorithm that bypasses the challenge caused by estimating \(\hat{\rho}_i\) to be plugged into \(v_i\). To be specific, if we only estimate \(\hat{\rho}_i\), then
Overestimating \({\rho}_i\) is problematic: When \(\rho_i\) (assumed to be positive without loss of generality) is overestimated by an error \(e>0\), i.e., \(\hat{\rho}_i=\rho_i+e\), we have \(1-\hat{\rho}_i^2\approx v_i-2\rho e < v_i\). Because the correctness of \(\delta\)-PAC usually requires at least \(\tilde{O}(v_i/\varepsilon_r^2)\) samples to distinguish a reward gap of width \(\varepsilon_r\), the sample size based on \(1-\hat{\rho}_i^2\) is less than required and the correctness then collapses due to lack of confidence. The underlying problem of naive plug-in is that \(1-\hat{\rho}_i^2\) is biased estimator for the residual variance \(v_i\), i.e., \(\mathbb{E}[1-\hat{\rho}_i^2]=1-\rho_i^2-\text{Var}(\hat{\rho}_i)<v_i\).
Conservative LCB of \(\hat{\rho}_i\) also fails: If we simply use a lower confidence bound of \(\hat{\rho}_i\) to avoid overestimating, sample complexity might explode when the proxy is informative. For example, we use \(\hat{\rho}_i-\text{rad}\) to plug into the residual variance where \(\text{rad}=\tilde{O}\left(s^{-1/2}\right)\) is the confidence radius under sample size \(s\). With similar order of confidence provided by \(\text{rad}\), \(\rho_i-\text{rad}/2 \leq \hat{\rho}_i \leq \rho_i+\text{rad}/2\) holds and implies \(\hat{\rho}_i-\text{rad}\leq\rho_i-\text{rad}/2\). Because we still want \(1-(\hat{\rho}_i-\text{rad})^2\) close enough to the real residual variance \(v_i\) from above (otherwise it cannot reduce sample complexity and might collapse to the case without proxy), we require, e.g., \(1-(\hat{\rho}_i-\text{rad})^2\leq (1+\kappa)v_i\) then \(1-({\rho}_i-\text{rad}/2)^2\leq (1+\kappa)v_i\). This implies \(s\geq\tilde{O}(\rho_i^2\kappa^{-2}v_i^{-2})\). When \(\rho_i\) is close to 1 and \(v_i\) is close to 0, the sample size \(s\) diverges to infinite for any fixed constant \(\kappa\).
Instead, the advantages of estimation and concentration on \(\hat{v}_i\) are
\(\hat{v}_i\) is an unbiased estimator for \(v_i\).
The \(\chi^2\)-distribution provides a scale-free concentration for \(\hat{v}_i/v\).
3.2 The PROBE Algorithm↩︎
For a batch of size \(s\ge 3\), define the residual-variance upper certificate given residual-variance estimation \(\hat{v}\) as \[U(s,\hat{v},\eta)= \begin{cases} \displaystyle \frac{\hat{v}}{1-2\sqrt{\frac{\log(1/\eta)}{s-2}}}, &\text{if }s-2>4\log(1/\eta),\\[2ex] +\infty, &\text{otherwise.} \end{cases} \label{eq:U95def}\tag{7}\] The lower tail of the \(\chi^2\)-distribution and Proposition 2 imply that \(v_i\le U(s,\hat{v}_i,\eta)\) with probability at least \(1-\eta\). The algorithm uses this residual-variance upper certificate rather than the plug-in estimate \(\hat{v}_i\).
We now present PROBE, short for PRoxy OLS for Best-arm Exploration, in Algorithm 1. The algorithm has two parts. First, it draws a calibration batch from every arm to initialize a
conservative upper bound \(U_{i,0}\) on the residual variance \(v_i\). Second, it runs phased elimination. In each round \(r\), active arms with smaller
residual-variance upper certificates \(U_{i,r}\) receive smaller fresh batches, and the OLS intercept from the fresh batch is used as the reward-mean estimate for elimination.
Remark 3 (Key design choices of PROBE). The correctness and efficiency of PROBE rest on the interplay of three design choices: a residual-variance upper certificate that replaces the plug-in estimate, a
monotone update that keeps this certificate valid throughout, and a one-round lag that decouples mean estimation from variance calibration.
Residual-variance upper certificate. The upper certificate \(U_{i,r}\) defined in 7 is pivotal because it sandwiches the residual variance as \(v_i\le U_{i,r-1}\le(1+\kappa)v_i\). The lower guarantee \(v_i\le U_{i,r-1}\) delivers correctness: the round-\(r\) batch is then large enough to estimate the OLS intercept \(\hat{p}_{i,r}\) to the target accuracy \(\varepsilon_r\), so no arm is eliminated before its mean has been estimated reliably. The upper guarantee \(U_{i,r-1}\le(1+\kappa)v_i\) delivers efficiency: the batch size exceeds the oracle allocation \((1-\rho_i^2)\varepsilon_r^{-2}\) by at most a factor \(1+\kappa\), so the total sample complexity stays controlled. Using a plug-in estimate \(\hat{v}_i\) in place of \(U_{i,r}\) would forfeit the lower guarantee, and hence correctness, for the reasons detailed in Remark 2.
Monotone certificate updates. The sandwich \(v_i\leq U_{i,r-1}\leq(1+\kappa)v_i\) on residual variance is established by induction from the calibration stage. Drawing \(t_{\rm cal}=2+\lceil64\log(8K/\delta)/\kappa^2\rceil\) samples per arm gives \(v_i\le U_{i,0}\le(1+\kappa)v_i\) at \(r=0\) for all arms with high probability, using the \(\chi^2\) concentration of \(\hat{v}_i/v_i\) in Proposition 2. Each subsequent batch produces a fresh residual-variance upper certificate \(\widetilde{U}_{i,r}\), and the monotone update \(U_{i,r}=\min\{U_{i,r-1},\widetilde{U}_{i,r}\}\) preserves both sides of the sandwich: taking the minimum of valid upper bounds keeps the lower guarantee \(v_i\le U_{i,r}\), while never worsening the initial factor-\((1+\kappa)\) tightness. Hence \(v_i\le U_{i,r}\le(1+\kappa)v_i\) holds throughout the algorithm.
One-round lag between mean and variance certificate. PROBE sizes the round-\(r\) batch through \(U_{i,r-1}\), computed from data collected before round
\(r\), while it estimates \(\hat{p}_{i,r}\) only from the fresh round-\(r\) observations. This deliberate lag decouples estimation on mean from
residual-variance upper certificate: since the accuracy of \(\hat{p}_{i,r}\) is certified by the fixed \(U_{i,r-1}\ge v_i\) from earlier rounds, Proposition 2 applies directly conditional on the fresh proxy values, and concentration of \(\hat{p}_{i,r}\) remains exactly Gaussian. Reusing a single batch for
both purposes would instead couple them and complicate the concentration of \(\hat{p}_{i,r}\) to follow a Student-\(t\) law rather than a Gaussian one, as in [22]. Instead, the lag yields a sharper concentration without changing the leading sample complexity as the geometrically growing phase lengths \(\varepsilon_r^{-2}\) absorb the one-round delay.
4 Sample Complexity↩︎
In this section, we show that PROBE matches the gap-dependent sample complexity reduced by the residual variance \(1-\rho_i^2\) up to the multiplicative factor \(1+\kappa\)
and an initial calibration cost.
4.1 \(\delta\)-PAC Analysis↩︎
Theorem 1. Under the proxy observation model 2 , PROBE is \(\delta\)-PAC. Moreover, with probability at least \(1-\delta\), its total number of arm pulls is at most \[\begin{align}
T\leq K\left(2+\left\lceil\frac{64\log(8K/\delta)}{\kappa^2}\right\rceil\right) +
\sum_{i=1}^K \Bigg(&\frac{1536(1+\kappa)(1-\rho_i^2)}{\Delta_i^2} \log\left(\frac{12K}{\delta}\log_2\left(\frac{3}{2\Delta_i}\right)\right)\\
&+\log\left(\frac{3}{2\Delta_i}\right) \left(64\log\left(\frac{12K}{\delta}\log_2\left(\frac{3}{2\Delta_i}\right)\right) +6\right)\Bigg)
\end{align}
\label{eq:probe-sample-complexity}\qquad{(1)}\]
The proof works on a clean event under which two properties hold uniformly over arms and rounds: every residual-variance upper certificate is valid and tight enough as \[v_i\le U_{i,r}\le (1+\kappa)v_i,\] and every OLS intercept used for elimination is accurate as \[|\hat{p}_{i,r}-p_i|\le \varepsilon_r/4.\] The clean event follows from three parts. First, the calibration event gives \(v_i\le U_{i,0}\le (1+\kappa)v_i\), then the monotone updates and the \(\chi^2\)-distribution gives \(v_i\le U_{i,r}\le (1+\kappa)v_i\) for all round \(r\geq0\) with induction. Second, the OLS intercept estimation \(\hat{p}_{i,r}\) is concentrated at \(p_i\). Third, the leverage event \(h_{i,r}\le 2/t_{i,r}\) happens (established by the first three terms in \(t_{i,r}\)), which further (by the final term in \(t_{i,r}\)) implies \(|\hat{p}_{i,r}-p_i|\le \varepsilon_r/4\) with the OLS intercept estimation.
Conditioned on this clean event which happens with probability at least \(1-\delta\), correctness follows from the standard phase-elimination argument. The best arm is never removed: if \(i^*\in S_r\), then the empirical reference arm \(i_r\) satisfies \[\hat{p}_{i_r,r}\le p_{i^*}+\varepsilon_r/4, \qquad \hat{p}_{i^*,r}\ge p_{i^*}-\varepsilon_r/4,\] so \(\hat{p}_{i_r,r}-\hat{p}_{i^*,r}\le \varepsilon_r/2<\varepsilon_r\). Conversely, a suboptimal arm \(i\) is eliminated once \(\varepsilon_r<2\Delta_i/3\) because \[\hat{p}_{i^*,r}-\hat{p}_{i,r} \ge \Delta_i-\varepsilon_r/2 > \varepsilon_r.\] Since the empirical reference arm has value at least \(\hat{p}_{i^*,r}\), arm \(i\) is removed if it is still active.
It remains to count samples. A suboptimal arm \(i\) is sampled only until the first round in which \(\varepsilon_r<2\Delta_i/3\), namely \(r\geq O(\log_2(1/\Delta_i))\), and the best arm is sampled only until the analogous round with \(\Delta_{i^*}=\min_{i\ne i^*}\Delta_i\). Conditioned on the clean event, arm \(i\) is sampled for \[t_{\rm cal} + \sum_{r=1}^{\log_2(1/\Delta_i)}t_{i,r} \leq \tilde{O}\left( \frac{1}{\kappa^2} + \sum_{r=1}^{\log_2(1/\Delta_i)} \frac{(1+\kappa)(1-\rho_i^2)}{4^{-r}} \right) \leq \tilde{O}\left( \frac{1}{\kappa^2} + \frac{(1+\kappa)(1-\rho_i^2)}{\Delta_i^2} \right)\] Summing this display for every arm \(i\) yields the bound in ?? . The full proof is deferred to Section 8 in the appendix.
Note that, Algorithm 1 can be easily extended to the \((\epsilon,\delta)\)-PAC setting. In the previous \(\delta\)-PAC setting, we focus on exact best-arm
identification, where the algorithm must return \(i^*\) with probability at least \(1-\delta\). In many applications, it is sufficient to identify an arm whose mean reward is within a
prescribed accuracy level \(\epsilon>0\) of the best arm. Formally, an algorithm is called \((\epsilon,\delta)\)-PAC if its recommendation \(\hat{i}\)
satisfies \[\text{Pr}(p_{\hat{i}}\ge p_{i^*}-\epsilon) \geq 1-\delta.\] To accommodate this relaxation, the modification of PROBE is minimal. The only change is the stopping rule: the algorithm is terminated
at round \(\min\{r:\varepsilon_r<{2\epsilon}/{3}\}\) even if \(|S_r|>1\) and an arbitrary arm in \(S_r\) is returned.
The sample complexity is obtained from Theorem 1 by replacing each gap \(\Delta_i\) with the effective gap \(\max\{\Delta_i,\epsilon\}\). In particular, with probability at least \(1-\delta\), the total number of arm pulls is at most \[\begin{align} T_{\epsilon}\leq K & \left(2+ \left\lceil\frac{64\log(8K/\delta)}{\kappa^2}\right\rceil\right) + \sum_{i=1}^K \Bigg( \frac{1536(1+\kappa)(1-\rho_i^2)}{\max\{\Delta_i,\epsilon\}^2} \log\left( \frac{12K}{\delta} \log_2\left(\frac{3}{2\max\{\Delta_i,\epsilon\}}\right) \right)\\ &+ \log_2\left(\frac{3}{2\max\{\Delta_i,\epsilon\}}\right) \left( 64\log\left( \frac{12K}{\delta} \log_2\left(\frac{3}{2\max\{\Delta_i,\epsilon\}}\right) \right)+6 \right) \Bigg). \end{align} \label{eq:probe-ep-delta}\tag{8}\] Ignoring constants and lower-order logarithmic terms, 8 gives \[\widetilde{O}\left( \frac{K}{\kappa^2} + \sum_{i=1}^K \frac{(1+\kappa)(1-\rho_i^2)}{\max\{\Delta_i,\epsilon\}^2} \right),\] preserving the same variance-adaptive improvement as the exact-identification guarantee.
4.2 Comparison with Benchmarks↩︎
We then compare the sample complexity of PROBE in Theorem 1 with two benchmarks given in 1 and 4 . The sample complexity in Theorem 1 can be expressed in leading dependence \[\sum_{i=1}^K \left(
\frac{(1+\kappa)(1-\rho_i^2)}{\Delta_i^2}\log\left(\frac{K}{\delta}\log\left( \frac{1}{\Delta_i}\right) \right) + \frac{\log(K/\delta)}{\kappa^2} \right)\] where the first term corresponds to the phase elimination stage and the second terms
corresponds to the calibration stage. Recall that the sample complexity of the fixed-confidence best-arm identification without proxy has leading dependence \[\sum_{i=1}^K
\frac{1}{\Delta_i^2}
\log\left(\frac{1}{\delta}\log\left(\frac{1}{\Delta_i}\right)\right).\] and the one with proxy and known correlation \(\rho_i\) has leading dependence \[\sum_{i=1}^K
\frac{1-\rho_i^2}{\Delta_i^2}
\log\left(\frac{1}{\delta}\log\left(\frac{1}{\Delta_i}\right)\right),\] PROBE pays only a multiplicative factor \(1+\kappa\) in the main gap-dependent term and an additive calibration cost of order \(\tilde{O}(\kappa^{-2})\).
The parameter \(\kappa\) captures the exploration and exploitation tradeoff facing unknown correlation: a tradeoff between the accuracy in residual variance \(v_i\) estimation
(exploitation) and the sample size in calibration for this accuracy (exploration). On one hand, \(\kappa\) quantifies the relative accuracy with which the algorithm learns an upper certificate \(U_{i,r}\) for \(v_i\). To be specific, PROBE maintains \(U_{i,r}\le (1+\kappa)v_i\) throughout the algorithm and provides tolerance towards
overestimating \(U_{i,r}\). This tolerance in upper bound further ensures that the batch size is not much larger than the oracle batch size based on the unknown residual variance \(v_i=1-\rho_i^2\) by a factor of \(1+\kappa\) during elimination. On the other hand, a lower tolerance in estimation accuracy requires more samples by an additive factor of \(\kappa^{-2}\) to establish the certificate reliably during calibration. In a word, \(\kappa\) represents the price of not knowing \(\rho_i\) and quantifies the
tradeoff between exploring and exploiting this correlation.
Thanks to the \(\chi^2\)-distribution of the residual variance, the additive exploration term is independent of the gaps \(\Delta_i\) and the correlations \(\rho_i\). We can fix it as a small constant value then the instance enjoys the same leading dependence as the known-\(\rho_i\) oracle \(\widetilde{O}\left( {(1-\rho_i^2)}/{\Delta_i^2}\right)\) for each arm \(i\), up to a constant multiplicative factor.
5 Numerical Experiments on Synthetic Data↩︎
We use synthetic experiments to isolate the statistical mechanism behind PROBE: a proxy is useful when its centered within-arm variation explains reward noise, even if its mean is biased or rank-reversed. We study two settings. The first is
a paired Gaussian benchmark in which the proxy is prescribed by hand and its correlation with the reward is a controlled parameter. The second is a nonlinear environment in which the proxy is a learned predictor, so that its correlation is an outcome of
model fit rather than a chosen quantity. Throughout, the performance metric is the fixed-confidence stopping time, measured by the total number of online reward pulls before stopping.
5.1 Gaussian Instance↩︎
Setup: We consider a synthetic Gaussian instance with \(K=8\) arms. The reward means are \((1.00,0.90,0.84,0.76,0.70,0.62,0.54,0.44)\), and the proxy means are \((0.00,0.25,0.50,0.75,1.00,1.25,1.50,1.75)\), following model 2 . Thus, the proxy means are deliberately ordered in the opposite direction from the reward means, so that the proxy cannot be useful through mean ranking alone. For each arm \(i\), we generate paired Gaussian observations with reward variance one, proxy variance one, and common reward-proxy correlation \(\rho\). We vary \(\rho \in \{0,0.2,0.4,0.6,0.8,0.9\}\).
We compare three classes of methods. The first is PROBE without proxy, which ignores the proxy and uses reward observations only and the calibration step in Algorithm 1 is skipped. The second is a
known-correlation oracle, which residualizes the reward using the true value of \(\rho\) and runs the same phase-elimination procedure with residual variance \(1-\rho^2\). The third is
PROBE with unknown correlation, and we test among \(\kappa \in \{0.1,0.2,0.5,1.0\}\). We set \(\delta=0.05\), repeat each configuration 1000 times, and report the average
stopping time and its ratio relative to PROBE witout proxy at the same value of \(\rho\).
Results: The results are shown in Figure 2. When \(\rho=0\), the proxy carries no information about the reward noise. In this case, proxy-assisted
PROBE pays a learning overhead. This is consistent with the role of calibration, and when there is no variance reduction to learn, the proxy can only add estimation cost leading to a sample complexity even higher than not using proxy. As \(\rho\) increases, the proxy becomes more useful, the stopping time decreases, and exceed the no proxy baseline after \(\rho\geq0.4\). Note that, although for all \(\rho\) the sample complexity of PROBE is higher than the known-correlation baseline, due to the multiplicative factor \(1+\kappa\), PROBE is very close to the
known-proxy baseline and the \(1-\rho^2\) theoretical reference.
These results confirm that the improvement comes from same-unit reward-proxy covariance rather than from the proxy mean itself. Although the proxy means are non-rank-preserving, PROBE substantially reduces the fixed-confidence stopping time
once the proxy explains a nontrivial portion of the reward noise. The empirical error rates remain close to the nominal level, with the largest observed error rate below \(5\%\) across all configurations.
5.2 Nonlinear Instance↩︎
Setup: We consider \(K=5\) arms, where arm \(i\) corresponds to an action \(a_i\) on an evenly spaced grid, \(a_i\in\{-1,-0.5,0,0.5,1\}\). Each sampled unit carries a latent feature vector \(u=(u_1,\dots,u_5)\) with independent standard normal entries. Given action \(a\) and features \(u\), the reward is driven by the nonlinear signal \[g(u,a)=0.9\sin(u_1+a)+0.55\,a\,u_2+0.35\,(u_3^2-1)(a+0.4)+0.4\cos(u_4-0.5a)+0.2\,u_5, \label{eq:ml-signal}\tag{9}\] which mixes sinusoidal, quadratic, and action-feature interaction terms. The realized reward of arm \(i\) on a unit is \[X_i = p_i + \frac{g(u,a_i)+0.45\,\xi-\mu_{a_i}}{\sigma_{a_i}},\qquad \xi\sim\mathcal{N}(0,1), \label{eq:ml-reward}\tag{10}\] where \(\xi\) is independent observation noise, and \((\mu_{a_i},\sigma_{a_i})\) are the mean and standard deviation of \(g(u,a_i)+0.45\xi\), estimated once from a large pretabulated sample. Subtracting \(\mu_{a_i}\) and dividing by \(\sigma_{a_i}\) standardizes each arm to unit reward variance and mean \(p_i\), matching the unit-variance model of the analysis. The arm means \(p=(1.00,0.90,0.84,0.76,0.70)\) decrease across arms, so arm 1 is the best. As in the Gaussian benchmark, the proxy means \(m=(0.00,0.30,0.60,0.90,1.20)\) increase across arms and are therefore non-rank-preserving.
The proxy is a regression model fitted to predict the reward from the feature map \(\phi(u,a)=(u,a,au_1,au_2,a^2)\). This map deliberately omits the terms through which \(g\) actually
depends on the features, namely the sinusoid \(\sin(u_1+a)\), the quadratic \(u_3^2\), the cosine \(\cos(u_4-0.5a)\), and the linear term \(u_5\). No model built on \(\phi\) can therefore reproduce the reward exactly, and the achievable reward-proxy correlation is limited by how much of \(g\) each
model recovers from \(\phi\). We fit three standard models of increasing flexibility on an offline training set of \(12{,}000\) units: linear regression, a depth-limited decision tree, and
gradient boosting. Each fitted model then scores a separate pool of \(40{,}000\) units per arm, and each score is paired with the realized reward on the same unit. We use \(\delta=0.05\) and
\(\kappa=1.0\), repeat each configuration 120 times, and report the stopping-time ratio relative to PROBE without proxy, together with the mean and minimum armwise correlations of each proxy.
Results: The three models differ only in how much of the nonlinear structure omitted from the feature map \(\phi(u,a)\) they can recover, and Figure 3 shows that this is exactly the order in which they shorten the stopping time. Gradient boosting, the most flexible model, recovers the most of the omitted structure in 9 and is the most strongly aligned with the reward, so it attains the largest sample saving. The decision tree and linear regression recover less and are correspondingly less aligned, so their savings are smaller and close to each other. All three proxies still identify the best arm in every repetition. The ordering matches the Gaussian benchmark and the theory: a proxy that explains more within-arm reward variation leaves a lower OLS residual variance and a shorter fixed-confidence stopping time, regardless of how the proxy is produced or whether its mean is rank-preserving.
Figure 3: Nonlinear synthetic instance. The left panel reports the mean and minimum armwise reward-proxy correlations. The right panel reports the corresponding stopping-time ratio of PROBE relative to reward-only
PROBE..
6 Numerical Experiments on Real Data: Auto-loan Replay Experiment↩︎
We now evaluate PROBE on a pricing policy-selection task built from real auto-loan data. This is the setting of the pricing example in Section 1: a lender must decide which of two pricing rules (two arms)
earns more revenue, each online sample is a costly interaction with a real applicant, and a machine-learning model supplies a cheap proxy for the outcome. The experiment tests whether the variance reduction isolated in the synthetic study still
materializes when the proxy is an LLM or a tabular predictor, and when the reward is a realized acceptance rather than Gaussian noise.
6.1 Data and Setup↩︎
We use the “On-Line Auto Lending” data set from the Center for Pricing and Revenue Management at Columbia University [49], first studied by [50]. Each record is an auto-loan application together with the contract offered to the applicant, described by customer and loan features (e.g., FICO score, approved amount, term, and competitor rate), the offered monthly-payment margin price, and a binary label indicating whether the applicant accepted the offer. We sort the applications chronologically and split them into a historical block of the first \(50{,}000\) records and a replay pool of the next \(40{,}000\) records. The historical block is used only to build the environment and to train the proxies, and the replay pool supplies the applicant contexts encountered online.
We fit an acceptance model \(q(C,P)\) for each context \(C\) and provided price \(P\) on the historical block and treat it as the ground-truth demand environment. At each arm pull on the pricing rule, a fresh applicant context \(C\) is drawn from the replay pool, and the pricing rule \(\pi\) computes the price \(P=\pi(C)\) to offer, which together generates a Bernoulli acceptance from \(q\) and realized to 0 or 1. The realized reward is the margin revenue, equal to the offered price when the applicant accepts and zero otherwise, with mean reward \[m(\pi) = \mathbb{E}_{C\sim D_{\rm context}}[\pi(C)\cdot q(C,\pi(C))]\] as an expectation over context (applicant) distribution. Rewards are scaled to unit variance to match the normalization of the analysis. We emphasize that \(q\) defines the environment only and is never revealed to the proxies.
We consider a best-arm identification on the two arms which are both fixed pricing rules. The historical rule \(\pi_{\mathrm{hist}}\) re-offers the price actually observed for each applicant. The improved rule \(\pi_{\mathrm{imp}}\) instead offers, for each applicant, the revenue-maximizing price among \(\{0.8,0.9,1.0,1.1,1.2\}\) times the historical price within the empirical price support, where the maximization uses the fitted environment. Under the fitted environment, the historical rule earns expected margin revenue \(133.0\) and the improved rule earns \(177.7\), so the identification target is the improved rule. The fixed-confidence task is to recover this pricing rule (arm) with probability at least \(1-\delta\).
Each proxy predicts an applicant’s acceptance probability from the context and offered price, which we convert to a proxy revenue score by multiplying by the offered price. This score is paired with the realized reward on the same applicant-rule pull. Intuitively, the proxy helps whenever it tracks which applicants are unusually likely or unlikely to accept, so that residualizing the reward on the centered proxy removes part of the predictable, context-driven variation in revenue and leaves a smaller residual for the mean estimate. What remains after residualization, the idiosyncratic acceptance shock, cannot be predicted from the context and sets a floor on the achievable variance reduction.
We compare five variants of PROBE. The baseline is reward-only PROBE, which ignores the proxy. Three variants use LLM proxies: Qwen2.5-7B [51] with in-context learning run locally, GPT-5.5 [52] with
in-context learning through its API, and Qwen2.5-7B fine-tuned on the historical block with LoRA, whose raw acceptance-classification score is used as the proxy. The last variant uses TabPFN [53], a pretrained tabular predictor, as a strong structured (non-LLM) proxy that calibrates how
PROBE behaves when the proxy is highly accurate.
We run every variant with \(\delta=0.05\) and \(\kappa=1.0\), and average over \(3{,}000\) independent replay repetitions. Besides the mean stopping time,
we report its median and 90th percentile, the empirical correctness at stopping, the sample saving relative to reward-only PROBE, and the absolute reward-proxy correlation on each arm. We report absolute correlations because the sign of the
association is absorbed by the OLS slope, so only its strength governs the residual-variance reduction.
| Method | Mean | Median | 90% | Correct | \(|\rho_{\mathrm{hist}}|\) | \(|\rho_{\mathrm{imp}}|\) | Saving |
|---|---|---|---|---|---|---|---|
| Reward-only PROBE | 105.2 | 105.2 | 106.7 | 1.000 | – | – | 0.0% |
| Qwen2.5-7B local ICL | 101.7 | 101.8 | 103.5 | 1.000 | 0.275 | 0.011 | 3.3% |
| GPT-5.5 API ICL | 95.6 | 95.7 | 97.6 | 1.000 | 0.264 | 0.327 | 9.1% |
| Qwen2.5-7B LoRA raw CLS | 52.0 | 52.3 | 54.5 | 1.000 | 0.725 | 0.707 | 50.5% |
| TabPFN structured proxy | 40.6 | 40.8 | 43.1 | 1.000 | 0.822 | 0.760 | 61.4% |
4pt
6.2 Results↩︎
Table 1 and Figure 4 report the results. Every variant identifies the improved rule in all \(3{,}000\) repetitions, so all methods are exactly
correct and none exhausts the sampling budget; the variants differ only in how quickly they stop. Reward-only PROBE stops after \(105.2\)k pulls on average. The two in-context-learning proxies help modestly,
cutting the stopping time to \(101.7\)k pulls (Qwen2.5-7B, \(3.3\%\) saving) and \(95.6\)k pulls (GPT-5.5, \(9.1\%\)). The
fine-tuned and structured proxies help far more: Qwen2.5-7B LoRA reaches \(52.0\)k pulls (\(50.5\%\) saving) and TabPFN reaches \(40.6\)k pulls (\(61.4\%\)). Figure 4 shows the same ordering as a uniform leftward shift of the entire stopping-time distribution, not merely of its mean.
The savings track the strength of the reward-proxy correlation, exactly as the theory predicts. The in-context proxies are weakly and unevenly aligned with the reward: the local Qwen proxy correlates at \(0.275\) on the
historical arm but essentially zero (\(0.011\)) on the improved arm, so it can residualize only one of the two policies and saves little. The fine-tuned and structured proxies are strongly aligned on both arms (\(0.725\) and \(0.707\) for LoRA; \(0.822\) and \(0.760\) for TabPFN), so PROBE removes a large predictable component
of revenue from each arm and needs far fewer pulls to separate the two policies.
7 Conclusion↩︎
In this paper, we study fixed-confidence best-arm identification when each costly reward pull is paired with a cheap proxy score whose marginal mean is known but whose correlation with the reward is not. We propose PROBE, a
phase-elimination algorithm that uses the proxy as a control variate and maintains a residual-variance upper certificate, from an ordinary least squares fit, that stays valid regardless of the unknown correlation. We prove that PROBE is \(\delta\)-PAC and attains the known-correlation oracle sample complexity up to a factor \(1+\kappa\) and an additive calibration cost of order \(\widetilde{O}(K/\kappa^2)\), and that it extends to the \((\epsilon,\delta)\)-PAC setting. Experiments on synthetic instances and on an auto-loan pricing replay with large language model and
tabular proxies confirm that its sample savings scale with the strength of the reward-proxy correlation, as the theory predicts.
Our results suggest two directions for future study. First, the proxy here is fixed throughout identification; understanding how an online-refined proxy, for instance one fine-tuned as reward-proxy pairs accumulate, interacts with variance-adaptive sampling is a natural next step. Second, the same control-variate mechanism may reduce sampling cost in other pure-exploration objectives, such as top-\(m\) identification and thresholding bandits.
Appendix
8 Proof of Theorem 1↩︎
We prove the result on the clean event. We first define the clean initialization on \(U_{i,0}\) for every arm \(i\) as \[\mathcal{E}_{\rm cal} = \left\{v_i\le U_{i,0}\le (1+\kappa)v_i\right\}.\] For every arm \(i\) and round \(r\), define the leverage event \[\mathcal{E}^h_{i,r}=\left\{h_{i,r}\le \frac{2}{t_{i,r}}\right\}.\] Define the variance-certificate update event \[\mathcal{E}^U_{i,r}=\left\{v_i\le \widetilde{U}_{i,r}\right\}.\] Finally, define the mean-estimation event \[\mathcal{E}^p_{i,r} = \left\{|\hat{p}_{i,r}-p_i|\le \sqrt{2v_ih_{i,r}\log(2/\delta_r)}\right\}.\] Combining these four events, we can define the global clean event \[\mathcal{E}= \mathcal{E}_{\rm cal} \cap \bigcap_{r\ge 1}\bigcap_{i=1}^K \left(\mathcal{E}^h_{i,r}\cap\mathcal{E}^U_{i,r}\cap\mathcal{E}^p_{i,r}\right),\] where events for inactive arms are treated as the whole sample space. Then, \(\mathcal{E}\) happens with probability at least \(1-\delta\) with proof given in Section 8.1.
Lemma 1. The event \(\mathcal{E}\) implies bounded upper bound \(U_{i,r}\) for all arm \(i\) and round \(r\geq0\) as \[v_i\leq U_{i,r}\leq (1+\kappa)v_i \label{eq:U95concentrate}\qquad{(2)}\] and the concentration in estimation \(\hat{p}_{i,r}\) for all arm \(i\) and round \(r\geq1\) as \[|\hat{p}_{i,r}-p_i| \leq \varepsilon_r/4. \label{eq:p95concentration}\qquad{(3)}\] Furthermore, \(\mathcal{E}\) happens with probability \[\text{Pr}(\mathcal{E})\ge 1-\delta.\]
Then, we prove conditioned on clean event \(\mathcal{E}\). Denote for every arm \(i\), \[R_i=\min\left\{r\ge 1:\;2^{-r}<\frac{2\Delta_i}{3}\right\},\] namely, \(R_i=\lceil\log(3/(2\Delta_i))\rceil\leq2\log(3/(2\Delta_i))\) because \(\Delta_i\leq 1\).
Elimination correctness. First, the best arm is never removed conditioned on \(\mathcal{E}\). Suppose \(i^*\in S_r\). Since \(i_r\) maximizes \(\hat{p}_{i,r}\) over \(S_r\), \[\hat{p}_{i_r,r}\le p_{i_r}+\frac{\varepsilon_r}{4}\le p_{i^*}+\frac{\varepsilon_r}{4},\] and \[\hat{p}_{i^*,r}\ge p_{i^*}-\frac{\varepsilon_r}{4}.\] Hence \[\hat{p}_{i_r,r}-\hat{p}_{i^*,r}\le \frac{\varepsilon_r}{2}<\varepsilon_r.\] Therefore the elimination condition \(\hat{p}_{i^*,r}<\hat{p}_{i_r,r}-\varepsilon_r\) cannot hold. By induction, the best arm remains active in every round.
Second, every suboptimal arm is removed once the round resolution is below its gap. Let \(i\ne i^*\). If \(r\ge R_i\), since the best arm is still active, we have \[\hat{p}_{i^*,r}-\hat{p}_{i,r} \ge \left(p_{i^*}-\frac{\varepsilon_r}{4}\right)-\left(p_i+\frac{\varepsilon_r}{4}\right) = \Delta_i-\frac{\varepsilon_r}{2} > \varepsilon_r.\] where the final inequality is based on \(\varepsilon_r\le \varepsilon_{R_i}<{2\Delta_i}/{3}\). Because \(\hat{p}_{i_r,r}\ge \hat{p}_{i^*,r}\), \[\hat{p}_{i,r}<\hat{p}_{i_r,r}-\varepsilon_r.\] Thus arm \(i\) is eliminated in round \(r\geq R_i\) if it is still active. Hence suboptimal arm \(i\) is sampled only through round \(R_i\), and the best arm is sampled only through round \(R_{i^*}\).
Sample complexity. The arm pull of each \(i\) at round \(r\) is upper bounded by \[\begin{align} t_{i,r} &\leq \left\lceil\frac{64U_{i,r-1}\log({2}/{\delta_r})}{\varepsilon_r^2}\right\rceil + \max\left\{ 3, \left\lceil1+16\log\frac{2}{\delta_r}\right\rceil, \left\lceil8\log\frac{4}{\delta_r}\right\rceil \right\}\\ &\leq \frac{64(1+\kappa)(1-\rho_i^2)\log({2}/{\delta_r})}{\varepsilon_r^2} + 1 + 2 + 16\log(2/\delta_r) \end{align}\] where first inequality is based on the definition of \(t_{i,r}\), the second inequality is \(U_{i,r-1}\le (1+\kappa)v_i=(1+\kappa)(1-\rho_i^2)\) implied by \(\mathcal{E}\), and \(\log(4/\delta_r)\le 2\log(2/\delta_r)\) and \(\log({2}/{\delta_r})\ge \log(32K/\delta)>1\). Since \(\log(2/\delta_r)\) is increasing in \(r\), \[\begin{align} \sum_{r=1}^{R_i}\frac{\log(2/\delta_r)}{\varepsilon_r^2} &\le \log(2/\delta_r)\sum_{r=1}^{R_i}4^r \leq \frac{4}{3}\log(2/\delta_{R_i})4^{R_i}\\ &\leq \frac{12\log(2/\delta_{R_i})}{\Delta_i^2}, \end{align}\] where the final inequality is based on the definition of \(R_i\), where either \(R_i=1\) or \(2^{-(R_i-1)}\ge 2\Delta_i/3\) which implies \(2^{-R_i}\ge {\Delta_i}/{3}\) in both cases with \(\Delta_i\le 1\)/. We further plug \(R_i\leq2\log(3/(2\Delta_i))\) into the definition of \(\delta_{R_i}\) as \[\log\frac{2}{\delta_{R_i}} = \log\frac{32KR_i^2}{\delta} \le \log\left( \frac{128K}{\delta} \left(\log_2\left(\frac{3}{2\Delta_i}\right)\right)^2 \right) \leq 2\log\left(\frac{12K}{\delta}\log_2\left(\frac{3}{2\Delta_i}\right)\right).\] The post-calibration samples of arm \(i\) are at most \[\sum_{r=1}^{R_i}t_{i,r} \le \frac{1536(1+\kappa)(1-\rho_i^2)}{\Delta_i^2} \log\left(\frac{12K}{\delta}\log_2\left(\frac{3}{2\Delta_i}\right)\right) +\log\left(\frac{3}{2\Delta_i}\right) \left(64\log\left(\frac{12K}{\delta}\log_2\left(\frac{3}{2\Delta_i}\right)\right) +6\right).\] Adding the \(t_{\rm cal}\) calibration samples for each arm proves the stated total sample bound.
8.1 Proof of Lemma 1↩︎
To prove Lemma 1, we need two additional lemmas with proofs in Section 8.2 and 8.3.
Lemma 2. For every \(\eta\in(0,1)\), if the sample size \(s\ge \max\{3, 1+16\log(2/\eta), 8\log(4/\eta)\}\), then the leverage factor follows \[\text{Pr}\left(h\le \frac{2}{s}\right)\ge 1-\eta.\]
Lemma 3. For each arm, draw \(t_{\rm cal}\) samples and compute the OLS residual variance \(\hat{v}_0\). Define \(U_0=U(t_{\rm cal},\hat{v}_0,\delta/(8K))\), then \[\text{Pr}\left(v\le U_0\le (1+\kappa)v\right)\ge 1-\frac{\delta}{4K}.\]
We first prove the implied statement of clean event \(\mathcal{E}\). To prove ?? for all \(r\geq0\), we prove by induction. Assuming \(v_i\leq U_{i,r}\leq(1+\kappa)v_i\), because \(\mathcal{E}^U_{i,r+1}\) implies \(v_i\leq\widetilde{U}_{i,r+1}\), and recall \(U_{i,r+1} = \min\{\widetilde{U}_{i,r+1}, U_{i,r}\}\), we get \[v_i\leq U_{i,r+1}.\] On the other hand, we know \[U_{i,r+1}=\min\{\widetilde{U}_{i,r+1}, U_{i,r}\}\leq U_{i,r} \leq U_{i,r}\leq(1+\kappa)v_i.\] Starting from \(r=0\) given by \(\mathcal{E}_{\rm cal}\), we show ?? holds for all \(r\geq0\).
Because we have proved \(v_i\leq U_{i,r-1}\) holds for \(r\geq1\), \(\mathcal{E}_{i,r}^p\) implies that \[\begin{align} |\hat{p}_{i,r}-p_i| &\le \sqrt{2v_ih_{i,r}\log(2/\delta_r)} \\ &\leq \sqrt{2U_{i,r-1}h_{i,r}\log(2/\delta_r)} \\ &\leq \sqrt{\frac{4U_{i,r-1}\log(2/\delta_r)}{t_{i,r}}}\leq\frac{\varepsilon_r}{4}, \end{align}\] where the third inequality is given by \(\mathcal{E}_{i,r}^h=\{h_{i,r}\leq2/t_{i,r}\}\) and the final inequality is based on definition \(t_{i,r}\geq {64U_{i,r-1}\log({2}/{\delta_r}})/{\varepsilon_r^2}\), and then ?? is proved.
We then prove the probability of the clean event \(\mathcal{E}\). By Lemma 3 and a union bound over the \(K\) arms, \[\text{Pr}(\mathcal{E}_{\rm cal}^c)\le \frac{\delta}{4}.\] Since \(t_{i,r}\) satisfies the lower bounds required in Lemma 2, the lemma gives \[\text{Pr}\left((\mathcal{E}^h_{i,r})^c\right) \le \delta_r .\] For the variance-certificate update event \(\mathcal{E}^U_{i,r}\), given current fresh proxy variables \(\{W_{i,r,l}\}_{l=1}^{t_{i,r}}\), Proposition 2 gives \({(t_{i,r}-2)\hat{v}_{i,r}}/{v_i} \sim \chi^2_{t_{i,r}-2}\), then \[\begin{align} &\text{Pr}\left((\mathcal{E}^U_{i,r})^c\right) \leq \text{Pr}\left(v_i\geq \frac{\hat{v}_{i,r}}{1-2\sqrt{\frac{\log(1/\delta_r)}{t_{i,r}-2}}} \right)\\ =& \text{Pr}\left( \frac{(t_{i,r}-2)\hat{v}_{i,r}}{v_i} \le (t_{i,r}-2) - 2\sqrt{(t_{i,r}-2)\log(1/\delta_r)} \right) \le \delta_r, \end{align}\] where the final inequality is based on the lower-tail part of Laurent-Massart inequality. For the mean-estimation event \(\mathcal{E}_{i,r}^p\), given current fresh proxy variables \(\{W_{i,r,l}\}_{l=1}^{t_{i,r}}\), the leverage \(h_{i,r}\) is fixed, and Proposition 2 gives \(\hat{p}_{i,r}-p_i \sim \mathcal{N}(0,v_i h_{i,r})\) and implies \[\text{Pr}\left( |\hat{p}_{i,r}-p_i| > \sqrt{2v_i h_{i,r}\log\frac{2}{\delta_r}} \right) \le 2\exp\left(-\log\frac{2}{\delta_r}\right) = \delta_r .\]
Finally, by a union bound and plugging into \(\delta_r=\delta/(16Kr^2)\), we have \[\begin{align} \text{Pr}(\mathcal{E}^c) &\le \text{Pr}(\mathcal{E}_{\rm cal}^c) + \sum_{r=1}^{\infty}\sum_{i=1}^K \left(\text{Pr}((\mathcal{E}^h_{i,r})^c) + \text{Pr}((\mathcal{E}^U_{i,r})^c) + \text{Pr}((\mathcal{E}^p_{i,r})^c)\right)\\ &\leq \frac{\delta}{4} + 3K\sum_{r=1}^{\infty}\delta_r = \frac{\delta}{4} + 3K\sum_{r=1}^{\infty}\frac{\delta}{16Kr^2}\\ &\leq = \frac{\delta}{4} + \frac{3\delta}{16}\cdot2 = \frac{\delta}{4} + \frac{3\delta}{8} <\delta. \end{align}\]
8.2 Proof of Lemma 2↩︎
We have \(\bar W\sim \mathcal{N}(0,1/s)\) and \[\sum_{l=1}^s(W_l-\bar W)^2\sim \chi^2_{s-1}.\] The Gaussian tail bound gives \[\text{Pr}\left(|\bar W| \leq \sqrt{\frac{2\log(4/\eta)}{s}}\right) \geq 1- \frac{\eta}{2}. \label{eq:124barWs124}\tag{11}\] Laurent-Massart inequality gives \[\text{Pr}\left(\sum_{l=1}^s(W_l-\bar W)^2 \geq (s-1)-2\sqrt{(s-1)\log(2/\eta)} \right) \geq 1- \frac{\eta}{2}. \label{eq:SWW}\tag{12}\] Since \(s\ge 1+16\log(2/\eta)\), we have \(2\sqrt{(s-1)\log(2/\eta)}\leq {(s-1)}/{2}\). Thus, on an event of probability at least \(1-\eta\), \[\begin{align} \frac{\bar{W}^2}{\sum_{l=1}^s(W_l-\bar{W}_s)^2} &\leq \frac{2\log(4/\eta)/s}{(s-1)/2} \\ &\leq \frac{2\log(4/\eta)/s}{s/3} = \frac{6\log(4/\eta)}{s^2}\\ &\leq \frac{1}{s}. \end{align}\] where the first inequality is by plugging into high probability events in 11 and 12 , the second inequality is because \((s-1)/2\ge s/3\) based on \(s\ge 3\), and the third inequality is based on definition \(s\ge 8\log(4/\eta)\). Therefore, we have \[\text{Pr}(h_s\le 2/s) \geq 1-\eta.\]
8.3 Proof of Lemma 3↩︎
Fix an arm and omit the arm index. By Proposition 2, \({(t_{\rm cal}-2)\hat{v}_0}/{v}\sim \chi^2_{t_{\rm cal}-2}\). Then, Laurent-Massart inequalities also states \[\text{Pr}\left(\frac{(t_{\rm cal}-2)\hat{v}_0}{v}\ge (t_{\rm cal}-2)+2\sqrt{(t_{\rm cal}-2)\log(8K/\delta)}+2\log(8K/\delta)\right)\le e^{-\log(8K/\delta)}=\frac{\delta}{8K},\] which provides the high probability upper bound on \(U_0/v\). Namely, \[\frac{\hat{v}_0}{v}\le 1+2\sqrt{\frac{\log(8K/\delta)}{t_{\rm cal}-2}}+2\frac{\log(8K/\delta)}{t_{\rm cal}-2}\] with probability at least \(1-\delta/(8K)\). Let \(y=\sqrt{{\log(8K/\delta)}/{(t_{\rm cal}-2)}}\), recall \(U_0=\hat{v}_0/(1-2y)\), we have \[\begin{align} \frac{U_0}{v}-1&\le \frac{1+2y+2y^2}{1-2y} - 1 = \frac{4y+2y^2}{1-2y} \le 8y\le \kappa \end{align}\] where the second inequality is because the definition of \(t_{\rm cal}\) implies \(t_{\rm cal}-2\ge {64\log(8K/\delta)}/{\kappa^2}\) and further \(y\leq\kappa/8\leq1/8\). Therefore we have \[\text{Pr}\left(U_0\le (1+\kappa)v \right) \geq 1-\frac{\delta}{8K}. \label{eq:U6040143k41v}\tag{13}\]
At the same time, Laurent-Massart inequalities state that \[\text{Pr}\left(\frac{(t_{\rm cal}-2)\hat{v}_0}{v}\le (t_{\rm cal}-2)-2\sqrt{(t_{\rm cal}-2)\log(8K/\delta)}\right)\le e^{-\log(8K/\delta)}=\frac{\delta}{8K},\] which provides the high probability lower bound on \(\hat{v}_0/v\). Namely, \[\frac{\hat{v}_0}{v}\ge 1-2\sqrt{\frac{\log(8K/\delta)}{t_{\rm cal}-2}}\] with probability at least \(1-\delta/(8K)\). Plug into the definition of \(U_0\), we have \(U_0\geq v\) with probability at least \(1-\delta/(8K)\). \[\text{Pr}\left(U_0\geq v \right) \geq 1-\frac{\delta}{8K}. \label{eq:v60U}\tag{14}\]
Combining the 13 and 14 with a union bound, we finish the proof.
9 Omitted Proofs↩︎
9.1 Proof of Proposition 1↩︎
Since \(Y_{i,l}-m_i\) has mean zero, for any \(\beta\in\mathbb{R}\), \[\mathbb{E}\left[X_{i,l}-\beta(Y_{i,l}-m_i)\right] = \mathbb{E}[X_{i,l}]-\beta\mathbb{E}[Y_{i,l}-m_i] = p_i.\] Thus the adjusted observation is unbiased for \(p_i\).
For the variance, using \(\text{Var}(X_{i,l})=1\), \(\text{Var}(Y_{i,l}-m_i)=\text{Var}(Y_{i,l})=1\), and \(\text{Cov}(X_{i,l},Y_{i,l}-m_i)=\text{Cov}(X_{i,l},Y_{i,l})=\rho_i\), we have \[\begin{align} \text{Var}\left(X_{i,l}-\beta(Y_{i,l}-m_i)\right) &= \text{Var}(X_{i,l}) +\beta^2\text{Var}(Y_{i,l}-m_i) -2\beta\text{Cov}(X_{i,l},Y_{i,l}-m_i)\\ &= 1+\beta^2-2\rho_i\beta. \end{align}\] Therefore the variance is minimized at \(\beta=\rho_i\), and the minimum variance is \(1-\rho_i^2\).
9.2 Proof of Proposition 2↩︎
Fix an arm \(i\) and omit the arm index. Conditional on \(W_1,\ldots,W_s\), the regression model 5 is a fixed-design Gaussian linear model \(X_l=p+\rho W_l+\epsilon_l\), where \(\epsilon_l\stackrel{i.i.d.}{\sim}\mathcal{N}(0,v)\). Let \[Z= \begin{pmatrix} 1 & W_1\\ \vdots & \vdots\\ 1 & W_s \end{pmatrix}, \qquad \theta= \begin{pmatrix} p\\ \rho \end{pmatrix}, \qquad X=(X_1,\ldots,X_s)^\top .\] Then, 5 can be rewritten as \[X=Z\theta+\epsilon, \qquad \epsilon\mid W_1,\ldots,W_s\sim \mathcal{N}(0,vI_s).\] The OLS estimator is \(\hat{\theta}=(Z^\top Z)^{-1}Z^\top X\). Therefore, conditional on \(W_1,\ldots,W_s\), \[\hat{\theta}-\theta = (Z^\top Z)^{-1}Z^\top \epsilon \sim \mathcal{N}\left(0,v(Z^\top Z)^{-1}\right).\] The intercept estimator \(\hat{p}\) is the first coordinate of \(\hat{\theta}\). Hence \[\hat{p}-p\mid W_1,\ldots,W_s \sim \mathcal{N}\left(0,v\left[(Z^\top Z)^{-1}\right]_{11}\right).\] where \[\begin{align} \left[(Z^\top Z)^{-1}\right]_{11} &= \begin{bmatrix}\begin{pmatrix} s & \sum_{l=1}^s W_l\\ \sum_{l=1}^s W_l & \sum_{l=1}^s W_l^2 \end{pmatrix}^{-1} \end{bmatrix}_{11}\\ &= \frac{\sum_{l=1}^s W_l^2}{s\sum_{l=1}^s W_l^2-\left(\sum_{l=1}^s W_l\right)^2} \\ &= \frac{1}{s} + \frac{\bar W^2}{\sum_{l=1}^s W_l^2-s\bar W^2} = h \end{align}\] Thus \[\hat{p}-p\mid W_1,\ldots,W_s\sim \mathcal{N}(0,vh).\]
It remains to prove the residual variance statement. Let \(P_Z=Z(Z^\top Z)^{-1}Z^\top\) be the projection matrix onto the column space of \(Z\). The residual vector is \[X-Z\hat{\theta} = (I-P_Z)X = (I-P_Z)\epsilon,\] because \((I-P_Z)Z\theta=0\). Conditional on \(W_1,\ldots,W_s\), \[\frac{1}{v}\|X-Z\hat{\theta}\|_2^2 = \frac{1}{v}\epsilon^\top(I-P_Z)\epsilon.\] The matrix \(I-P_Z\) is symmetric and idempotent. Since the design has two linearly independent columns whenever \(S_{WW}>0\), we have \[\operatorname{rank}(I-P_Z)=s-2.\] For a Gaussian vector \(\epsilon\sim \mathcal{N}(0,vI_s)\), the quadratic form with a symmetric idempotent matrix of rank \(s-2\) satisfies \[\frac{1}{v}\epsilon^\top(I-P_Z)\epsilon\sim\chi^2_{s-2}.\] By the definition 6 , \(\hat{v}=\|X-Z\hat{\theta}\|_2^2/(s-2)\), we obtain \[\frac{(s-2)\hat{v}}{v}\sim\chi^2_{s-2}.\] This completes the proof.
10 Repeated Proxy Calls↩︎
Consider a limited number of proxy query for \(d\) times (where \(d\) might be determined by the latency or concurrency in real world sequential decisions), the real reward \(X_{i,l}\) is paired with a proxy vector \((Y_{i,l}^{(1)},\dots,Y_{i,l}^{(d)})\), which follow a multivariate gaussian distribution similar to 2 as \[\begin{pmatrix} X_{i,l}\\ Y_{i,l}^{(1)}\\ \vdots\\ Y_{i,l}^{(d)} \end{pmatrix} \sim \mathcal{N}\!\left( \begin{pmatrix} p_i\\ m_i\mathbf{1}_d \end{pmatrix}, \begin{pmatrix} 1 & \rho_i\mathbf{1}_d^\top\\ \rho_i\mathbf{1}_d & \Sigma_d \end{pmatrix} \right),\qquad \rho_i\in(-1,1), \label{eq:proxy-model-vector}\tag{15}\] where \(\Sigma_d\) is a positive semi-definite matrix with unit diagonal. Similar to the single proxy case, we can introduce \(W_{i,l}^{(j)}=Y_{i,l}^{(j)}-m_i\) for all \(j\in[d]\) then \(\mathbb{E}[W_{i,l}^{(j)}]=0\) and \(\text{Var}(W_{i,l}^{(j)})=1\). Assuming \(Y_{i,l}^{(1)},\dots,Y_{i,l}^{(d)}\) are conditional i.i.d given \(X_{i,l}\), the repeated proxies \[\bar W_{i,l}^{(d)}=\frac{1}{d}\sum_{j=1}^d W_{i,l}^{(j)},\] can be reduced to a single averaged proxy, i.e., using all \(d\) proxies as linear control variates is equivalent to using their average as a single proxy.
Proposition 3. Under the multivariate Gaussian model in 15 , and assume \(Y_{i,l}^{(1)},\ldots,Y_{i,l}^{(d)}\) are conditionally i.i.d. given \(X_{i,l}\) for all \(j\in[d]\). Then, for every \(j\ne k\), \(\text{Cov}\left(Y_{i,l}^{(j)},Y_{i,l}^{(k)}\right)=\rho_i^2\), or equivalently \[\Sigma_d=(1-\rho_i^2)I_d+\rho_i^2\mathbf{1}_d\mathbf{1}_d^\top.\] Moreover, using the single averaged proxy \(\bar W_{i,l}^{(d)}\) instead of the full vector. The effective correlation and residual variance are \[\rho_i^{(d)} = \frac{\sqrt d\,\rho_i}{\sqrt{1+(d-1)\rho_i^2}}, \qquad v_i^{(d)}=\frac{1-\rho_i^2}{1+(d-1)\rho_i^2}.\]
We omit the arm \(i\) and pull indices \(l\), and write \(X=X_{i,l}\), \(Y^{(j)}=Y_{i,l}^{(j)}\), and \(W^{(j)}=Y^{(j)}-m_i\). Then \(\mathbb{E}[W^{(j)}]=0\), \(\text{Var}(X)=\text{Var}(W^{(j)})=1\), and \(\text{Cov}(X,W^{(j)})=\text{Cov}(X,Y^{(j)})=\rho_i\) .
We first identify the covariance matrix of the repeated proxies. Extracting the principal sub-matrix indexed by \((X,W^{(j)},W^{(k)})\) from the covariance matrix in 15 , by Schur complement, we have \[\begin{align} 0 = \text{Cov}(W^{(j)},W^{(k)}\mid X) &= \text{Cov}(W^{(j)},W^{(k)}) - \text{Cov}(W^{(j)},X)\text{Var}(X)^{-1}\text{Cov}(X,W^{(k)})\\ &= \text{Cov}(W^{(j)},W^{(k)}) - \rho_i \cdot 1 \cdot \rho_i, \end{align}\] we obtain \(\text{Cov}(W^{(j)},W^{(k)})=\rho_i^2\) for all \(j\neq k\) and \(j,k\in[d]\). Equivalently, for \(W=:(W^{(1)},\ldots,W^{(d)})^\top\), \[\text{Cov}(W)=(1-\rho_i^2)I_d+\rho_i^2\mathbf{1}_d\mathbf{1}_d^\top.\]
We now show that the full vector proxy is equivalent to the averaged proxy. Consider the best linear control variate based on the full vector \(W\), \[\alpha^* \in \arg\min_{\alpha\in\mathbb{R}^d} \text{Var}\left(X-\alpha^\top W\right).\] Since \(\mathbb{E}[W]=0\), the solution is \[\alpha^*=\text{Cov}(W)^{-1}\text{Cov}(W,X).\] Recall that \[\text{Cov}(W)=(1-\rho_i^2)I_d+\rho_i^2\mathbf{1}_d\mathbf{1}_d^\top, \qquad \text{Cov}(W,X)=\rho_i\mathbf{1}_d.\] By the Sherman-Morrison formula, \[\text{Cov}(W)^{-1} = \frac{1}{1-\rho_i^2}I_d - \frac{\rho_i^2}{(1-\rho_i^2)(1+(d-1)\rho_i^2)} \mathbf{1}_d\mathbf{1}_d^\top.\] Therefore, \[\alpha^* = \left( \frac{1}{1-\rho_i^2}I_d - \frac{\rho_i^2}{(1-\rho_i^2)(1+(d-1)\rho_i^2)} \mathbf{1}_d\mathbf{1}_d^\top \right)\rho_i\mathbf{1}_d = \frac{\rho_i}{1+(d-1)\rho_i^2}\mathbf{1}_d.\] Thus the optimal vector control term satisfies \[(\alpha^*)^\top W = \frac{\rho_i}{1+(d-1)\rho_i^2} \sum_{j=1}^d W^{(j)} = \frac{d\rho_i}{1+(d-1)\rho_i^2}\bar W^{(d)}.\] Therefore the optimal linear control variate based on the full proxy vector only depends on the averaged proxy \(\bar W^{(d)}\), which proves the claimed equivalence.
Proposition 3 shows that the averaged proxy \(\bar W_{i,l}^{(d)}\) is a sufficient summary of the whole proxy vector for the control-variate adjustment: the best linear adjustment built from \((Y_{i,l}^{(1)},\dots,Y_{i,l}^{(d)})\) coincides with the one built from their average alone. Consequently, the \(d\)-proxy observation reduces exactly to a single-proxy observation of the form 2 , with the correlation and residual variance replaced by the effective quantities \(\rho_i^{(d)}\) and \(v_i^{(d)}\). Therefore, we are able to develop the algorithm and its analysis for a single proxy per pull without loss of generality, and recover the guarantee for any fixed \(d\) by substituting \(\rho_i^{(d)}\) and \(v_i^{(d)}\) for \(\rho_i\) and \(v_i\).