Neural Bandit with Arm Group Graph

Contextual bandits are a specific type of multi-armed bandit (MAB) problem where the learner has access to the contextual information (contexts) related to arms at each round, and the learner is required to make recommendations based on past contexts and received rewards. A variety of models and algorithms have been proposed and successfully applied on real-world problems, such as online content and advertising recommendation (Li et al., 2010; Wu et al., 2016), clinical trials (Durand et al., 2018; Villar et al., 2015) and virtual support agents (Sajeev et al., 2021).

In this paper, we focus on exploiting the accessible arm information to improve the performance of bandit algorithms. Among different types of contextual bandit algorithms, upper confidence bound (UCB) algorithms have been proposed to balance between exploitation and exploration (Auer et al., 2002; Chu et al., 2011; Valko et al., 2013). For conventional UCB algorithms, they are either under the ”pooling setting” (Chu et al., 2011) where one single UCB model is applied for all candidate arms, or the ”disjoint setting” (Li et al., 2010) where each arm is given its own estimator without the collaboration across different arms. Both settings have their limitations: applying only one single model may lead to unanticipated estimation error when some arms exhibit distinct behaviors (Wu et al., 2019b, 2016); on the other hand, assigning each arm its own estimator neglects the mutual impacts among arms and usually suffers from limited user feedback (Gentile et al., 2017; Ban and He, 2021b).

To deal with this challenge, adaptively assigning UCB models to arms based on their group information can be an ideal strategy, i.e., each group of arms has one estimator to represent its behavior. This modeling strategy is linked to ”arm groups” existing in real-world applications. For example, regarding the online movie recommendation scenario, the movies (arms) with the same genre can be assigned to one (arm) group. Another scenario is the drug development, where given a new cancer treatment and a patient pool, we need to select the best patient on whom the treatment is most effective. Here, the patients are the arms, and they can be naturally grouped by their non-numerical attributes, such as the cancer types. Such group information is easily accessible, and can significantly improve the performance of bandit algorithms. Although some works  (Li et al., 2016; Deshmukh et al., 2017) have been proposed to leverage the arm correlations, they can suffer from two common limitations. First, they rely on the assumption of parametric (linear / kernel-based) reward functions, which may not hold in real-world applications (Zhou et al., 2020). Second, they both neglect the correlations among arm groups. We emphasize that the correlations among arm groups also play indispensable roles in many decision-making scenarios. For instance, in online movie recommendation, with each genre being a group of movies, the users who like ”adventure movies” may also appreciate ”action movies”. Regarding drug development, since the alternation of some genes may lead to multiple kinds of tumors (Shao et al., 2019), different types of cancer can also be correlated to some extent.

To address these limitations, we first introduce a novel model, AGG (Arm Group Graph), to formulate non-linear reward assumptions and arm groups with correlations. In this model, as arm attributes are easily accessible (e.g., movie’s genres and patient’s cancer types), the arms with the same attribute are assigned into one group, and represented as one node in the graph. The weighted edge between two nodes represents the correlation between these two groups. In this paper, we assume the arms from the same group are drawn from one unknown distribution. This also provides us with an opportunity to model the correlation of two arm groups by modeling the statistical distance between their associated distributions. Meantime, the unknown non-parametric reward mapping function can be either linear or non-linear.

Then, with the arm group graph, we propose the AGG-UCB framework for contextual bandits. It applies graph neural networks (GNNs) to learn the representations of arm groups with correlations, and neural networks to estimate the reward functions (exploitation). In particular, with the collaboration across arm groups, each arm will be assigned with the group-aware arm representation learned by GNN, which will be fed into a fully-connected (FC) network for the estimation of arm rewards. To deal with the exploitation-exploration dilemma, we also derive a new upper confidence bound based on network-gradients for exploration. By leveraging the arm group information and modeling arm group correlations, our proposed framework provides a novel arm selection strategy for dealing with the aforementioned challenges and limitations. Our main contributions can be summarized as follows:

• •

  First, motivated by real-world applications, we introduce a new graph-based model in contextual bandits to leverage the available group information of arms and exploit the potential correlations among arm groups.

• •

  Second, we propose a novel UCB-based neural framework called AGG-UCB  for the graph-based model. To exploit the relationship of arm groups, AGG-UCB  estimates the arm group graph with received contexts on the fly, and utilizes GNN to learn group-aware arm representations.

• •

  Third, we prove that AGG-UCB can achieve a near-optimal regret bound in the over-parameterized neural works, and provide convergence analysis of GNN with fully-connected layers, which may be of independent interest.

• •

  Finally, we conduct experiments on publicly available real data sets, and demonstrate that our framework outperforms state-of-the-art techniques. Additional studies are conducted to understand the properties of the proposed framework.

The rest of this paper is organized as following. In Section 2, we briefly discuss related works. Section 3 introduces the new problem settings, and details of our proposed framework AGG-UCB  will be presented in Section 4. Then, we provide theoretical analysis for AGG-UCB in Section 5. After presenting experimental results in Section 6, we finally conclude the paper in Section 7. Due to the page limit, readers may refer to our arXiv version of the paper for the supplementary contents (https://arxiv.org/abs/2206.03644).

In this section, we briefly review the related work on contextual bandits. Lin-UCB (Chu et al., 2011) first formulates the reward estimation through a linear regression with the received context and builds a confidence bound accordingly. Kernel-UCB (Valko et al., 2013) further extends the reward mapping to the Reproducing Kernel Hilbert Space (RKHS) for the reward and confidence bound estimation under non-linear settings. Besides, there are algorithms under the non-linear settings. Similarly, CGP-UCB (Krause and Ong, 2011) models the reward function through a Gaussian process. GCN-UCB (Upadhyay et al., 2020) applies the GNN model to learn each context an embedding for the linear regression. Then, Neural-UCB (Zhou et al., 2020) proposes to apply FC neural network for reward estimations and derive a confidence bound with the network gradient, which is proved to be a success, and similar ideas has been applied to some other models (Ban et al., 2021; Zhang et al., 2020; Ban and He, 2021a). (Ban et al., 2022b) assigns another FC neural network to learn the confidence ellipsoid for exploration. Yet, as these works consider no collaboration among estimators, they may suffer from the performance bottleneck in the introduction.

To collaborate with different estimators for contextual bandits, various approaches are proposed from different perspectives. User clustering algorithms (Gentile et al., 2014; Li et al., 2019; Ban and He, 2021b; Ban et al., 2022a) try to cluster user with alike preferences into user groups for information sharing while COFIBA (Li et al., 2016) additionally models the relationship of arms. Then, KMTL-UCB (Deshmukh et al., 2017) extends Kernel-UCB to multi-task learning settings for a refined reward and confidence bound estimation. However, these works may encounter with performance bottlenecks as they incline to make additional assumptions on the reward mapping by applying parametric models and neglect the available arm group information.

GNNs (Kipf and Welling, 2016; Wu et al., 2019a; Klicpera et al., 2018; Fu and He, 2021b) are a kind of neural models that deal with tasks on graphs, such as community detection (You et al., 2019), recommender systems (Wei et al., 2020) and modeling protein interactions (Fu and He, 2021a). GNNs can learn from the topological graph structure information and the input graph signal simultaneously, which enables AGG-UCB  to cooperate with different arm groups by sharing information over the arm group neighborhood.

We suppose a fixed pool $\mathcal{C}=\{1,\dots,N_{c}\}$ for arm groups with the number of arm groups being $\lvert\mathcal{C}\rvert=N_{c}$, and assume each arm group $c\in\mathcal{C}$ is associated with an unknown fixed context distribution $\mathcal{D}_{c}$. At each time step $t$, we will receive a subset of groups $\mathcal{C}_{t}\subseteq\mathcal{C}$. For each group $c\in\mathcal{C}_{t}$, we will have the set of sampled arms $\mathcal{X}_{c,t}=\{\boldsymbol{x}^{(1)}_{c,t},\cdots\boldsymbol{x}^{(n_{c,t})}_{c,t}\}$ with the size of $\lvert\mathcal{X}_{c,t}\rvert=n_{c,t}$. Then, $\forall i\in[n_{c,t}]=\{1,\dots,n_{c,t}\}$, we suppose $\boldsymbol{x}^{(i)}_{c,t}\sim\mathcal{D}_{c}$ with the dimensionality $\boldsymbol{x}^{(i)}_{c,t}\in\mathbb{R}^{d_{x}}$. Therefore, in the $t$-th round, we receive

With $\boldsymbol{W}^{*}\in\mathbb{R}^{N_{c}\times N_{c}}$ being the unknown affinity matrix encoding the true arm group correlations, the true reward $r_{c,t}^{(i)}$ for arm $\boldsymbol{x}^{(i)}_{c,t}$ is defined as

where $h(\cdot)$ represents the unknown reward mapping function, and $\epsilon$ is the zero-mean Gaussian noise. For brevity, let $\boldsymbol{x}_{t}$ be the arm we select in round $t$ and $r_{t}$ be the corresponding received reward.

Our goal is recommending arm $\boldsymbol{x}_{t}$ (with reward $r_{t}$) at each time step $t$ to minimize the cumulative pseudo-regret $R(T)=\sum_{t=1}^{T}\mathbb{E}\left[(r_{t}^{*}-r_{t})\right]$ where $\mathbb{E}[r_{t}^{*}]=\max_{(c\in\mathcal{C}_{t},i\in[n_{c,t}])}\big{[}h(\boldsymbol{W}^{*},\boldsymbol{x}^{(i)}_{c,t})\big{]}$. At each time step $t$, the overall set of received contexts is defined as $\mathcal{X}_{t}=\{\boldsymbol{x}^{(i)}_{c,t}\}_{c\in\mathcal{C}_{t},i\in[n_{c,t}]}$. Note that one arm is possibly associated with multiple arm groups, such as a movie with multiple genres. In other words, for some $c,c^{\prime}\in\mathcal{C}_{t}$, we may have $\mathcal{X}_{c,t}\cap\mathcal{X}_{c^{\prime},t}\neq\emptyset$.

In order to model the arm group correlations, we maintain an undirected graph $\mathcal{G}_{t}=(V,E,W_{t})$ at each time step $t$, where each arm group from $\mathcal{C}$ is mapped to a corresponding node in node set $V$. Then, $E=\{e(c_{i},c_{j})\}_{c_{i},c_{j}\in\mathcal{C}}$ is the set of edges, and $W_{t}$ represents the set of edge weights. Note that by definition, $\mathcal{G}_{t}$ will stay as a fully-connected graph, and the estimated arm group correlations are modeled by the edge weights connecting pairs of nodes. For a node $v\in V$, we denote the augmented $k$-hop neighborhood $\widetilde{\mathcal{N}}_{k}(v)=\mathcal{N}_{k}(v)\cup\{v\}$ as the union node set of its $k$-hop neighborhood $\mathcal{N}_{k}(v)$ and node $v$ itself. For the arm group graph $\mathcal{G}_{t}$, we denote $\boldsymbol{A}_{t}\in\mathbb{R}^{N_{c}\times N_{c}}$ as the adjacency matrix (with added self-loops) of the given arm group graph and $\boldsymbol{D}_{t}\in\mathbb{R}^{N_{c}\times N_{c}}$ as its degree matrix. For the notation consistency, we will apply a true arm group graph $\mathcal{G}^{*}$ instead of $\boldsymbol{W}^{*}$ in Eq. 2 to represent the true arm group correlation.

In this section, we start with an overview to introduce our proposed AGG-UCB framework. Then, we will show our estimation method of arm group graph before mentioning the related group-aware arm embedding. Afterwards, the two components of our proposed framework, namely the aggregation module and the reward-estimation module, will be presented.

In this section, we provide the theoretical analysis for our proposed framework. For the sake of analysis, at each time step $t$, we assume each arm group $c\in\mathcal{C}$ would receive one arm $\boldsymbol{x}_{c,t}$, which makes $\lvert\mathcal{X}_{1}^{t}\rvert=\dots=\lvert\mathcal{X}_{N_{c}}^{t}\rvert=t$. We also apply the adjacency matrix $\boldsymbol{A}_{t}$ instead of $\boldsymbol{S}_{t}$ for aggregation, and set its elements $[\boldsymbol{A}_{t}]_{ij}=\frac{1}{t\cdot N_{c}}\sum_{\tau=1}^{t}\phi_{k_{\mathcal{G}}}(x_{c_{i},\tau})^{\intercal}\phi_{k_{\mathcal{G}}}(x_{c_{j},\tau})$ for arm group similarity between group $c_{i},c_{j}\in\mathcal{C}$. Here, $\phi_{k_{\mathcal{G}}}(\cdot)$ is the kernel mapping given an RBF kernel $k_{\mathcal{G}}$. With $\mathcal{G}^{*}$ being the unknown true arm group graph, its adjacency matrix elements are $[\boldsymbol{A}^{*}]_{ij}=\frac{1}{N_{c}}\mathbb{E}_{x_{i}\sim\mathcal{D}_{c_{i}},x_{j}\sim\mathcal{D}_{c_{j}}}(\phi_{k_{\mathcal{G}}}(x_{i})^{\intercal}\phi_{k_{\mathcal{G}}}(x_{j}))$. Note that the norm of adjacency matrices $\lVert\boldsymbol{A}^{*}\rVert_{2},\lVert\boldsymbol{A}_{t}\rVert_{2}\leq 1$ since $\langle\phi_{k_{\mathcal{G}}}(x),\phi_{k_{\mathcal{G}}}(x^{\prime})\rangle\leq 1$ for any $x,x^{\prime}\in\mathbb{R}^{d_{x}}$, which makes it feasible to aggregate over $k$-hop neighborhood without the explosion of eigenvalues. Before presenting the main results, we first introduce the following background.

Note that with sufficient network width $m$, we will have $\boldsymbol{\Theta}_{t-1}^{*}\approx\boldsymbol{\Theta}_{0}^{*},\forall t\in[T]$, and we will include more details in the full version of the paper. Following the analogous ideas from previous works (Zhou et al., 2020; Ban et al., 2021), this lemma formulates the expected reward as a linear function parameterized by the difference between randomly initialized network parameter $\boldsymbol{\Theta}_{0}$ and the parameter $\boldsymbol{\Theta}_{t-1}^{*}$, which lies in the confidence set with the high probability (Abbasi-Yadkori et al., 2011). Then, regarding the activation function $\sigma(\cdot)$, we have the following assumption on its continuity and smoothness.

Note that Assumption 5.2 is mild and applicable on many activation functions, such as Sigmoid. Then, we proceed to bound the regret for a single time step $t$.