Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
74,216 characters · 16 sections · 62 citation commands
The Role of Contextual Information in Best Arm Identification
\begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let \begingroup \def\fnsymbol{footnote}{\fnsymbol{footnote}} \def\hbox to 0pt{$^{\@thefnmark}$\hss}{\hbox to 0pt{$^{\@thefnmark}$\hss}} \deffootnote[1.7em]{1.6em}{2em}{$^\thefootnotemark$} \@maketitle \@thanks \endgroup \setcounter{footnote}{0} \let\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax\relax \let\@maketitle\relax \gdef\@thanks\gdef\@author\gdef\@title\let\thanks\relax
This paper studies best-arm identification (BAI) with contextual information in stochastic multi-armed bandit (MAB) problems. We define the best arm as the arm with the maximum marginalized mean reward, where the expectation is defined over the context distribution, not on a specific context. We call this setting contextual BAI. The goal is to identify the best arm with a fixed confidence level and a smaller sample complexity defined by the probably approximately correct (PAC) framework. The instance-specific sample complexity of BAI without contextual information is now well understood. There exists an instance-specific lower bound Kaufman2016complexity,Garivier2016 and optimal algorithms whose performance guarantee matches the lower bound Kaufman2016complexity, Garivier2016, degenne2019non; however, that of contextual BAI has never been elucidated.
Formally, we consider the following setting. At each time $t=1,2,\dots$, an agent observes a context (covariate) $X_t\in\mathcal{X}$ and chooses an arm $A_t \in [K] = \{1,\dots, K\}$, where $\mathcal{X}$ denotes the context space. Then, the agent immediately receives a reward (or outcome) $R_t$ linked to the arm $A_t$. This setting is called the bandit feedback or Rubin causal model Neyman,Rubin1974; that is, a reward in round $t$ is $R_t=\sum^K_{a=1}\mathbbm{1}[A_t = a]R_{t, a}$, where $R_{t, a}$ is a potential independent (random) reward. We assume that $X_t$ is independent and identically distributed (i.i.d.) over $[T]$ and denote the distribution of $X_t$ by $\zeta$. Given the context $x \in \mathcal{X}$, we denote the reward distributions of the potential outcomes as $\boldsymbol{p} = (p_{1, x}, p_{2, x}, \ldots, p_{K, x})$ and their means as $\boldsymbol{\mu} = (\mu_{1, x}, \mu_{2, x}, \ldots, \mu_{K, x})$. Let $\mathcal{V} = (\boldsymbol{p}, \zeta)$ (this can be written as $ \nu = ((\mu_{a,x}), (\zeta_x))$ when the rewards follow a distribution that belongs to a single parameter exponential family, and the contexts are finite) be a bandit problem. Let $\mathbb{P}_{\mathcal{V}}$ (resp. $\mathbb{P}_{\nu}$ ) and $\mathbb{E}_{\mathcal{V}}$ (resp. $(\mathbb{E}_{\nu})$) be the probability and expectations under model $\mathcal{V}$ (resp. $\nu$), respectively. Then, $\mu_a = \mathbb{E}_{X \sim \zeta}[\mu_{a, X}] = \mathbb{E}_{X \sim \zeta}[\mathbb{E}_{\mathcal{V}}[R_{t, a} | X]] = \mathbb{E}_{\mathcal{V}}[R_{t, a}]$ is the average reward marginalized over $\mathcal{X}$. We assume that $\mathcal{V}$ belongs to a class $\Omega = \{ (\boldsymbol{p}, \zeta): \exists a^* \in [K] \;s.t.\; \forall a \neq a^*, \mu_{a^*} > \mu_{a} \}$; that is, the best arm $a^*(\mathcal{V}) = \operatorname*{arg\,max}_a \mu_a$ is uniquely defined. Let $p_{a, x}$ and $q_{a, x}$ be two absolutely continuous probability distributions (w.r.t. the Lebesgue measure) of $R_{t, a}$, given $X_{t} = x$. We define the Kullback--Leibler (KL) divergence from $p_{a, x}$ to $ q_{a, x}$ as
We assume that for all $(\boldsymbol{p}, \zeta), (\boldsymbol{q}, \zeta) \in \Omega$, if $p_{a, x}\neq q_{a, x}$, then $0 < \mathrm{KL}(p_{a, x}, q_{a, x}) < +\infty$. For distributions that belong to the single parameter exponential family, we introduce the KL divergence from the distribution with mean $\mu$ to the distribution with mean $\nu$ as $\mathrm{kl}(\mu, \nu)$. Furthermore, for the Bernoulli distributions, we denote the KL divergence by $d(\mu, \nu) =\mu\log (\mu/\nu) + (1-\mu)\log((1-\mu)/(1-\nu))$ with the convention that $d(0, 0) = d(1,1) = 0$.
Let $\mathcal{F}_t = \sigma(X_1, A_1, R_1, \ldots, X_t, A_t, R_t, X_{t+1})$ and $\mathcal{G}_t = \sigma(X_1, A_1, R_1, \ldots, X_t, A_t, R_t)$ be the sigma-algebras generated by the observations up to immediately before the selection of the arm at time $t+1$ and all observations up to time $t$, respectively. The strategy or algorithm of the best arm identification consists of the following three elements: a sampling, stopping, and decision rules. A sampling rule selects from which arm we collect the sample each time based on past observation ($ A_t$ is $\mathcal{F}_{t-1}$-measurable). The stopping rule determines when to stop sampling based on the past observation. We denote $\tau$ as this time; $\tau$ is the stopping time with respect to the filtration $(\mathcal{G}_t)_{t\ge1}$. The decision rule estimates the best arm $\hat{a}_\tau$ based on observation up to time $\tau$ ($\hat{a}_\tau$ is $\mathcal{G}_\tau$-measurable).
We focus on the fixed confidence setting; that is, with a given admissible failure probability $\delta \in (0, 1)$, the algorithm is guaranteed to have $\mathbb{P}( \hat{a}_\tau \neq \operatorname*{arg\,max}_a \mu_a) \leq \delta$. We define $\delta$-PAC to formalize this property:
Later, we propose algorithms that are $\delta$-PAC.
We reemphasize that although we can use contextual information, our primary interest is not in the mean reward conditioned on each context. Similar problems are frequently considered in the literature on causal inference that mainly discusses the efficient estimation of causal parameters. The assigned treatment (chosen arm) and observed outcomes for each treatment (reward) and covariate (context) are given therein. Here, we are not interested in the distribution of the covariate; rather we are interested in the estimation of the expected value of the outcome of the treatment marginalized over the covariate distribution; that is, the average treatment effect (ATE) imbens_rubin_2015. For this setting, Laan2008TheCA and Hahn2011 proposed experimental design methods to estimate the ATE more efficiently by assigning treatments based on the covariate. According to their results, even if the covariates are marginalized, the variance of the estimator can be reduced with the help of the covariate information. Karlan2014 applied the method of Hahn2011 to test how donors respond to new information about the effectiveness of charity. These studies have been attempted to be improved by Meehan2018 and kato2020efficienttest.
For each $x\in\mathcal{X}$, we define allocations for each arm with the context $x$ as $\Sigma^K_x = \{w_{a, x}\in\mathbb{R}^k_+: w_{1, x} + \cdots + w_{K, x} = 1\}$. Let $\mathcal{W}$ be all possible such allocations. We denote by $N_x(t)$ and $N_{a, x}(t)$ the number of times we observe context $x$, and we choose arm $a$ given context $x$; that is, $N_x(t) = \sum_{s=1}^t \mathds{1}\{ X_s = x\}$ and $N_{a, x}(t) = \sum_{s=1}^t \mathds{1}\{X_s = x, A_s = a\}$, respectively.
\paragraph{Main results.} We briefly summarize our contributions.
First, we establish the instance-specific lower bound on contextual BAI for both continuous and finite context cases. The derived lower bound formula has smaller sample complexity than that of lower bound formula in Garivier2016, suggesting that a faster BAI may be possible.
Then, we propose optimal algorithms for two cases: (i) two-armed Gaussian bandits where the arms and context jointly follow the multivariate normal distribution; and (ii) MAB with reward distributions belonging to the single parameter exponential family and finite contexts. We prove that the sample complexity upper bounds of the proposed algorithms asymptotically match the lower bounds.
\paragraph{Organization.} This paper is organized as follows. In Section (ref), we derive the general instance-specific lower bounds for contextual BAI for a case with continuous contexts. Then, in Section (ref), we discuss an optimal algorithm for two-armed Gaussian bandits with continuous contexts. Section (ref) focuses on the lower bound when the number of contexts is finite, and the reward distributions are from the single parameter exponential family. In Section (ref), for the finite context case, we obtain the optimal allocations for each pair of contexts and actions by simplifying the lower bound formula. In Section (ref), in the same setting of Section (ref), we show an optimal algorithm. We describe details of the sampling, stopping, and decision rules that are the core of the proposed algorithm and demonstrate that the algorithm is $\delta$-PAC. We further confirm that the sample complexity of the proposed algorithm is asymptotically optimal. Section (ref) presents the results of our numerical experiments.
\paragraph{Related work.} The stochastic MAB problem is a classical abstraction of the sequential decision-making problem Thompson1933,Robbins1952,Lai1985. BAI is a paradigm of the MAB problem, where we consider pure exploration to find the best arm. Several strategies and efficiency metrics have been proposed for BAI bechhofer1968sequential,Paulson1964,Mannor2004,EvanDar2006,Bubeck2011,Gabillon2012,Karnin2013,Garivier2016,Jamieson2014. BAI with linear bandits Soare2014,Xu2018,Tao2018,Fiez2019,jedra2020optimal, BAI with multiple queries, and the partition identification problem Juneja2019 are different directions for the generalization of BAI.
Our setting is a generalization of BAI without contextual information. We can use the side information (explicitly or implicitly) at each round. There have been limited studies that address pure exploration in contextual bandits. Tekin2015, GuanJiang2018, and Deshmukh2018 also consider BAI with contextual information; however, they do not discuss the instance-specific optimality. After this study, Qin2022 also considers a related topic.
From the causal inference perspective, contextual BAI is closely related to a (semiparametric) experimental design for efficient ATE estimation Laan2008TheCA,Hahn2011,Karlan2014,Athey2016,Meehan2018. The goal of efficient ATE estimation by adaptive experimentation is often in choosing the best treatment (arm) via hypothesis testing. Therefore, it can be considered as a case where the proposed method should be applied, especially when there are multiple treatments (arms).
Russac2021 also addresses a similar problem independently of us. Their problem setting is the same as ours in that they can observe discrete contexts. However, they are considering a slightly different problem than best arm identification, i.e., A/B/n testing, where they consider the comparison with a designated control arm. In that problem setting, optimal allocation is uniquely obtained, and they do not have to consider multiple candidates of optimal allocations as we do. Besides, we also derive the result for the case of continuous contexts, which they do not address. On the other hand, they discuss the problem more generally by considering four situations, (a) active mode, (b) proportional mode, (c) agnostic mode, and (d) oblivious mode, depending on how the decision is made. The (b) proportional mode discussed by them is closer to the setting discussed in this paper. In these senses, our results and theirs, while similar, are independent and parallel, and correspond to complementary studies.
In this section, we provide the instance-specific sample complexity lower bounds for general contextual BAI. The proof is based on standard change-of-measure arguments Kaufman2016complexity. However, the derivations must consider the possibly continuous context distributions, which are non-trivial.
Based on the lower bound, we find that the contextual information either helps or does not harm the BAI. Our result is the same as those of existing studies on fixed-confidence BAI without contextual information, except that we can obtain help from the existence of the contextual information. At first glance, it does not necessarily seem advantageous to use contextual information as the marginalized mean reward is not directly related to the contextual information. However, the lower bound with contextual information (see Section (ref)) is strictly lower than the sample complexity derived by Kaufman2016complexity and Garivier2016.
Assume $\mathcal{X} = \mathbb{R}$. Then, we present the non-asymptotic sample complexity lower bound.
We provide the proof of Theorem (ref) in Appendix (ref).
\paragraph{Efficiency gains from the context use.} In Figure (ref), we illustrate the efficiency gain by using contextual information. We consider a two-armed, one-dimensional context $X_t\in\mathbb{R}$. Suppose that $(R_{t,1} \ R_{t,2} \ X_t)^\top$ follows a multivariate normal distribution with mean vector $(1\ 0\ 0)^\top$. We assume that the variances of $R_{t,1}$, $R_{t,2}$, and $X_t$ are $1$. We investigate the variation in the theoretical sample complexity by varying the correlation coefficients between $X_t$ and $R_{t,1}$ and $X_t$ and $R_{t,2}$, which are denoted as $\rho_{1\mathcal{X}}\in[0,1]$ and $\rho_{2\mathcal{X}}\in[0,1]$, respectively. Note that we omit the other domains due to symmetry with the current domain. Note that when ignoring (marginalizing) the context, arm $1$ follows $\mathcal{N}(1, 1)$ and arm $2$ follows $\mathcal{N}(0, 1)$, where $\mathcal{N}(\mu, \sigma^2)$ denotes a normal distribution with a mean $\mu$ and variance $\sigma^2$. Here, for $\delta=0.05$, we calculate the sample complexity lower bounds of the standard setting of BAI from the result of Garivier2016 and those of the contextual case from our results. We denote the former as $\ell$ and the latter as $\widetilde{\ell}$. Then, we compute the sample complexity gain ($1 - \widetilde{\ell}/ \ell$) for different pairs of $(\rho_{1\mathcal{X}}, \rho_{2\mathcal{X}})$ and illustrate it in Figure (ref).
In this section, we provide an example for the case of continuous contexts and prove the upper bound of the sample complexity. We consider the following two-armed bandit problem. For each $t$, $R_{t,1}$, $R_{t,2}$, and $X_t\in\mathbb{R}$ are drawn from the following Gaussian distributions $\mathcal{N}(\mu_1, \sigma^2_1)$, $\mathcal{N}(\mu_2, \sigma^2_2)$, and $\mathcal{N}(\mu_{\mathcal{X}}, \sigma_{\mathcal{X}}^2)$, respectively ($\mu_1 > \mu_2$). Assume that the vector $(R_{t,1}, R_{t,2}, X_t)$ forms a multivariate Gaussian distribution. We denote $ \textnormal{Cov}(R_{t,1}, X_t) = \sigma_{\mathcal{X}1}$ and $ \textnormal{Cov}(R_{t,2}, X_t) = \sigma_{\mathcal{X}2}$. Suppose the algorithm knows that $(R_{t,1}, R_{t,2}, X_t)$ form a multivariate Gaussian distribution, knows the values of $\sigma^2_1$, $\sigma^2_2$, $ \mu_{\mathcal{X}}$, $ \sigma_{\mathcal{X}}$, $\sigma_{\mathcal{X}1}$, and $ \sigma_{\mathcal{X}2}$, and does not know the values of $\mu_1$ and $\mu_2$. Let $ \widetilde{\Omega}$ be a set of all such problems. Given an observation $X_t = x$, we have conditional distributions of $R_{t,1}$ and $R_{t,2}$ where for each $a \in \{1, 2\}$, $ R_{t,a} \sim \mathcal{N}\left(\mu_a + \frac{\sigma_{\mathcal{X}a}}{\sigma_{\mathcal{X}}^2}(x - \mu_{\mathcal{X}}), \sigma^2_a - \frac{\sigma_{\mathcal{X}1}^2}{\sigma_{\mathcal{X}}^2}\right) =\mathcal{N}\left(\mu_a + \frac{\rho_{\mathcal{X}a} \sigma_a}{\sigma_{\mathcal{X}}}(x - \mu_{\mathcal{X}}), \sigma^2_a(1 - \rho^2_{\mathcal{X}a})\right)$. Here, $\rho_{\mathcal{X}a}$ is the correlation coefficient between the context and arm $a\in\{1,2\}$. We denote $\sigma_1'^2 = \sigma^2_1 - \frac{\sigma_{\mathcal{X}1}^2}{\sigma_{\mathcal{X}}^2}$ and $ \sigma_2'^2 = \sigma^2_2 - \frac{\sigma_{\mathcal{X}2}^2}{\sigma_{\mathcal{X}}^2}$. From our lower bound in Theorem (ref), we can derive the following lower bound for this specific problem. We give the proof in Appendix (ref).
Note that when $\sigma_{\mathcal{X}1}^2>0$ or $\sigma_{\mathcal{X}2}^2>0$, $ \sigma_1 + \sigma_2> \sigma_1' + \sigma_2'$; that is, the value of the lower bound derived in Theorem (ref) is strictly smaller than that of the lower bound derived by Kaufman2016complexity, $\frac{2(\sigma_1+ \sigma_2)^2}{(\mu_1 - \mu_2)^2}\mathrm{kl}(\delta, 1-\delta)$. Let $\alpha = {\sigma_1'}/{(\sigma_1' + \sigma_2')}$. We also note that the simple $\alpha$-elimination algorithm by Kaufman2016complexity with $\alpha$ achieves the lower bound as well as a strictly better sample complexity than that given in Kaufman2016complexity. We give the proof in Appendix (ref).
Hence, $\alpha$-elimination is optimal for this problem. The details of $\alpha$-elimination with contextual information is shown in Appendix (ref). The pseudo-code is shown in Algorithm (ref). Thus, apparently irrelevant contextual information improves optimal sample complexity.
Although we derived the optimal algorithm for BAI with continuous contexts in the previous sections, it requires some assumptions that may not be practical, e.g., multivariate normal distribution and known variance. We also consider a more practical algorithm by considering BAI with finite contexts. In this section, we consider a lower bound when the number of contexts is finite. For $\nu = ((\mu_{a,x})_{a,x}, (\zeta_x))$, we suppose that $\mathcal{X}$ is finite ($\zeta$ follows the multinomial distribution), and for each arm $a$ and context $x$, arm distribution belongs to the canonical one-parameter exponential family Cappe2013,Kaufman2016complexity,Garivier2016,Juneja2019:
where $\lambda$ is some reference measure on $\mathbb{R}$, $b:\Pi\mapsto \mathbb{R}$ is a convex, twice differential function, and $\Pi \subset \mathbb{R} $ is a parameter space. Note that a distribution $p_\pi \in \mathcal{P}$ can be parameterized by its mean $\dot{b}(\pi)$. As discussed in Cappe2013,Garivier2016, the KL divergence from $p_\pi$ to $p_{\pi'}$ is given by
For each arm $a$ and context $x$ pair, we represent the unique distribution in $\mathcal{P}$ by $(\mu_{a,x})$. We further write the multinomial contextual distribution by $(\zeta_x)$.
We denote by $\Theta$ a set of BAI problems with finite contexts and the single parameter (canonical) exponential family. The lower bound is given in the following theorem.
As an intuition behind $T^\star(\nu)$, the probability of misidentification is roughly $\exp (-\tau \left(T^\star(\nu)\right)^{-1})$; that is, larger $\left(T^\star(\nu)\right)^{-1}$ means a strategy with smaller sample complexity.
We note specific properties of this lower bound. From the results in Garivier2016, we know that when the optimal arm is unique, the expected value of the sampling budget of the optimal BAI algorithm does not diverge; rather it is less than or equal to the order of $\log(1/\delta)$. Therefore, from the assumption of the proposed model, $T^\star(\nu)$ is finite under certain regularity conditions; for example, the context marginalized distribution of the reward $ R_t$ is sub-Gaussian.
To derive the lower bound, we show the following lemma, which is an extension of Lemma 1 of Kaufman2016complexity.
The proof is provided in Appendix (ref). Here, we offer the proof sketch of Theorem (ref) as follows. \paragraph{Proof sketch.} From Lemma (ref) with $\mathcal{E} =\{\hat{a}_\tau = {a}^*(\nu)\}$, for each $\nu \in \Theta$ and $\nu' \in \mathrm{Alt}(\nu)$, we have
where, for the last inequality, we use the definition of the $\delta$-PAC algorithm and monotonicity of the KL divergence. Then, for each $\nu \in \Theta$, for some $(w_{a, x})_{a \in [K], x \in \mathcal{X}} \in \mathcal{W}$, we can obtain $\mathrm{kl}(\delta, 1- \delta) \le \operatorname{\mathbb{E}}_{\nu}[\tau_\delta] \sup_{w \in \mathcal{W}} \inf_{((\lambda_{a,x}), \zeta) \in \mathrm{Alt}(\nu)} \sum_{x \in \mathcal{X}} \zeta_x \sum_{a \in [K]} w_{a, x} \mathrm{kl}(\mu_{a, x}, \lambda_{a, x})$.
In Section (ref), we explain that the lower bound with contextual information is smaller than or equal to the lower bound without contextual information shown by Garivier2016.
In this section, we first provide a simplification of the lower bound derived in Section (ref). Then, we examine the characteristics of the optimal allocations used in the proof. It becomes apparent that the set of optimal allocations is, in general, not unique. Therefore, we define the notion of convergence to the set and prove that the estimated optimal allocations converge to the set of optimal allocations (even though they might not converge to a point).
Without loss of generality, let $a^*(\nu) = 1$. First, we show a simpler equivalence form for the optimization problem $T^\star(\nu)^{-1}$ in the following theorem.
We provide the proof in Appendix (ref).
Moreover, we can further simplify the constraint in the minimization problem. We define \[f_a((\lambda_{x,a})) = \sum_{x \in \mathcal{X}}\zeta_x\Bigg\{ w_{1, x}\mathrm{kl}(\mu_{1, x}, \lambda_{1, x}) + w_{a, x}\mathrm{kl}(\mu_{a, x}, \lambda_{a, x})\Bigg\}.\] Then, we show the following lemma.
We provide the proof in Appendix (ref).
Let $2^{\mathcal{W}}$ be a power set of ${\mathcal{W}}$. We define a point-to-set map $\Phi: \Theta \to 2^{\mathcal{W}}$; that is, the set of all optimal allocations for the bandit problem $\nu$ as $$ \Phi(\nu) = \Big\{\bm w \in \mathcal{W} \mid m(\bm w, \nu) = \max_{\boldsymbol{w}' \in \mathcal{W}} m(\boldsymbol{w}', \nu)\Big\}, $$ where
The interpretation of $m(\bm w, \nu)$ is that, unlike the corresponding part in Garivier2016, we can further minimize the lower bound by choosing an optimal allocation from a wider domain than the case without contextual information as long as the constraints are satisfied. For example, let us consider a case where two arms $a$ and $b$, and two contexts $1$ and $2$ are given. Here, under certain circumstances, one needs to think about saving the allocations to arm $a$ in context $1$, allocating more to arm $a$ in context $2$, and get more budget to arm $b$ in context $1$. Thus, solving $m(\bm w, \nu)$ is inherently different from optimizing the allocations separately for each context; that is, a case where we apply a BAI algorithm without contextual information for each discrete context such as Garivier2016.
From this simplified formula of the lower bound, we obtain the following lemmas. We provide the proofs in Appendix (ref)-(ref).
Note that the reason why $\nu$ is in $\mathbb{R}^{|\mathcal{X}|(K + 1)}$ is that we include $\zeta\in\mathbb{R}^{|\mathcal{X}|}$ in $\nu$ with $(\mu_{a,x})\in\mathbb{R}^{|\mathcal{X}|K}$.
The set of the optimal allocations is not, in general, unique. Therefore, we introduce the notion of convergence, where the metric is defined as the minimum distance from the point to the set.
Using this definition of convergence, we obtain the following lemmas. We provide the proofs in Appendix (ref)--(ref).
Here, we show that the lower bound with contextual information is smaller than or equal to the lower bound without contextual information shown by Garivier2016. For simplicity of discussion, we consider a two-armed bandit case. Let us denote the lower bound without contextual information by $\Gamma^\star(\nu)\mathrm{kl}(\mathbb{P}_{\nu}(\mathcal{E}), \mathbb{P}_{\nu'}(\mathcal{E}))$, where $\Gamma^\star(\nu)$ is defined as the same quantity as $T^\star(\bm{\mu})$ in Garivier2016. Let us also denote the optimal allocation in Garivier2016 by $\gamma^*_1$ and $\gamma^*_2$ and one of the optimal allocations in ours by $(w^*_{1, x})$ and $(w^*_{2, x})$. Then, $\Gamma^\star(\nu)^{-1} \leq T^\star(\nu)^{-1}$ holds as follows:
where for $(a)$, we use the convexity of the KL divergence. Next, we discuss when the equality holds. For brevity, we consider a case with only two contexts. Let us denote the optimal $\lambda_1$ in the case without contextual information by $\lambda^*_1$ (note that $\lambda_1 = \lambda_2$) and the optimal $(\lambda_{1,1}, \lambda_{1,2})$ and $(\lambda_{2,1}, \lambda_{2,2})$ in the case with contextual information by $(\lambda^*_{1,1}, \lambda^*_{1,2})$ and $(\lambda^*_{2,1}, \lambda^*_{2,2})$. Then, the equality holds only if the following three conditions simultaneously hold:
We believe that it is difficult to summarize these conditions in a simpler form, but except for cases where the expected reward does not change among contexts, situations satisfying these conditions are extremely limited.
In this section, we propose an optimal algorithm for contextual BAI, called the Contextual Track-and-Stop (CTS) algorithm for the case of finite context. The strategy is an extension of the Track-and-Stop (TS) algorithm by Garivier2016 for contextual BAI. We further prove that the proposed algorithm is $\delta$-PAC.
Recall that the optimal algorithm of BAI with fixed confidence Garivier2016 consists of sampling, stopping, and decision rules. We follow the same path for the contextual BAI. We show the pseudo-code of the proposed CTS algorithm in Algorithm (ref). There, the empirical averages $\hat{\mu}_{a,x}(t)$ and $\hat{\zeta}_x(t)$ are defined as for each $a \in [K]$ and $x \in \mathcal{X}$, $\hat{\mu}_{a,x}(t) = (\sum_{s=1}^tR_{t}\mathds{1}\{A_s = a, X_s = s\})/N_{a,x}(t)$ and $\hat{\zeta}_x(t) = (\sum_{s=1}^t\mathds{1}\{X_s = s\})/N_x(t)$. Our procedure is similar to TS with D-tracking, proposed by Garivier2016. However, incorporating contextual information is a non-trivial extension of their method. The algorithm consists of sampling, stopping, and decision rules. The detail of the sampling rule is described in the following Section (ref). The stopping rule, in particular, for determining the threshold $\beta(t, \delta)$, is described in Section (ref) when the reward distributions are Bernoulli and in Section (ref) when the reward distributions belong to the canonical one-parameter exponential family.
Our proposed algorithm consists of sampling, stopping, and recommendation rules. In the sampling rule, we use the forced exploration, which is an extension of D-tracking of Garivier2016 and is known to be empirically superior to their C-tracking. To estimate the optimal weights, we solve an empirically approximated optimization problem ((ref)) by applying optimization solvers directly. Several methods are proposed to solve the maximin problem more efficiently, such as the application of no-regret learning algorithms in degenne2019non. However, we cannot use them directly for solving contextual BAI, in which we have a different form of the maximin problem than that of BAI without context. jedra2020optimal (BAI with linear models) and Russac2021 (A/B/n texting with contextual information) also directly solve the maximin problem. In the stopping rule, we use the criterion proposed by Kaufmann2021, which refines the stopping rule of Garivier2016. Then, we recommend an arm with the maximum sample average of the reward.
To design an algorithm with minimal sample complexity, the sampling rule should match the optimal proportions of the arm draws; that is, an allocation in the set $\Phi(\nu)$. Because $\mu_{a,x}$ and $\zeta_x$ are unknown, our sampling rule tracks, in round $t$, the optimal allocations in the plug-in estimate $\Phi(\hat{\nu}(t))$, where $\hat{\nu}(t) = ((\hat{\mu}_{a,x}(t))_{a \in [K], x \in \mathcal{X}}, (\hat{\zeta}_x(t))_{x \in \mathcal{X}})$.
The design of our tracking rule is equivalent to computing a sequence of allocations $(w_{a,x}(t))_{t\ge 1}$. The only requirement we actually impose on this sequence is the following condition:
This condition is sufficient to guarantee the asymptotic optimality of the algorithm. We introduce a set $\varphi^g_{x,t} = \{ a : N_{a, x}(t) < \sqrt{N_{x}(t)} - K/2\}$ consisting of the context-action pairs that are poorly explored. Then, in round $t$, after observing a context $X_t\in\mathcal{X}$, our sampling rule $(A_t)$ is sequentially defined as
We offer the following lemma under this sampling rule. The proof is provided in Appendix (ref).
This lemma shows that the sampling rule can keep the allocation close to the optimal allocations. Thus, we can ensure that the sampling rule defined by ((ref)) (sampling rule) satisfies ((ref)) (allocation convergence).
To compute $w_{a,x}(t)$ in ((ref)), we need to solve the minimax problem defined in ((ref)) with the estimated parameters. If the number of contexts and arms is very large, it may be difficult to solve. However, except for such an extreme case, we can solve the problem by using minimax optimization based on the convex optimization algorithm in a short time. The computation is similar to that in jedra2020optimal.
We remark that the application of the original TS algorithm Garivier2016 for each context separately is not optimal for contextual BAI. Our problem setting makes finding the best allocations difficult, which is quite different from running BAI in parallel for each context. It is necessary to find good allocations of each arm to the right context, and the allocations among contexts are entangled. For example, to achieve our derived lower bound, one needs to think about saving the allocations to an arm $a$ in context $1$, then allocating more to arm $a$ in context $2$, and getting more budget to another arm $b$ in context $1$. In contrast, when separately applying the original TS algorithm, we cannot attain such an optimal allocation.
In this subsection, we present the stopping rule, in particular the threshold for the Bernoulli bandit model. We aim to design an algorithm that stops as early as possible while maintaining the failure probability less than or equal to $\delta$. We demonstrate that the stopping rule using the generalized likelihood ratio test (GLRT) for contextual BAI is $\delta$-PAC when the exploration ratio is properly tuned. Such a stopping rule is also known as Chernoff's stopping rule Chernoff1959. Although the approach for deriving the threshold is inspired by and similar to that of Garivier2016, our computation with the contextual information is more involved.
We consider a case where the reward $R_{t,a}$ follows a Bernoulli distribution conditioned on $X_t = x$. Here, the likelihood is given as
Then, for all pairs of the arms, $a,b\in [K]$, the GLRT statistic is given as
where $\overline{\xi}_a(t) = \sum_{x \in \mathcal{X}} \hat{\zeta}_x(t)\xi_{a,x}$. Note that the maximizer of
is equivalent to that of
We denote the maximizers by $(\widetilde{\xi}_{a,x}(t))$ and $(\widetilde{\xi}_{b,x}(t))$. Similarly, we denote the solution of the maximization problem in the denominator by $(\widetilde{\xi}^\dagger_{a,x}(t))$ and $(\widetilde{\xi}^\dagger_{b,x}(t))$.
In the numerator, if $\sum_{x\in\mathcal{X}}\hat{\zeta}_x(t)\hat{\mu}_{a,x}(t) \geq \sum_{x\in\mathcal{X}}\hat{\zeta}_x(t)\hat{\mu}_{b,x}(t)$, then the maximum likelihood estimator falls within the optimization constraint; that is, $\widetilde{\xi}_{a,x}(t) = \hat{\mu}_{a,x}(t)$ and $\widetilde{\xi}_{b,x}(t) = \hat{\mu}_{b,x}(t)$. Therefore, our remaining problem is to compute the denominator. Because $\sum_{x\in\mathcal{X}}\hat{\zeta}_x(t)\hat{\mu}_{a,x}(t) \geq \sum_{x\in\mathcal{X}}\hat{\zeta}_x(t)\hat{\mu}_{b,x}(t)$ does not satisfy the constraint condition in the denominator, it is hard to obtain the closed-form expression of the denominator and we need to solve the optimization problem numerically. Given the solutions, $(\widetilde{\xi}^\dagger_{a,x}(t))$ and $(\widetilde{\xi}^\dagger_{b,x}(t))$, the GLRT statistic $Z_{a,b}(t)$ is equal to
By multiplying $Z_{a,b}(t)$ by $-1/t$, we can find that solving the maximization problem is equal to solving the inner minimization problem of ((ref)), or equivalently the problem defined in ((ref)), by letting $\zeta_xw_{1,x}=\frac{N_{a,x}}{t}$, $\zeta_xw_{b,x}=\frac{N_{a,x}}{t}$, $\mu_{1,x} = \hat{\mu}_{a,x}(t)$, and $\mu_{a,x} = \hat{\mu}_{b,x}(t)$. From Lemma (ref), the constraint $\sum_{x\in\mathcal{X}}{\zeta}_x \xi_{a, x} \leq \sum_{x\in\mathcal{X}}{\zeta}_x \xi_{b, x}$ holds with equality; that is, \[Z_{a,b}(t) = -tL_{a,b}\left( \left(\hat{\mu}_{a,x}(t), \hat{\mu}_{b,x}(t), {\zeta}_x t, {N_{a,x}(t)}/N_{x}(t), {N_{b,x}(t)}/N_{x}(t) \right)_{x \in \mathcal{X}}\right).\] It is also easy to observe that when $\sum_{x\in\mathcal{X}}{\zeta}_x \hat{\mu}_{a,x}(t) \leq \sum_{x\in\mathcal{X}} {\zeta}_x \hat{\mu}_{b,x}(t)$, then $Z_{a,b}(t) = -Z_{b,a}(t)$.
Using the GLRT statistic, we use the following stopping rule:
where $\beta(t, \delta)$ is the threshold of the GLRT statistic $Z_{a,b}(t)$ (exploration rate), which controls the failure probability under the stopping rule.
Next, we determine $\beta(t, \delta)$ such that the proposed algorithm is $\delta$-PAC. We present the following theorem to decide the threshold $\beta(t, \delta)$ in the stopping rule.
The proof is provided in Appendix (ref). The proof with contextual information is accomplished by using the fact that joint distribution of the contexts and the rewards is the Multinomial distribution. This theorem confirms that the proposed algorithm is $\delta$-PAC when $\beta(t, \delta)=\log \left((2t (K-1))/\delta\right)v$. We note that this threshold does not depend on the cardinality of $\mathcal{X}$.
For the Bernoulli bandit, we derive the stopping and recommendation rule by using the fact that the rewards and finite contexts jointly follow a Multinominal distribution. We cannot use this property when the conditional rewards follow different distributions such as a Gaussian distribution. For example, when the rewards follow a Gaussian distribution, the rewards and contexts jointly follow a Gaussian mixture model, not a Gaussian distribution. This fact makes derivation of the $\delta$-PAC threshold difficult. However, if the contextual distribution is known, we can extend the existing results, such as Garivier2016 and Kaufmann2021, to derive the threshold.
We consider a case where for each $a \in [K]$ and $x \in \mathcal{X}$, the reward $R_{t,a}$ follows a distribution that belongs to the canonical one-parameter exponential family ((ref)) conditioned on $X_t = x$ and the context $X_{t}$ follows a multinomial distribution with known parameters; that is, we treat the estimator $\hat{\zeta}_x$ as the true value $\zeta_x$ in our proposed CTS algorithm. Similarly to the Bernoulli case, the likelihood of the observations $(\underline{R}_{a,x}(t))_{x \in \mathcal{X}}, \forall a \in [K]$ and $\underline{X}(t)$ regarding arm $a$ is given as follows.
Then, for all pairs of the arms, $a,b\in [K]$, the GLRT statistic is given as
where $\overline{\xi}_a(t) = \sum_{x \in \mathcal{X}} \zeta_x\xi_{a,x}$. For the numerator optimization problem, from the definition of the single parameter exponential family, the maximizer of
is equivalent to the maximizer of the optimization problem
As for the case of a Bernoulli bandit model, using the notation, we compute the GLRT statistic $Z_{a,b}(t)$ as follows. Now, suppose that $\sum_{x\in\mathcal{X}}\zeta_x\hat{\mu}_{a,x}(t) \geq \sum_{x\in\mathcal{X}} \zeta_x\hat{\mu}_{b,x}(t)$. Then, $\widetilde{\xi}_{a,x}(t) = \hat{\mu}_{a,x}(t)$ in the denominator and $\widetilde{\xi}_{b,x}(t) = \hat{\mu}_{b,x}(t)$. We numerically solve the optimization problem in the denominator and obtain the solutions, $(\widetilde{\xi}^\dagger_{a,x}(t))$ and $(\widetilde{\xi}^\dagger_{b,x}(t))$. Then, $Z_{a,b}(t)$ is equal to
A similar argument can be made when $\sum_{x\in\mathcal{X}}\zeta_x\hat{\mu}_{a,x}(t) < \sum_{x\in\mathcal{X}} \zeta_x\hat{\mu}_{b,x}(t)$ by reversing the sign of the constraint.
Next, we define the stopping rule using the GLRT statistic $Z_{a,b}(t)$ as follows. \[\tau_\delta = \inf\Big\{t\in\mathbb{N}: Z(t) := \max_{a\in[K]}\min_{b\in[K]\backslash\{a\}} Z_{a,b}(t) > \beta(t,\delta) \Big\},\] where we decide the threshold $\beta(t,\delta)$ later. Let $\hat{\mu}_c(t) = \sum_{x \in \mathcal{X}} \zeta_x \hat{\mu}_{c,x}(t)$ for $c\in\mathcal{A}$. If $\sum_{x \in \mathcal{X}} \zeta_x \mu_{a,x} = \mu_a \leq \mu_{b} = \sum_{x \in \mathcal{X}} \zeta_x \mu_{b,x}$ and $\hat{\mu}_a > \hat{\mu}_b$, then
Then, we decompose the probability $\mathbb{P}_{\nu}\left(\tau_\delta<\infty, \,\hat{a}_{\tau_\delta} \neq a^*\right) $ as
Thus, if we choose a threshold $\beta(t,\delta)$ such that the upper bound of the last equation ((ref)) is $\delta$, we can guarantee that the algorithm is $\delta$-PAC.
Using the results of Kaufmann2021, which refines existing deviation bounds and the threshold in Garivier2016, we can guarantee that the algorithm is $\delta$-PAC with a tight threshold. We use the following theorem from Kaufmann2021.
Let us define the threshold $\beta(t,\delta)$ as
Using this threshold, the following guarantee can be obtained.
In this section, we address the upper bound of the sample complexity of the proposed CTS algorithm.
First, we demonstrate that the sample complexity asymptotically matches the lower bound almost surely for a case where the reward follows a Bernoulli bandit model.
We provide the proof in Appendix (ref).
We now provide an upper bound on the expected number of the stopping times $\operatorname{\mathbb{E}}[\tau_\delta]$. The following theorem states that the proposed CTS algorithm asymptotically matches the sample complexity lower bound derived from Theorem (ref). The proof of this result is provided in Appendix (ref).
As well as the case with a Bernoulli bandit model, we can also show that an upper bound on the expected number of the stopping times $\operatorname{\mathbb{E}}[\tau_\delta]$ matches the lower bound almost surely for a case where the reward follows a distribution that belongs to a canonical one-parameter exponential family, and the parameters of the context distribution are known.
In this section, we investigate the behavior of the proposed algorithms. First, we examine the performance of $\alpha$-elimination using contextual information. As in Section (ref), we generate samples $\{(R_{t, 1}\ R_{t, 2}\ X_t)^\top\}^T_{t=1}$ from the multivariate distribution with the mean vector $(1 \ 0 \ 0)^\top$. We denote the variances of $R_{t, 1}$, $R_{t, 2}$, and $X_t$ as $\sigma^2_{1}$, $\sigma^2_2$, and $\sigma^2_{\mathcal{X}}$. Let the correlation coefficient between $R_{t,1}$ and $X_t$ be $\rho_{1 \mathcal{X}}$, and the correlation coefficient between $R_{t,2}$ and $X_t$ be $\rho_{2\mathcal{X}}$. We fix $\sigma^2_2 = 1$, $\sigma^2_{\mathcal{X}}=1$, and $\rho_{2\mathcal{X}}=0.5$. We investigate the performance of the proposed method by varying the combination of the variance $\sigma^2_1$ and correlation coefficient $\rho_{1\mathcal{X}}$. We choose $\sigma^2_1$ from $\{1,2\}$ and $\rho_{1\mathcal{X}}$ from $\{-0.9, -0.5, 0, 0.5, 0.9\}$. For the case with $\sigma^2_1=1$, the $\alpha$-elimination without contextual information of Kaufman2016complexity results in an allocation of $\alpha=0.5$ (uniform sampling). For the case with $\sigma^2_1=2$, it results in an allocation of $\alpha=\sqrt{2}/(\sqrt{2} + \sqrt{1})$. Conversely, the proposed $\alpha$-elimination with contextual information uses different allocations for each correlation coefficient. We conducted $1000$ trials with $\delta=0.05$ and display the realized stopping time (sample complexity) in Figure (ref) using box plots, where the right figure shows the results with $\sigma^2_1=1$ and the left shows the results with $\sigma^2_1=2$. In Figure (ref), we compare the proposed algorithm with different $\rho_{1\mathcal{X}}$ with the $\alpha$-elimination (without context). The results demonstrate that when using contextual information, the proposed $\alpha$-elimination can stop earlier than the original $\alpha$-elimination. We note the fact that the proposed algorithm can stop earlier, even though the allocation is also $0.5$ when $\rho_{1\mathcal{X}}$ is $0$. Here, the stopping threshold $\beta$ used in the proposed algorithm is less than that used in the original algorithm, while maintaining the $\delta$-PAC property. Note that for all cases, the realized $\delta$ does not exceed $0.05$.
Next, we compare the performance of the proposed CTS algorithm to the TS algorithm for BAI without contextual information Garivier2016. For a Bernoulli bandit model, we consider a sample scenario with marginalized mean rewards $\{\mu_1, \mu_2, \mu_3, \mu_4\} = \{0.3, 0.21, 0.2, 0.19\}$, which is the same as a scenario used in Garivier2016. Suppose that there exist two contexts $X_t \in\{1,2\}$, where the conditional mean rewards are given as $\{\mu_{1,1}, \mu_{2,1}, \mu_{3,1}, \mu_{4,1}\} = \{0.5, 0.01, 0.4, 0.01\}$ and $\{\mu_{1,2}, \mu_{2,2}, \mu_{3,2}, \mu_{4,2}\} = \{0.1, 0.41, 0., 0.37\}$. The context $1$ and $2$ appear with probability $0.5$, respectively. Because $\beta(t,\delta)$ can be determined by us within the range suggested in Theorems (ref)--(ref), and because the role of $\beta(t,\delta)$ does not change considerably between the CTS and TS algorithms, we display the value of the GLRT statistic in Figure (ref). The earlier this value becomes large, the smaller the sample complexity that can be achieved under a properly specified $\beta(t,\delta)$. This figure indicates that the CTS algorithm achieves a smaller sample complexity than TS, as suggested by the theoretical results. Conversely, the reason why the CTS algorithm indicates a smaller GLRT statistic compared with TS in the early rounds is likely because the number of parameters to be estimated is proportional to the number of contexts; thus it requires more time to converge in finite samples. In Appendix (ref), we present more details and additional results under different settings.
This paper proposed contextual BAI, where contextual information can be used to identify marginalized mean rewards. We noted that even contextual information that is not immediately related to the parameter that we wish to identify could help us solve the task more efficiently. We proposed the CTS algorithm as an algorithm when the rewards follow Bernoulli distributions, and confirmed that it performs better theoretically and experimentally when contextual information is provided. We also found that when the rewards and context follow a multivariate normal distribution in the two-armed bandit problem, we could improve the efficiency of BAI without changing the conventional algorithm. These properties have not been discussed to date. We consider that these results are related to semiparametric inference and the James--Stein shrinkage estimator; however, it is a future task to clarify their relationship
The authors thank Alexandre Proutière for detailed discussions.
\vskip 0.2in