Regularized OFU: an Efficient UCB Estimator for Non-linear Contextual Bandit

Contextual bandit langford2007epoch ; bubeck2012regret is a basic sequential decision-making problem which is extensively studied and widely applied in machine learning. At each time step in contextual bandit, an agent is presented with a context and the agent chooses an action according to the context. After that the agent will receive a reward conditioned on the context and the selected action. The goal of the agent is to maximize its cumulative reward. One key challenge in contextual bandit is exploration and exploitation (EE) trade-off. In order to maximize cumulative reward, the agent exploits collected data to take the action with high estimated reward while it also explores the undiscovered areas to learn knowledge.

Optimism in Face of Uncertainty (OFU) is a powerful principle for EE trade-off. When facing uncertainty, OFU first optimistically guesses how good each action could be and then takes the action with highest guess. In philosophy, it exploits collected knowledge by estimating the expected reward and enforces exploration of the areas with large uncertainty via optimism. OFU algorithms can be divided into two categories, i.e., confidence bonus method auer2002finite and confidence set method jaksch2010near . In confidence bonus method, the uncertainty of each action is measured according to its historical frequency, which is added to its mean value as the estimated reward; in confidence set method, the reward function is constrained to a subset with high confidence and the maximum value over the confidence set is taken as the estimated value. These OFU algorithms have been proved to achieve the (near-)optimal regret bounds on tabular and linear problems, including multi-armed bandit (MAB) auer2002finite , linear contextual bandit (LCB) chu2011contextual and kernel contextual bandit (KCB) valko2013finite , and have also been successfully applied in practice.

However, the representation power of the tabular and linear models often becomes insufficient when facing many real-world contextual bandit problems where the scale of the underlying problem is huge or even infinite, and the reward mechanism is complex. Therefore, in practice, non-linear deep neural networks (DNN) are used to approximate the reward functions. While DNN models work well to fit training data, it remains challenging to quantify its uncertainty on new data. Therefore, existing OFU methods cannot be efficiently applied to non-linear contextual bandit. For confidence set method, it is difficult to compute the confidence set of the non-linear deep neural networks and hence hard to obtain the optimal action. For the confidence bound method, it is unclear how to design the frequency for non-linear reward functions. Overall, estimating uncertainty of the reward function approximated by non-linear functions, e.g., DNN, becomes a major obstacle when applying OFU algorithms to real-world applications.

There are several works zhou2020neural ; filippi2010parametric ; zahavy2019deep ; gu2021batched that attempt to generalize OFU to non-linear contextual bandit. One is NeuralUCB zhou2020neural that takes the gradient of each observation as the random feature and constructs the upper confidence bound accordingly. It has been shown that NeuralUCB achieves sublinear regret on neural tangent kernel contextual bandit (NTKCB). However, at each time step, NeuralUCB requires to compute an inversion of a $p\times p$ matrix where $p$ is the number of model parameters. As the time complexity of matrix inversion is $O(p^{2.373})$ davie2013improved , the computational cost is prohibited for large-scale DNN models. Therefore, it is still largely under-explored to develop OFU algorithms that are able to achieve good EE trade-off with reasonable computational cost for non-linear contextual bandit.

Motivated by these challenges, in this paper, we propose a novel and general OFU method, called Regularized OFU (ROFU). In ROFU, we measure the uncertainty of the reward function by a differentiable regularizer, and then we calculate the optimistic estimation by maximizing the reward function with the regularizer penalizing its uncertainty. We prove that ROFU enjoys near-optimal regret on many contextual bandit problems with relatively simple reward mechanism, including MAB, LCB, KCB and NTKCB. Importantly, in contrast to existing OFU algorithms, ROFU admits an efficient implementation with gradient-based optimizer as the optimization problem in ROFU is unconstrained and differentiable, which has a time complexity linear to $p$. Moreover, our algorithm easily extends to general deep neural network models.

We summarize our contributions as follows:

• •

  We show that ROFU achieves (near-)optimal regret on contextual bandits with simple reward functions, including MAB, LCB, KCB and NTK bandit. These results theoretically justify the efficiency of ROFU on EE trade-off.

• •

  We propose an algorithm which can efficiently approximate the upper confidence bound with standard gradient-based optimizer.

• •

  We empirically evaluate ROFU on complex contextual bandits with reward functions beyond linear, kernel and NTK. We show that ROFU also provides efficient UCB estimation for popular DNN architectures including MLP, CNN and ResNet. Moreover, our algorithm enjoys a smaller regret than strong baselines on real-world non-linear contextual bandit problems introduced by deep_bayesian_bandit_showdown .

In this work, we consider contextual bandit with general reward functions.

The most studied contextual bandit is linear contextual bandit, where the reward is generated from a linear function, i.e., $f_{\theta^{*}}(x,a)\mathrel{\mathop{\mathchar 58\relax}}=\phi(x,a)^{\top}\theta^{*}$ where $\phi(x,a)$ is a known feature map. For more complex problems, we have to use non-linear models, such as DNN, to represent the reward function.

In practice, given dataset $D\mathrel{\mathop{\mathchar 58\relax}}=\{(x_{t},a_{t},r_{t})\}_{t\leq|D|}$, we can estimate $\theta^{*}$ by minimizing a loss function, i.e., $\bar{\theta}\mathrel{\mathop{\mathchar 58\relax}}=\operatorname*{arg\,min}_{\theta}\mathcal{L}(\theta;D)$. In this work, we focus on the mean squared loss and its variants,

This is because MSE is a natural choice for regression tasks. Generally speaking, OFU algorithms trade off exploration and exploitation with an optimistic estimation on the reward for each action, and then select $a_{t}=\operatorname*{arg\,max}_{a}\text{OFU}(x_{t},a)$. In the literature, OFU algorithms can be divided into two categories. The first one optimistically estimates the reward by adding a confidence bonus to $f_{\bar{\theta}}$, i.e.,

We call this OFU method Bonus-OFU. In simple cases, the bonus can be constructed using concentration bounds. For example, at time step $\tau$, in UCB1 for MAB auer2002finite , the bonus is $\sqrt{2\log\tau/n(a)}$, where $n(a)$ is the number of pulls on action $a$; in LinUCB for LCB chu2011contextual , the bonus is $\sqrt{\phi(x,a)^{\top}(\sum_{t<\tau}\phi(x_{t},a_{t})\phi(x_{t},a_{t})^{\top}+{\bm{I}})^{-1}\phi(x,a)}$. Both UCB1 and LinUCB are proved to have (near-)optimal regret bounds.

The second kind of OFU algorithms optimistically estimates the reward by the best possible value over a confidence set of the reward functions, i.e.,

where $\Theta_{\delta}\mathrel{\mathop{\mathchar 58\relax}}=\{\theta\mathrel{\mathop{\mathchar 58\relax}}{\mathbb{P}}(D|\theta)>\delta\}$ and ${\mathbb{P}}(D|\theta)$ is the likelihood of $D$ given $\theta$. We call this OFU method Set-OFU.

For convenience, we suppose $\theta$ is a $p$-dimensional vector. Both Bonus-OFU and Set-OFU can effectively trade-off EE, however, they require a lot of computational resources when $f_{\theta}$ is a complex neural network: firstly, existing Bonus-OFU algorithms, such as NeuralUCB zhou2020neural , need to calculate the inversion of a $p\times p$ matrix, which is expensive especially for large DNN models; secondly, it is also hard to efficiently compute the confidence set of DNN parameters and even harder to maximize the value over it to compute the $\text{OFU}^{S}$. The computational cost significantly limits the application of OFU methods for complex tasks.

In the next section, we will propose a new OFU method, named Regularized OFU, which achieves a near-optimal regret bound on simple tasks, and more importantly, admits a very efficient implementation with a gradient-based optimizer.

Different from Bonus-OFU and Set-OFU, we measure the uncertainty of reward with a (differentiable) function, and compute the optimistic reward estimation under the regularization of the uncertainty measure function. We call the reward estimation regularized OFU (ROFU) and take action accordingly. In round $t$, given context and action pair $(x_{t},a)$, ROFU is defined as follows and the algorithm is presented in Alg. 1:

where $D_{t-1}$ is the set of collected data in first $t-1$ rounds, $\mathcal{R}$ is a regularizer, $\eta$ is the weight of the regularizer, $g$ is a monotonically increasing function with $g(0)=0$ and $\theta_{t-1}$ is trained by minimizing some objective function, i.e., $\theta_{t-1}\mathrel{\mathop{\mathchar 58\relax}}=\operatorname*{arg\,min}_{\theta}\mathcal{L}(\theta;D_{t-1})$. We simply take action $a_{\text{ROFU}}\mathrel{\mathop{\mathchar 58\relax}}=\operatorname*{arg\,max}_{a\in A}\text{OFU}^{R}(x_{t},a)$ to interact with the environment.

In Eq. (4), we first solve an unconstrained optimization problem to calculate $\hat{\theta}(x_{t},a)$. Then, in Eq. (5), we add a bonus to $f_{\theta_{t-1}}(x_{t},a)$ according to the difference between $f_{\hat{\theta}(x_{t},a)}(x_{t},a)$ and $f_{\theta_{t-1}}(x_{t},a)$. From Eq. (4) and (5), we can see that , while ROFU still adds a bonus on $f_{\theta_{t-1}}(x,a)$ as Bonus-OFU, it does not need to explicitly construct the upper confidence bound via concentration inequalities. Instead, we construct the upper confidence bound by solving a regularized optimization problem which is more general. In comparison to Set-OFU, our optimization problem is unconstrained, so Eq. (4) can be efficiently computed by standard gradient-based optimizers. In Sec. 3.2, we will propose an algorithm which can efficiently approximate $\text{OFU}^{R}$ with a few steps of gradient descent.

An important question is whether selecting action with maximum $\text{OFU}^{R}$ achieves good EE trade-off, or whether $\text{OFU}^{R}(x_{t},a)$ is a reasonable upper confidence bound over $f_{\theta_{t-1}}(x_{t},a)$. To provide more insights, consider the case where $\mathcal{L}=\mathcal{R}$. It is easy to see that $\forall(x,a),f_{\hat{\theta}(x,a)}(x,a)\geq f_{\theta_{t-1}}(x,a)$. Thus, given $g$ is monotonically increasing and $g(0)=0$, we have $\text{OFU}^{R}(x,a)$ is an upper bound of $f_{\theta_{t-1}}(x,a)$. Intuitively, if we can find a parameter $\hat{\theta}$ that increases the value of $f_{\theta}(x,a)$ without significantly increasing $\mathcal{R}(\theta;D)$, then the uncertainty on the reward of $(x,a)$ would be large. Consequently, a large bonus would be added to such an action by Eq. (5). Moreover, $g$ and $\eta$ cooperatively control the scale of the confidence bonus: the difference between $f_{\hat{\theta}(x,a)}(x,a)$ and $f_{\theta_{t-1}}(x,a)$ increases as $\eta$ decreases and $g$ controls the contribution of such difference to the final bonus. In Section 3.1, we show that with proper $g,\eta,\mathcal{L}$ and $\mathcal{R}$ we can achieve near-optimal regret bound on contextual bandit when the reward function is simple.

The contextual bandit problem has been extensively studied and applied in machine learning. Besides OFU, the most used and studied contextual bandit algorithms can be mainly divided into two categories: $\epsilon$-greedy sutton2018reinforcement and Thompson sampling thompson1933likelihood .

$\epsilon$-greedy greedily selects the action $a=\operatorname*{arg\,max}_{a^{\prime}}f_{\theta_{t-1}}(x,a^{\prime})$ with probability $1-\epsilon$ and randomly selects an action to explore with probability $\epsilon$. While being sub-optimal on EE trade-off, $\epsilon$-greedy is simple, general and computationally efficient. Thus, $\epsilon$-greedy is the most widely used algorithm in complex tasks.

The basic idea of Thompson sampling thompson1933likelihood is to maintain a posterior over parameters and make decisions according to the samples from the posterior. However, Thompson sampling can be applied only if we can access the posterior and use approximate inference methods which can be computationally inefficient and may significantly hurt the performance of regret deep_bayesian_bandit_showdown ; phan2019thompson .

There are several works extending OFU and Thompson sampling to the case of non-linear contextual bandits. NeuralUCB zhou2020neural and Neural-TS zhang2021neural extend OFU and TS to NTK bandits and they provide regret bounds in the order of $\tilde{O}(\sqrt{T})$. Our analysis for NTK bandit is largely inspired by their work. However, as mentioned in Sec. 2, both NeuralUCB and NeuralTS need to calculate the inversion of a matrix with dimension equal to the number of parameters in DNN, which is computationally inefficient for large DNN models. Moreover, the neural-linear model is studied in zahavy2019deep ; deep_bayesian_bandit_showdown , where they take the output of the second last layer of a DNN as feature map. These algorithms sometimes work well, but have no theoretical guarantee on regret.

We now empirically evaluate ROFU. For simplicity, in all our experiments, we set $\eta=1,g(w)=\sqrt{w}$, $\mathcal{L}(\theta;D)=\text{MSE}(\theta;D)$ and $\mathcal{R}(\theta;D)=|D|\text{MSE}(\theta;D)$. That is, $\theta_{t-1}$ is trained to minimize $\text{MSE}(\theta;D)$ with standard optimizer ¹¹1We train $\theta_{t-1}$ with stochastic gradient descent starting from $\theta_{t-2}$ in all the experiments. and $\hat{\theta}(x,a)$ is trained to maximize $f_{\theta}(x,a)-|D|\text{MSE}(\theta;D)$ using Alg. 2.

In this work, we propose an OFU variant, called ROFU, which can be applied to non-linear contextual bandit. We show that the regret of ROFU is (near-)optimal on various contextual bandit models. Moreover, we propose an efficient algorithm to approximately compute the upper confidence bound. Thus, ROFU is efficient in both computation and EE trade-off, which are empirically verified by our experimental results.

EE trade-off is a fundamental problem that lies in the heart of sequential decision making. However, the huge computational cost of (near-)optimal EE trade-off algorithms significantly limits the application, especially on complex domain. We hope our method could inspire more algorithms to efficiently trade-off EE for sequential decision-making tasks beyond contextual bandit, such as deep reinforcement learning.