Best-Arm Identification with Generative Proxy

In many real-world decision-making scenarios, such as pricing experiments or recommendation systems, obtaining a reward observation is expensive. However, researchers often have access to cheap, auxiliary 'proxy' scores (e.g., from LLMs or predictive models) that are correlated with the true reward. This paper asks: how can we leverage these unreliable, cheap proxies to identify the best arm with fewer costly reward samples, while ensuring the identification remains statistically valid?

The authors propose PROBE (PRoxy OLS for Best-arm Exploration), a phase-elimination algorithm. The core idea is to treat the proxy as a control variate. By centering the proxy (using an offline-estimated mean) and performing an ordinary least squares (OLS) regression, the algorithm effectively subtracts the predictable part of the reward noise. A key technical challenge is that the reward-proxy correlation is unknown and must be learned online. To avoid the pitfalls of 'plug-in' estimates—which can be anti-conservative and lead to incorrect arm identification—PROBE maintains a one-sided upper certificate on the residual variance using the chi-square distribution properties of the OLS fit. This ensures that the algorithm remains delta-PAC (probably approximately correct) regardless of the true correlation.

PROBE successfully reduces the number of costly online samples required for best-arm identification. The theoretical analysis shows that the algorithm achieves the oracle sample complexity—the performance one would expect if the correlation were known in advance—up to a constant multiplicative factor and a small, fixed calibration cost. Numerical experiments, including a real-world auto-loan pricing replay, demonstrate that PROBE scales its sample savings according to the strength of the reward-proxy correlation, achieving up to a 61% reduction in stopping time compared to a reward-only baseline.