A UCB Bandit Algorithm for General ML-Based Estimators

TL;DR

ML-UCB introduces a generalized UCB framework that integrates arbitrary machine learning estimators by calibrating exploration to the estimator's learning curve. By modeling the Mean Squared Error as $\text{MSE}(n) = C n^{-s}$ and using a learning-curve–driven $\psi$-UCB bound, the approach achieves sublinear regret under a generalized $\psi$-decay with rate $s$ and provides a model-agnostic way to combine ML with bandits. Theoretical guarantees are complemented by experiments on online matrix factorization for collaborative filtering, showing superior regret performance to LinUCB across settings and demonstrating robust exploration with varying $s$. This framework enables principled, adaptive integration of diverse ML models into sequential decision problems, with practical impact in recommender systems and beyond.

Abstract

We present ML-UCB, a generalized upper confidence bound algorithm that integrates arbitrary machine learning models into multi-armed bandit frameworks. A fundamental challenge in deploying sophisticated ML models for sequential decision-making is the lack of tractable concentration inequalities required for principled exploration. We overcome this limitation by directly modeling the learning curve behavior of the underlying estimator. Specifically, assuming the Mean Squared Error decreases as a power law in the number of training samples, we derive a generalized concentration inequality and prove that ML-UCB achieves sublinear regret. This framework enables the principled integration of any ML model whose learning curve can be empirically characterized, eliminating the need for model-specific theoretical analysis. We validate our approach through experiments on a collaborative filtering recommendation system using online matrix factorization with synthetic data designed to simulate a simplified two-tower model, demonstrating substantial improvements over LinUCB

Figures (4)

• Figure 1: Flowchart of the ML-UCB Algorithm
• Figure 2: Learning curve analysis: $\log(\text{MSE})$ vs $\log(n)$, where $n$ is the number of training samples. (a) The full training trajectory shows an initial plateau (cold start) followed by power-law decay. The overall slope is $-0.27$, but this is dominated by the cold-start phase. (b) Focusing on the stable regime (last 20% of training), we observe a cleaner linear relationship with slope $s \approx 0.97$, indicating $\text{MSE} = O(n^{-0.97})$ convergence.
• Figure 3: Comprehensive comparison of ML-UCB ($s=1.0$, $\alpha=10.0$) vs LinUCB ($\alpha=1.0$ and $\alpha=1.4$) over 33,333 iterations. Top row: (left) Cumulative regret showing ML-UCB's superior performance, (center-left) regret rate $R(t)/t$ decreasing over time, (center-right) final regret comparison, (right) regret difference. Bottom row: (left) optimal selection accuracy, (center-left) training MSE learning curves on log scale, (center-right) smoothed instantaneous regret, (right) summary statistics.
• Figure 4: Effect of convergence rate $s$ on ML-UCB performance. We compare three settings: $s=0.272$ (full trajectory slope), $s=0.5$ (conservative), and $s=0.97$ (stable regime). Smaller $s$ values produce larger exploration bonuses, while larger $s$ values enable faster transition to exploitation. All experiments use identical ground truth matrices for fair comparison.

Theorems & Definitions (6)

• Example 4.2: CGF of Gaussian
• Theorem 4.3: CGF of Sample Mean
• Theorem 4.4: Concentration Inequality
• Theorem 4.5: $\psi$-UCB Regret
• Theorem 6.1
• Proof 6.2