Fixed-Budget Best-Arm Identification with Heterogeneous Reward Variances

The problem of best-arm identification (BAI) in the fixed-budget setting is a pure exploration bandit problem which can be briefly described as follows. An agent interacts with a stochastic multi-armed bandit with $K$ arms and its goal is to identify the arm with the highest mean reward within a fixed budget $n$ of arm pulls [Bubeck et al., 2009, Audibert et al., 2010]. This problem arises naturally in many applications in practice, such as online advertising, recommender systems, and vaccine tests [Lattimore and Szepesvari, 2019]. It is also common in applications where observations are costly, such as Bayesian optimization [Krause et al., 2008]. Another commonly studied setting is fixed-confidence BAI [Even-Dar et al., 2006, Soare et al., 2014]. Here the goal is to identify the best arm within a prescribed confidence level while minimizing the budget. Some works also studied both settings [Gabillon et al., 2012, Karnin et al., 2013, Kaufmann et al., 2016].

Our work can be motivated by the following example. Consider an A/B test where the goal is to identify a movie with the highest average user rating from a set of $K$ movies. This problem can be formulated as BAI by treating the movies as arms and user ratings as stochastic rewards. Some movies get either unanimously good or bad ratings, and thus their ratings have a low variance. Others get a wide range of ratings, because they are rated highly by their target audience and poorly by others; and hence their ratings have a high variance. For this setting, we can design better BAI policies that take the variance into account. Specifically, movies with low-variance ratings can be exposed to fewer users in the A/B test than movies with high-variance ratings.

An analogous synthetic example is presented in Figure 1. In this example, reward variances increase with mean arm rewards for a half of the arms, while the remaining arms have very low variances. The knowledge of the reward variances can be obviously used to reduce the number of pulls of arms with low-variance rewards. However, in practice, the reward variances are rarely known in advance, such as in our motivating A/B testing example, and this makes the design and analysis of variance-adaptive BAI algorithms challenging. We revisit these two examples in our empirical studies in Section 5.

We propose and analyze two variance-adaptive BAI algorithms: $\tt SHVar$ and $\tt SHAdaVar$. $\tt SHVar$ assumes that the reward variances are known and is a stepping stone for our fully-adaptive BAI algorithm $\tt SHAdaVar$, which estimates them. $\tt SHAdaVar$ utilizes high-probability upper confidence bounds on the reward variances. Both algorithms are motivated by sequential halving ($\tt SH$) of Karnin et al. [2013], a near-optimal solution for fixed-budget BAI with homogeneous reward variances.

Our main contributions are:

• •

  We design two variance-adaptive algorithms for fixed-budget BAI: $\tt SHVar$ for known reward variances and $\tt SHAdaVar$ for unknown reward variances. $\tt SHAdaVar$ is only a third algorithm for this setting [Gabillon et al., 2011, Faella et al., 2020] and only a second that can be implemented as analyzed [Faella et al., 2020]. The key idea in $\tt SHAdaVar$ is to solve a budget allocation problem with unknown reward variances by a greedy algorithm that overestimates them. This idea can be applied to other elimination algorithms in the cumulative regret setting [Auer and Ortner, 2010] and is of independent interest to the field of optimal experiment design [Pukelsheim, 1993].

• •

  We prove upper bounds on the probability of misidentifying the best arm for both $\tt SHVar$ and $\tt SHAdaVar$. The analysis of $\tt SHVar$ extends that of Karnin et al. [2013] to heterogeneous variances. The analysis of $\tt SHAdaVar$ relies on a novel lower bound on the number of pulls of an arm that scales linearly with its unknown reward variance. This permits an analysis of sequential halving without requiring a closed form for the number of pulls of each arm.

• •

  We evaluate our methods empirically on Gaussian bandits and the MovieLens dataset [Lam and Herlocker, 2016]. In most settings, $\tt SHVar$ and $\tt SHAdaVar$ outperform all prior algorithms.

The paper is organized as follows. In Section 2, we present the fixed-budget BAI problem. We present our algorithms in Section 3 and analyze them in Section 4. The algorithms are empirically evaluated in Section 5. We review prior works in Section 6 and conclude in Section 7.

We use the following notation. Random variables are capitalized, except for Greek letters like $\mu$. For any positive integer $n$, we define $[n]=\left\{1,\dots,n\right\}$. The indicator function is denoted by $\mathds{1}\!\left\{\cdot\right\}$. The $i$-th entry of vector $v$ is $v_{i}$. If the vector is already indexed, such as $v_{j}$, we write $v_{j,i}$. The big O notation up to logarithmic factors is $\tilde{O}$.

We have a stochastic bandit with $K$ arms and denote the set of arms by $\mathcal{A}=[K]$. When the arm is pulled, its reward is drawn i.i.d. from its reward distribution. The reward distribution of arm $i\in\mathcal{A}$ is sub-Gaussian with mean $\mu_{i}$ and variance proxy $\sigma^{2}_{i}$. The best arm is the arm with the highest mean reward,

Without loss of generality, we make an assumption that the arms are ordered as $\mu_{1}>\mu_{2}\geq\ldots\geq\mu_{K}$. Therefore, arm $i_{*}=1$ is a unique best arm. The agent has a budget of $n$ observations and the goal is to identify $i_{*}$ as accurately as possible after pulling all arms $n$ times. Specifically, let $\hat{I}$ denote the arm returned by the agent after $n$ pulls. Then our objective is to minimize the probability of misidentifying the best arm $\mathbb{P}\left(\hat{I}\neq i_{*}\right)$, which we also call a mistake probability. This setting is known as fixed-budget BAI [Bubeck et al., 2009, Audibert et al., 2010]. When observations are costly, it is natural to limit them by a fixed budget $n$.

Another commonly studied setting is fixed-confidence BAI [Even-Dar et al., 2006, Soare et al., 2014]. Here the agent is given an upper bound on the mistake probability $\delta$ as an input and the goal is to attain $\mathbb{P}\left(\hat{I}\neq i_{*}\right)\leq\delta$ at minimum budget $n$. Some works also studied both the fixed-budget and fixed-confidence settings [Gabillon et al., 2012, Karnin et al., 2013, Kaufmann et al., 2016].

A near-optimal solution for fixed-budget BAI with homogeneous reward variances is sequential halving [Karnin et al., 2013]. The key idea is to sequentially eliminate suboptimal arms in $\log_{2}K$ stages. In each stage, all arms are pulled equally and the worst half of the arms are eliminated at the end of the stage. At the end of the last stage, only one arm $\hat{I}$ remains and that arm is the estimated best arm.

The main algorithmic contribution of our work is that we generalize sequential halving of Karnin et al. [2013] to heterogeneous reward variances. All of our algorithms can be viewed as instances of a meta-algorithm (Algorithm 1), which we describe in detail next. Its inputs are a budget $n$ on the number of observations and base algorithm $\tt Alg$. The meta-algorithm has $m$ stages (line 2) and the budget is divided equally across the stages, with a per-stage budget $n_{s}=\left\lfloor n/m\right\rfloor$ (line 5). In stage $s$, all remaining arms $\mathcal{A}_{s}$ are pulled according to $\tt Alg$ (lines 6–8). At the end of stage $s$, the worst half of the remaining arms, as measured by their estimated mean rewards, is eliminated (lines 9–12). Here $Y_{s,t,i}$ is the stochastic reward of arm $i$ in round $t$ of stage $s$, $I_{s,t}\in\mathcal{A}_{s}$ is the pulled arm in round $t$ of stage $s$, $N_{s,i}$ is the number of pulls of arm $i$ in stage $s$, and $\hat{\mu}_{s,i}$ is its mean reward estimate from all observations in stage $s$.

The sequential halving of Karnin et al. [2013] is an instance of Algorithm 1 for $\tt Alg=\tt SH$. The pseudocode of $\tt SH$, which pulls all arms in stage $s$ equally, is in Algorithm 2. We call the resulting algorithm $\tt SH$. This algorithm misidentifies the best arm with probability [Karnin et al., 2013]

where

is a complexity parameter and $\Delta_{i}=\mu_{1}-\mu_{i}$ is the suboptimality gap of arm $i$. The bound in (1) decreases as budget $n$ increases and problem complexity $H_{2}$ decreases.

$\tt SH$ is near optimal only in the setting of homogeneous reward variances. In this work, we study the general setting where the reward variances of arms vary, potentially as extremely as in our motivating example in Figure 1. In this example, $\tt SH$ would face arms with both low and high variances in each stage. A variance-adaptive $\tt SH$ could adapt its budget allocation in each stage to the reward variances and thus eliminate suboptimal arms more effectively.

This section comprises three analyses. In Section 4.1, we bound the probability that $\tt SHVar$, an algorithm that knows reward variances, misidentifies the best arm. In Section 4.2, we provide an alternative analysis that does not rely on the closed form in (4). Finally, in Section 4.3, we bound the probability that $\tt SHAdaVar$, an algorithm that learns reward variances, misidentifies the best arm.

All analyses are under the assumption of Gaussian reward noise. Specifically, the reward of arm $i$ in round $t$ of stage $s$ is distributed as $Y_{s,t,i}\sim\mathcal{N}(\mu_{i},\sigma_{i}^{2})$.

In this section, we empirically evaluate our proposed algorithms, $\tt SHVar$ and $\tt SHAdaVar$, and compare them to algorithms from prior works. We choose the following baselines: uniform allocation ($\tt Unif$), sequential halving ($\tt SH$) [Karnin et al., 2013], gap-based exploration ($\tt GapE$) [Gabillon et al., 2011], gap-based exploration with variance ($\tt GapE\mathchar 45\relax V$) [Gabillon et al., 2011], and variance-based rejects ($\tt VBR$) [Faella et al., 2020].

$\tt Unif$ allocates equal budget to all arms and $\tt SH$ was originally proposed for homogeneous reward variances. Neither $\tt Unif$ nor $\tt SH$ can adapt to heterogenuous reward variances. $\tt GapE$, $\tt GapE\mathchar 45\relax V$ and $\tt VBR$ are variance-adaptive BAI methods from related works (Section 6). In $\tt GapE$, we use $H$ from Theorem 1 of Gabillon et al. [2011]. In $\tt GapE\mathchar 45\relax V$, we use $H$ from Theorem 2 of Gabillon et al. [2011]. Both $\tt GapE$ and $\tt GapE\mathchar 45\relax V$ assume bounded reward distributions with support $[0,b]$. We choose $b=\max_{i\in\mathcal{A}}\mu_{i}+\sigma_{i}\sqrt{\log n}$, since this is a high-probability upper bound on the absolute value of $n$ independent observations from $\mathcal{N}(\mu_{i},\sigma_{i}^{2})$. In $\tt SHAdaVar$, we set $\delta=0.05$, and thus our upper bounds on reward variances hold with probability $0.95$. In $\tt VBR$, $\gamma=1.96$, which means that the mean arm rewards lie between their upper and lower bounds with probability $0.95$. Faella et al. [2020] showed that $\tt VBR$ performs well with Gaussian noise when $\gamma\approx 2$. All reported results are averaged over $5\,000$ runs.

$\tt GapE$ and $\tt GapE\mathchar 45\relax V$ have $O(\exp[-cn/H])$ error bounds on the probability of misidentifying the best arm, where $n$ is the budget, $H$ is the complexity parameter, and $c=1/144$ for $\tt GapE$ and $c=1/512$ for $\tt GapE\mathchar 45\relax V$. Our error bounds are $O(\exp[-c^{\prime}n/H^{\prime}])$, where $H^{\prime}$ is a comparable complexity parameter and $c^{\prime}=1/(4\log_{2}K)$. Even for moderate $K$, $c\ll c^{\prime}$. Therefore, when $\tt SHVar$ and $\tt SHAdaVar$ are implemented as analyzed, they provide stronger guarantees on identifying the best arm than $\tt GapE$ and $\tt GapE\mathchar 45\relax V$. To make the algorithms comparable, we set $H$ of $\tt GapE$ and $\tt GapE\mathchar 45\relax V$ to $Hc/c^{\prime}$, by increasing their confidence widths. Since $H$ is an input to both $\tt GapE$ and $\tt GapE\mathchar 45\relax V$, note that they have an advantage over our algorithms that do not require it.

Best-arm identification has been studied extensively in both fixed-budget [Bubeck et al., 2009, Audibert et al., 2010] and fixed-confidence [Even-Dar et al., 2006] settings. The two closest prior works are Gabillon et al. [2011] and Faella et al. [2020], both of which studied fixed-budget BAI with heterogeneous reward variances. All other works on BAI with heterogeneous reward variances are in the fixed-confidence setting [Lu et al., 2021, Zhou and Tian, 2022, Jourdan et al., 2022].

The first work on variance-adaptive BAI was in the fixed-budget setting [Gabillon et al., 2011]. This paper proposed algorithm $\tt GapE\mathchar 45\relax V$ and showed that its probability of mistake decreases exponentially with budget $n$. Our error bounds are comparable to Gabillon et al. [2011]. The main shortcoming of the analyses in Gabillon et al. [2011] is that they assume that the complexity parameter is known and used by $\tt GapE\mathchar 45\relax V$. Since the complexity parameter depends on unknown gaps and reward variances, it is typically unknown in practice. To address this issue, Gabillon et al. [2011] introduced an adaptive variant of $\tt GapE\mathchar 45\relax V$, $\tt A\mathchar 45\relax GapE\mathchar 45\relax V$, where the complexity parameter is estimated. This algorithm does not come with any guarantee.

The only other work that studied variance-adaptive fixed-budget BAI is Faella et al. [2020]. This paper proposed and analyzed a variant of successive rejects algorithm [Audibert et al., 2010]. Since $\tt SH$ of Karnin et al. [2013] has a comparable error bound to successive rejects of Audibert et al. [2010], our variance-adaptive sequential halving algorithms have comparable error bounds to variance-adaptive successive rejects of Faella et al. [2020]. Roughly speaking, all bounds can be stated as $\exp[-n/H]$, where $H$ is a complexity parameter that depends on the number of arms $K$, their variances, and their gaps.

We propose variance-adaptive sequential halving for fixed-budget BAI. Our algorithms have state-of-the-art performance in our experiments (Section 5). They are conceptually simpler than prior works [Gabillon et al., 2011, Faella et al., 2020] and can be implemented as analyzed, unlike Gabillon et al. [2011].

We study best-arm identification in the fixed-budget setting where the reward variances vary across the arms. We propose two variance-adaptive elimination algorithms for this problem: $\tt SHVar$ for known reward variances and $\tt SHAdaVar$ for unknown reward variances. Both algorithms proceed in stages and pull arms with higher reward variances more often than those with lower variances. While the design and analysis of $\tt SHVar$ are of interest, they are a stepping stone for $\tt SHAdaVar$, which adapts to unknown reward variances. The novelty in $\tt SHAdaVar$ is in solving an optimal design problem with unknown observation variances. Its analysis relies on a novel lower bound on the number of arm pulls in BAI that does not require closed-form solutions to the budget allocation problem. Our numerical simulations show that $\tt SHVar$ and $\tt SHAdaVar$ are not only theoretically sound, but also competitive with state-of-the-art baselines.

Our work leaves open several questions of interest. First, the design of $\tt SHAdaVar$ is for Gaussian reward noise. The reason for this choice is that our initial experiments showed quick concentration and also robustness to noise misspecification. Concentration of general random variables with unknown variances can be analyzed using empirical Bernstein bounds [Maurer and Pontil, 2009]. This approach was taken by Gabillon et al. [2011] and could also be applied in our setting. For now, to address the issue of Gaussian noise, we experiment with non-Gaussian noise in Section 5.2. Second, while our error bounds depend on all parameters of interest as expected, we do not provide a matching lower bound. When the reward variances are known, we believe that a lower bound can be proved by building on the work of Carpentier and Locatelli [2016]. Finally, our algorithms are not contextual, which limits their application because many bandit problems are contextual [Li et al., 2010, Wen et al., 2015, Zong et al., 2016].