Multi-Armed Bandits With Machine Learning-Generated Surrogate Rewards

1. Problem Statement

The paper addresses the Multi-Armed Bandit (MAB) problem, a framework for sequential decision-making under uncertainty where an agent must balance exploration (trying new actions) and exploitation (choosing the best known action) to minimize cumulative regret.

The Core Challenge:
Traditional MAB algorithms rely solely on online data collected during the decision-making process. However, in many real-world scenarios (e.g., clinical trials, recommendation systems), rich offline/historical data (covariates, user features) is available before the online phase begins.

• The Opportunity: Machine Learning (ML) models can be trained on this offline data to predict outcomes, generating surrogate rewards.
• The Difficulty: These ML-generated surrogates are often biased. They may not preserve the ranking of true rewards (e.g., an ML model might predict a new drug is effective when it is not, or vice versa) because they often rely on extrapolation to unseen environments.
• The Gap: Existing methods either assume offline data is unbiased (rarely true) or require knowledge of the bias bound. There is a lack of algorithms that can leverage correlated but biased surrogate rewards to reduce regret without prior knowledge of the bias magnitude or the covariance structure.

2. Methodology: MLA-UCB

The authors propose the Machine Learning-Assisted Upper Confidence Bound (MLA-UCB) algorithm. The core innovation is integrating Prediction-Powered Inference (PPI) techniques into the UCB framework to debias the mean reward estimation.

A. The Estimator

Instead of using the raw sample mean of true rewards, MLA-UCB constructs a machine learning-assisted mean estimator ($\hat{\mu}^{MLA}$) for each arm $k$:
$$\hat{\mu}^{MLA}_{k,n} = \bar{R}_{k,n} - \hat{\lambda}_{k,n} (\bar{\hat{R}}_{k,n} - \bar{\hat{R}}^{off}_{k})$$
Where:

• $\bar{R}_{k,n}$: Sample mean of true rewards (online).
• $\bar{\hat{R}}_{k,n}$: Sample mean of surrogate rewards (online).
• $\bar{\hat{R}}^{off}_{k}$: Sample mean of surrogate rewards (offline).
• $\hat{\lambda}_{k,n}$: A data-driven weight (plug-in estimate) derived from the sample covariance between true and surrogate rewards.

This estimator effectively treats the offline surrogate mean as a "control variate." It subtracts the bias observed in the online surrogate data relative to the offline surrogate data, correcting the true reward estimate.

B. Confidence Bounds

To ensure theoretical guarantees, the authors derive a rigorous Upper Confidence Bound (UCB) for $\hat{\mu}^{MLA}_{k,n}$.

• Distributional Analysis: Under a bivariate Gaussian assumption (true and surrogate rewards), they prove that the estimation error can be decomposed into two independent components: one related to the uncertainty of the offline surrogate mean and another related to the regression residual.
• Variance Reduction: The confidence bound scales with the variance of the residuals, which is reduced by a factor of $(1 - \rho_k^2)$, where $\rho_k$ is the correlation between true and surrogate rewards.
• No Prior Knowledge: Crucially, the algorithm does not require prior knowledge of the variance of true rewards, the variance of surrogates, or the covariance matrix. It estimates these online using sample statistics.

C. Extension to Non-Gaussian and Batched Settings

The paper extends the method to:

1. Batched Rewards: Where decisions are applied to a batch of $m$ users simultaneously.
2. Non-Gaussian Distributions: Using self-normalized moderate deviation bounds (Cramér-type inequalities) to construct valid confidence intervals even when rewards are not Gaussian, provided the batch size is sufficiently large to induce approximate normality.

3. Key Contributions

1. New MAB Setting: Formalizes the setting of MAB with arbitrarily biased but correlated ML-generated surrogate rewards. This relaxes the unrealistic assumption that offline data must be unbiased or that the bias is bounded by a known constant.
2. MLA-UCB Algorithm: Proposes an algorithm that:
   • Works with any reward prediction model (linear, neural networks, etc.).
   • Requires no prior knowledge of the bias or covariance structure.
   • Achieves asymptotic optimality.
3. Theoretical Guarantees:
   • Regret Upper Bound: Proves that the cumulative regret is strictly lower than classical UCB when $\rho_k \neq 0$. The leading term of the regret is reduced by a factor proportional to $(1 - \rho_k^2)$.
   • Regret Lower Bound: Derives a matching lower bound, proving that MLA-UCB is asymptotically optimal among uniformly good policies.
   • Robustness: Shows that even if the surrogate mean is completely misaligned with the true mean (high bias), the algorithm still improves performance as long as there is non-zero correlation.
4. Batched & Non-Gaussian Theory: Extends the theory to batched settings with non-Gaussian rewards, providing computable confidence bounds using self-normalized statistics.

4. Results

Theoretical Results

• Gaussian Case: If the offline sample size $N_k$ is poly-logarithmic in the horizon $T$ (i.e., $N_k = \Omega((\log T)^{4/3})$), the expected regret is bounded by:
  $$E[Reg_T] \leq \sum_{k \neq k^*} \frac{2\Delta_k \log T}{\log\left(1 + \frac{\Delta_k^2}{\sigma_k^2(1-\rho_k^2)}\right)}$$
  This demonstrates that the effective variance is reduced from $\sigma_k^2$ to $\sigma_k^2(1-\rho_k^2)$.
• Optimality: The algorithm achieves the information-theoretic lower bound for this setting.

Empirical Results

The authors validated the method through extensive simulations and real-world case studies:

1. Synthetic Gaussian Data: MLA-UCB significantly outperformed classical UCB. Regret reduction was linear in $\rho_k^2$ and saturated as offline sample size increased.
2. ML-Generated Non-Gaussian Surrogates: Using various models (Linear Regression, SVR, Neural Networks, Decision Trees) to generate surrogates for a non-linear reward function, MLA-UCB consistently reduced regret compared to UCB, even when the Gaussian assumption was violated.
3. Real-World Case Studies:
   • Language Model Selection: Used open-source LLMs (Llama) as surrogates to select the best proprietary model (GPT-4o, Claude). MLA-UCB reduced regret compared to standard selection methods.
   • Video Recommendation: Used a gradient boosting model trained on historical user-item interactions to generate surrogates for new video recommendations. MLA-UCB improved user engagement (watch ratio) and reduced regret.

5. Significance

This paper bridges the gap between offline machine learning and online sequential decision-making.

• Practical Impact: It provides a rigorous framework for practitioners to utilize vast amounts of historical data and cheap-to-generate ML predictions to accelerate learning in online environments, even when those predictions are imperfect or biased.
• Theoretical Advancement: It resolves the challenge of "biased control variates" in bandits, showing that correlation is sufficient for variance reduction without needing to know the bias.
• Generality: The method is model-agnostic and applicable to a wide range of domains, from clinical trials (predicting long-term outcomes) to digital advertising and recommendation systems.

In summary, MLA-UCB transforms the "noise" of biased ML predictions into a signal that accelerates the discovery of optimal actions, offering a provably optimal solution for the modern era of data-rich, AI-assisted decision-making.