Best-Arm Identification with Generative Proxy

2026-07-07

AI Summary

This work addresses fixed-confidence best-arm identification in settings where reward observations are costly, proposing the PROBE algorithm that leverages cheap, reward-correlated proxy scores to aid decision-making. PROBE learns the unknown reward-proxy relationship online and employs a control variate technique to transform the problem into a heteroscedastic identification task. It directly maintains upper confidence bounds on residual variance via ordinary least squares estimation, guaranteeing δ-PAC correctness. Theoretical analysis shows that its sample complexity incurs only a constant-factor overhead plus calibration cost compared to the optimal strategy with known correlation. Experiments on synthetic data and a real-world loan pricing task demonstrate substantial sample savings, with the degree of reduction positively correlated with the strength of the proxy-reward relationship.

Abstract

Best-arm identification is a canonical model for data-driven decision-making, but in many applications each reward observation is costly. Motivated by the growing availability of cheap predictions from machine learning and large language models, we study fixed-confidence best-arm identification in which each costly reward pull is paired with a cheap but correlated proxy score. The marginal mean of the proxy can be estimated offline and is treated as known, whereas its correlation $ρ$ with the reward, which governs how much the proxy helps, is unknown and must be learned online in pair with real rewards. We show that a control-variate adjustment turns this model into a heteroscedastic identification problem whose oracle sample complexity improves by residual variance $1-ρ^2$. The central difficulty is that the correlation must be learned from the same costly samples that identification consumes online, and that a plug-in estimate of the residual variance is anti-conservative and can compromise correctness. We propose PROBE (PRoxy OLS for Best-arm Exploration), a phase-elimination algorithm that directly maintains an upper certificate on the residual variance with an ordinary least squares fit, whose exact chi-square law keeps the certificate valid regardless of the unknown correlation. We prove that PROBE is $δ$-PAC and attains the known-correlation oracle sample complexity up to a constant multiplicative factor and a constant additive calibration cost. The guarantee extends to the $(ε,δ)$-PAC setting under minimal changes to the algorithm. Numerical experiments on synthetic instances and on an auto-loan pricing replay with large language model and tabular proxies confirm that the sample savings of PROBE scale with the strength of the reward-proxy correlation, exactly as the theory predicts.

Problem

• best-arm identification
• proxy score
• correlation learning
• fixed-confidence
• sample complexity

Innovation

• best-arm identification
• proxy score
• control variates
• heteroscedasticity
• PAC learning