SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Covariance adaptive best arm identification
El Mehdi Saad
Université Paris-Sacalay, CentraleSupélec
Laboratoire des Signaux et Systemes
el-mehdi.saad@centralesupelec.fr
Gilles Blanchard
Institut de Mathématique d'Orsay
gilles.blanchard@universite-paris-saclay.fr
Nicolas Verzelen
Mistea, INRAE
Institut Agro, Université de Montpellier
nicolas.verzelen@inrae.fr
We consider the problem of best arm identification in the multi-armed bandit model, under fixed confidence. Given a confidence input  $\delta$ , the goal is to identify the arm with the highest mean reward with a probability of at least  $1 - \delta$ , while minimizing the number of arm pulls. While the literature provides solutions to this problem under the assumption of independent arms distributions, we propose a more flexible scenario where arms can be dependent and rewards can be sampled simultaneously. This framework allows the learner to estimate the covariance among the arms distributions, enabling a more efficient identification of the best arm. The relaxed setting we propose is relevant in various applications, such as clinical trials, where similarities between patients or drugs suggest underlying correlations in the outcomes. We introduce new algorithms that adapt to the unknown covariance of the arms and demonstrate through theoretical guarantees that substantial improvement can be achieved over the standard setting. Additionally, we provide new lower bounds for the relaxed setting and present numerical simulations that support their theoretical findings.
1 Introduction and setting
Best arm identification (BAI) is a sequential learning and decision problem that refers to finding the arm with the largest mean (average reward) among a finite number of arms in a stochastic multi-armed bandit (MAB) setting. An MAB model  $\nu$  is a set of  $K$  distributions in  $\mathbb{R}$ :  $\nu_{1},\ldots ,\nu_{K}$ , with means  $\mu_1,\dots ,\mu_K$ . An "arm" is identified with the corresponding distribution index. The observation consists in sequential draws ("queries" or "arm pulls") from these distributions, and each such outcome is a "reward". The learner's goal is to identify the optimal arm  $i^{*}:= \operatorname {Arg}\operatorname *{Max}_{i\in [[K]]}\mu_{i}$ , efficiently. There are two main variants of BAI problems: The fixed budget setting [1, 8], where given a fixed number of queries  $T$ , the learner allocates queries to candidates arms and provides a
guess for the optimal arm. The theoretical guarantee in this case takes the form of an upper bound on the probability  $p_T$  of selecting a sub-optimal arm. The second variant is the fixed confidence setting [14, 20], where a confidence parameter  $\delta \in (0,1)$  is given as an input to the learner, and the objective is to output an arm  $\psi \in [[K]]$ , such that  $\mathbb{P}(\psi = i^*) \geq 1 - \delta$ , using the least number of arm pulls. The complexity in this case corresponds to the total number of queries made before the algorithm stops and gives a guess for the best arm that is valid with probability at least  $1 - \delta$  according to a specified stopping rule. In this paper we specifically focus on the fixed confidence setting.
The problem of best-arm identification with fixed confidence was extensively studied and is well understood in the literature [16, 19, 20, 14]. However, in these previous works, the problem was considered under the assumption that all observed rewards are independent. More precisely, in each round  $t$ , a fresh sample (reward vector), independent of the past,  $(X_{1,t},\ldots ,X_{K,t})$  is secretly drawn from  $\nu$  by the environment, and the learner is only allowed to choose one arm  $A_{t}$  out of  $K$  (and observe its reward  $X_{A_t,t}$ ). We relax this setting by allowing simultaneous queries. Specifically, we consider the MAB model  $\nu$  as a joint probability distribution of the  $K$  arms, and in each round  $t$  the learner chooses a subset  $C_t\subset [[K]]$  and observes the rewards  $(X_{i,t})_{i\in C_t}$  (see the Game Protocol 1). The high-level idea of our work is that allowing multiple queries per round opens up opportunities to estimate and leverage the underlying structure of the arms distribution, which would otherwise remain inaccessible with one-point feedback. This includes estimating the covariance between rewards at the same time point. It is important to note that our proposed algorithms do not require any prior knowledge of the covariance between arms.
Throughout this paper, we consider two cases: bounded rewards (in Section 4) and Gaussian rewards (in Section 5), with the following assumptions:
Assumption 1. Suppose that:
- IID assumption with respect to  $t$ :  $(X_{t})_{t\geq 1} = (X_{1,t},\ldots ,X_{K,t})_{t\geq 1}$  are independent and identically distributed variables following  $\nu$ .
- There is only one optimal arm:  $\left|\operatorname{Arg}\operatorname{Max}_{i\in [[K]]}\mu_i\right| = 1$ .
Assumption 2 -  $\mathcal{B}$ . Bounded rewards: The support of  $\nu$  is in  $[0,1]^K$ .
Assumption 2 -  $\mathcal{G}$ . Gaussian rewards:  $\nu$  is a multivariate normal distribution.
Protocol 1 The Game Protocol
Parameters:  $\delta$
while [condition] do
Choose a subset  $C \subseteq [[K]]$ .
The environment reveals the rewards  $(X_{i})_{i\in C}$
Output the selected arm:  $\psi$
A round corresponds to an iteration in Protocol 1. Denote by  $i^{*} \in [[K]]$  the optimal arm. The learner algorithm consists of: a sequence of queried subsets  $(C_t)_t$  of  $[[K]]$ , such that subset  $C_t$  at round  $t$  is chosen based on past observations, a halting condition to stop sampling (i.e. a stopping time written  $\tau$ ) and an arm  $\psi$  to output after stopping the sampling procedure. The theoretical guarantees take the form of a high probability upper-bound on the total number  $N$  of queries made through the game:
$$
N := \sum_ {t = 1} ^ {\tau} \left| C _ {t} \right|. \tag {1}
$$
2 Motivation and main contributions
The query complexity of best arm identification with independent rewards, when arms distributions are  $\sigma$ -sub-Gaussian is characterized by the following quantity:
$$
H (\nu) := \sum_ {i \in [ K ] \backslash \{i ^ {*} \}} \frac {\sigma^ {2}}{\left(\mu_ {i ^ {*}} - \mu_ {i}\right) ^ {2}}. \tag {2}
$$
Some  $\delta$ -PAC algorithms [16] guarantee a total number of queries satisfying  $\tau = \tilde{\mathcal{O}}(H(\nu) \log(1/\delta))$  (where  $\tilde{\mathcal{O}}(.)$  hides a logarithmic factor in the problem parameters).
Recall that the number of queries required for comparing the means of two variables, namely  $X_{1} \sim \mathcal{N}(\mu_{1}, 1)$  and  $X_{2} \sim \mathcal{N}(\mu_{2}, 1)$ , within a tolerance error  $\delta$ , is approximately  $\mathcal{O}\big(\log(1/\delta)/(\mu_{1}-\mu_{2})^{2}\big)$ . An alternative way to understand the sample complexity stated in equation (2) is to view the task of identifying the optimal arm  $i^{*}$  as synonymous with determining that the remaining arms  $\llbracket K\rrbracket \setminus \{i^{*}\}$  are suboptimal. Consequently, the cost of identifying the best arm corresponds to the sum of the comparison costs between each suboptimal arm and the optimal one.
In practical settings, the arms distributions are not independent. This lack of independence can arise in various scenarios, such as clinical trials, where the effects of drugs on patients with similar traits or comparing drugs with similar components may exhibit underlying correlations. These correlations provide additional information that can potentially expedite the decision-making process [18].
In such cases, Protocol 1 allows the player to estimate the means and the covariances of arms. This additional information naturally raises the following question:
Can we accelerate best arm identification by leveraging the (unknown) covariance between the arms?
We give two arguments for a positive answer to the last question: When allowed simultaneous queries (more than one arm per round as in Protocol 1), the learner can adapt to the covariance between variables. To illustrate, consider the following toy example for 2 arms comparison:  $X_{1} \sim \mathcal{N}(\mu_{1}, 1)$  and  $X_{2} = X_{1} + Y$  where  $Y \sim \epsilon \mathcal{N}(1, 1)$  for some  $\epsilon > 0$ . BAI algorithms in one query per round framework require  $\mathcal{O}((1 + \epsilon)^{2} \log(1/\delta)/\epsilon^{2})$  queries, which gives  $\mathcal{O}(\log(1/\delta)/\epsilon^{2})$  for small  $\epsilon$ . In contrast, when two queries per round are possible, the learner can perform the following test  $\mathcal{H}_{0}: "\mathbb{E}[X_{1} - X_{2}] > 0"$  against  $\mathcal{H}_{1}: "\mathbb{E}[X_{1} - X_{2}] \leq 0"$ . Therefore using standard test algorithms adaptive to unknown variances, such as Student's  $t$ -test, leads to a number of queries of the order  $\mathcal{O}(\mathrm{Var}(X_{1} - X_{2}) \log(1/\delta)/\epsilon^{2}) = \mathcal{O}(\log(1/\delta))$ . Hence a substantial improvement in the sample complexity can be achieved when leveraging the covariance. We can also go one step further to even reduce the sample complexity in some settings. Indeed, we can establish the sub-optimality of an arm  $i$  by comparing it with another sub-optimal arm  $j$  faster than when comparing it to the optimal arm  $i^{*}$ . To illustrate consider the following toy example with  $K = 3$ : let  $X_{1} \sim \mathcal{N}(\mu_{1}, 1)$ ,  $X_{2} = X_{1} + Y$  where  $Y \sim \epsilon \mathcal{N}(1, 1)$  and  $X_{3} \sim \mathcal{N}(\mu_{1} + 2\epsilon, 1)$ , with  $X_{1}, Y$  and  $X_{3}$  independent.  $X_{3}$  is clearly the optimal arm. Eliminating  $X_{1}$  with a comparison with  $X_{3}$  requires  $\mathcal{O}(\log(1/\delta)/\epsilon^{2})$  queries, while comparing  $X_{1}$  with the sub-optimal arm  $X_{2}$  requires  $\mathcal{O}(\mathrm{Var}(X_{1} - X_{2}) \log(1/\delta)/\epsilon^{2}) = \mathcal{O}(\log(1/\delta))$ .
The aforementioned arguments can be adapted to accommodate bounded variables, such as when comparing correlated Bernoulli variables. These arguments suggest that by utilizing multiple queries per round (as shown in Protocol 1), we can expedite the identification of the best arm compared to the standard one-query-per-setting approach. Specifically, for variables bounded by 1, our algorithm ensures that the cost of eliminating a sub-optimal arm  $i \in [[K]]$  is given by:
$$
\min  _ {j \in [ K ]: \mu_ {j} > \mu_ {i}} \left\{\frac {\operatorname {V a r} \left(X _ {j} - X _ {i}\right)}{\left(\mu_ {j} - \mu_ {i}\right) ^ {2}} + \frac {1}{\mu_ {j} - \mu_ {i}} \right\}. \tag {3}
$$
The additional term  $1 / (\mu_j - \mu_i)$  arises due to the sub-exponential tail behavior of the sum of bounded variables. To compare the quantity in (3) with its counterpart in the independent case, it is important to note that the minimum is taken over a set that includes the best arm, and that  $\operatorname{Var}(X_{i^*} - X_i) \leq 2(\operatorname{Var}(X_{i^*}) + \operatorname{Var}(X_i))$ . Consequently, when the variables are bounded, the quantity (3) is no larger than a numerical constant times  $H(\nu)$  (its independent case counterpart in 2), with potentially a significant improvements if there is positive correlation between an arm  $j \in [[K]]$  with a higher mean and arm  $i$  (so that  $\operatorname{Var}(X_j - X_i)$  is small). As a result, our algorithm for bounded variables has a sample complexity, up to a logarithmic factor, that corresponds to the sum of quantities (3) over all sub-optimal arms.
We expand our analysis to encompass the scenario where arms follow a Gaussian distribution. In this context, our procedure ensures that the cost of eliminating a sub-optimal arm  $i \in [[K]]$  is given by:
$$
\min_{\substack{j\in \llbracket K\rrbracket :\\ \mu_{j} > \mu_{i}}}\Bigg\{\frac{\operatorname{Var}(X_{j} - X_{i})}{(\mu_{j} - \mu_{i})^{2}}\lor 1\Bigg\}.
$$
Similar to the setting with bounded variables, the aforementioned quantity is always smaller than its counterpart for independent arms when all variables  $X_{j}$  have a unit variance.
In Section 6, we present two lower bound results for the sample complexity of best arm identification in our multiple query setting, specifically when the arms are correlated. Notably, these lower bounds pinpoint how the optimal sample complexity may decrease with the covariance between arms, making them the first of their kind in the context of best arm identification, to the best of our knowledge. The presented lower bounds are sharp, up to a logarithmic factor, in the case where all arms are positively correlated with the optimal arm. However, it remains an open question to determine a sharper lower bound applicable to a more general class of distributions with arbitrary covariance matrices.
In Section B of the appendix, we introduce a new algorithm that differs from the previous ones by performing comparisons between the candidate arm and a convex combination of the remaining arms, rather than pairwise comparisons. We provide theoretical guarantees for the resulting algorithm.
Lastly, we conduct numerical experiments using synthetic data to assess the practical relevance of our approach.
Best arm identification: BAI in the fixed confidence setting was studied by [11], [22], and [12], where the objective is to find  $\epsilon$ -optimal arms under the PAC ("probably approximately correct") model. A summary of various optimal bounds for this problem is presented in [8, 20]. Prior works on fixed confidence BAI [25] and [17] developed strategies adaptive to the unknown variances of the arms. In contrast, our proposed algorithms demonstrate adaptability to all entries of the covariance matrix of the arms. In particular, we establish that in the worst-case scenario, where the arms are independent, our guarantees align with the guarantees provided by these previous approaches.
Covariance in the Multi-Armed Bandits model: Recently, the concept of leveraging arm dependencies in the multi-armed bandit (MAB) model for best-arm identification was explored in [15]. However, their framework heavily relies on prior knowledge, specifically upper bounds on the conditional expectation of rewards from unobserved arms given the chosen arm's reward. In a similar vein, a game protocol that allows simultaneous queries was examined in [21]. However, their objective differs from ours as their focus is on identifying the most correlated arms, whereas our primary goal is to identify the arm with the highest mean reward.
The extension of the standard multi-armed bandit setting to multiple-point bandit feedback was also considered in the stochastic combinatorial semi-bandit problem (2, 9, 10 and 13). At each round  $t \geq 1$ , the learner pulls  $m$  out of  $K$  arms and receives the sum of the pulled arms rewards. The objective is to minimize the cumulative regret with respect to the best choice of arms subset. [27] proposed an algorithm that adapts to the covariance of the covariates within the same arm. While this line of research shares the intuition of exploiting the covariance structure with our paper, there are essential differences between the two settings. In the combinatorial semi-bandit problem, the learner receives the sum of rewards from all selected arms in each round and aims to minimize cumulative regret, necessitating careful exploration during the game. In contrast, our approach does not impose any constraint on the number of queried arms per round, and our focus is purely on exploration.
Simultaneous queries of multiple arms was also considered in the context of graph-based bandit problems [7]. However, in these studies, it is assumed that the distributions of arms are independent.
Model selection racing: Racing algorithms for model selection refers to the problem of selecting the best model out of a finite set efficiently. The main idea consists of early elimination of poorly performing models and concentrating the selection effort on good models. This idea was seemingly first exploited in [23] through Hoeffding Racing. It consists of sequentially constructing a confidence interval for the generalization error of each (non-eliminated) model. Once two intervals become disjoint, the corresponding sub-optimal model is discarded. Later [25] presented an adaptive stopping algorithm using confidence regions derived with empirical Bernstein concentration inequality (3, 24). The resulting algorithm is adaptive to the unknown marginal variances of the models. Similarly, [4] presented a procedure centered around sequential hypothesis testing to make decisions between two possibilities. In their setup, they assume independent samples and use Bernstein concentration inequality tailored for bounded variables to adapt to the variances of the two variables under consideration.
While the idea of exploiting the possible dependence between models was shown [6, 26] to empirically outperform methods based on individual performance monitoring, there is an apparent lack of theoretical guarantees. This work aims to develop a control on the number of sufficient queries for
reliable best arm identification, while being adaptive to the unknown correlation of the candidate arms.
4 Algorithm and main theorem for bounded variables
In this section, we focus on the scenario where the arms are bounded by 1. Let us establish the notation we will be using throughout. For each arm  $i \in [[K]]$  associated with the reward variable  $X_{i}$ , we denote by  $\mu_{i}$  the mean of  $X_{i}$ . Additionally, for any pair of arms  $i, j \in [[K]]$ , we denote  $V_{ij}$  as the variance of the difference between  $X_{i}$  and  $X_{j}$ , i.e.,  $V_{ij} := \operatorname{Var}(X_{i} - X_{j})$ . We remind the reader that the quantities  $V_{ij}$  are unknown to the learner.
Algorithm 2 is developed based on the ideas introduced in Section 2, which involves conducting sequential tests between every pair  $(i,j) \in [[K]] \times [[K]]$  of arms. The key element to adapt to the covariance between arms is the utilization of the empirical Bernstein's inequality [24] for the sequence of differences  $(X_{i,t} - X_{j,t})_t$  for  $i,j \in [[K]]$ . To that end, we introduce the following quantity:
$$
\hat {\Delta} _ {i j} (t, \delta) := \hat {\mu} _ {j, t} - \hat {\mu} _ {i, t} - \frac {3}{2} \alpha (t, \delta) \sqrt {2 \hat {V} _ {i j , t}} - 9 \alpha^ {2} (t, \delta),
$$
where,  $\hat{\mu}_{i,t}$  represents the empirical mean of the samples obtained from arm  $i$  up to round  $t$ , and  $\widehat{V}_{ij,t}$  denotes the empirical variance associated with the difference variable  $(X_{i} - X_{j})$  up to round  $t$ . The term  $\alpha(t,\delta)$  is defined as  $\alpha(t,\delta) \coloneqq \sqrt{\log(1 / \delta_t) / (t - 1)}$ , and  $\delta_t$  is given by  $\delta / (2K^2 t(t + 1))$ .
By leveraging the empirical Bernstein's inequality (restated in Theorem K.1), we can establish that if  $\hat{\Delta}_{ij}(t,\delta) > 0$  at any time  $t$ , then with a probability of at least  $1 - \delta$ , we have  $\mu_i < \mu_j$ . This observation indicates that arm  $i$  is sub-optimal. Furthermore, when  $\mu_j > \mu_i$  for  $i,j \in [[K]]$ , a sufficient sample sizes ensuring that the quantity  $\hat{\Delta}_{ij}(t,\delta)$  is positive is proportional to:
$$
\Lambda_ {i j} := \left(\frac {V _ {i j}}{(\mu_ {i} - \mu_ {j}) ^ {2}} + \frac {1}{\mu_ {j} - \mu_ {i}}\right) \log \left(\frac {K}{(\mu_ {j} - \mu_ {i}) \delta}\right). \tag {4}
$$
Moreover, we show in Lemma E.5 in the appendix that the last quantity is necessary, up to a smaller numerical constant factor, in order to have  $\hat{\Delta}_{ij}(t,\delta) > 0$ .
Algorithm 2 follows a successive elimination approach based on the tests  $\hat{\Delta}_{ij}(t,\delta) > 0$ . Our objective is to ensure that any sub-optimal arm  $i$  is queried at most  $\min_j\Lambda_{ij}$  times, where the minimum is taken over arms with means larger than  $\mu_i$ . However, this approach poses a challenge: the arm  $j^*$  that achieves the minimum of  $\Lambda_{ij}$  may be eliminated early in the process, and the algorithm is then constrained to compare arm  $i$  with other arms costing larger complexity  $\Lambda_{ij}$ . To address this limitation, it is useful to continue querying arms even after deciding their sub-optimality through the pairwise tests.
In Algorithm 2, we introduce two sets at round  $t$ :  $S_{t}$ , which contains arms that are candidates to be optimal, initialized as  $S_{1} = [[K]]$ , and  $C_t$ , which represents the set of arms queried at round  $t$ . Naturally,  $C_t$  contains  $S_{t}$  and also includes arms that were freshly eliminated from  $S_{t}$ , as we hope that these arms will help in further eliminating candidate arms from  $S_{t}$  more quickly.
An important consideration is how long the algorithm should continue sampling an arm that has been eliminated from  $S_{t}$ . In Theorem E.2, we prove that when arm  $j$  is eliminated at round  $t$ , it is sufficient to keep sampling it up to round  $82t$  (the constant 82 is discussed in Remark 1). The rationale behind this number of additional queries is explained in the sketch of the proof of Theorem 4.1. It suggests that if arm  $j$  fails to eliminate another arm after round  $82t$ , then the arm that eliminated  $j$  from  $S_{t}$  can ensure faster eliminations for the remaining arms in  $S_{t}$ .
Theorem 4.1. Suppose Assumption 1 and  $2 - \mathcal{B}$  holds. Consider Algorithm 2, with input  $\delta \in (0,1)$ . We have with probability at least  $1 - \delta$ : the algorithm identifies the best arm and the total number of queries  $N$  satisfies:
$$
N \leq c \log \left(K \Lambda \delta^ {- 1}\right) \sum_ {i \in [   [ K ]   ] \backslash \{i ^ {*} \}} \min  _ {\substack {j \in [   [ K ]   ]; \\ \mu_ {j} > \mu_ {i}}} \left\{\frac {V _ {i j}}{\left(\mu_ {i} - \mu_ {j}\right) ^ {2}} + \frac {1}{\mu_ {j} - \mu_ {i}} \right\}, \tag{5}
$$
where  $\Lambda = \max_{i\in [[K]]}\min_{\substack{j\in [[K]]:\\ \mu_j > \mu_i}}\left\{\frac{V_{ij}}{(\mu_i - \mu_j)^2} +\frac{1}{\mu_j - \mu_i}\right\}$  and  $c$  is a numerical constant.
Algorithm 2 Pairwise-BAI
1: Input  $\delta$    
2: Initialization:   
3: Query all arms for 2 rounds and compute the empirical means vector  $\hat{\mu}_t,t\gets 3$    
4:  $S\gets [[K]],\quad /\ast \mathrm{Set}$  of candidate arms\*/   
5:  $C\gets [[K]],\quad /\ast \mathrm{Set}$  of queried arms\*/   
6: while  $|S| > 1$  do   
7: Jointly query all arms in  $C$    
8: Update  $\hat{\mu}_t$  and compute  $\max_{j\in C}\hat{\Delta}_{ij}(t,\delta)$  for each  $i\in S$    
9: for  $i\in S$  do   
10: if  $\max_{j\in C}\hat{\Delta}_{ij}(t,\delta) > 0$  then   
11: Eliminate  $i$  from  $S:S\gets S\setminus \{i\}$  /\* i is sub-optimal \*/   
12: Mark  $i$  for elimination from  $C$  at round: 82 t.   
13: end if   
14: end for   
15: Increment t.   
16: end while   
17: Return  $S$
Moreover, if we omit line 12 from Algorithm 2, that is we do not query non-candidate arms, we have with probability at least  $1 - \delta$ : the algorithm identifies the best arm and the total number of queries  $N$  satisfies:
$$
N \leq c \log \left(K \Lambda \delta^ {- 1}\right) \sum_ {i \in [   [ K ]   ] \backslash \{i ^ {*} \}} \left\{\frac {V _ {i i ^ {*}}}{\left(\mu_ {i ^ {*}} - \mu_ {i}\right) ^ {2}} + \frac {1}{\mu_ {i ^ {*}} - \mu_ {i}} \right\}, \tag {6}
$$
where  $\Lambda$  is defined above and  $c$  is a numerical constant.