A Simple and General Approach to top-$k$ Selection

Introduction

This paper studies the problem of a decision maker who faces a large pool of items with unknown quality, can assess them through noisy and costly samples, and wants to identify and select the best item(s). This is commonly referred to as a ranking and selection task in the operations literature [5, 4, 25], and as a best-arm(s) identification task in the multi-armed bandit literature [32, 34, 44]. This problem arises in many domains, including clinical trials [49, 10], A/B/n testing [46, 30], machine-learning model selection and hyperparameter optimization [48, 26, 36], and simulation optimization [17]. In each setting, the key question is how to adaptively allocate a fixed sampling budget to maximize the quality of the final selection.

A central goal of this paper is to treat the problem with generality so that the method is useful across a range of settings. Two aspects of this generality are worth discussing up front. The first concerns how samples are collected. The most adaptive option is to allocate one sample at a time and wait to observe the outcome before choosing the next. However, this can be impractical to implement. For example, in clinical trials, one may not be able to wait to observe an outcome before assigning the next treatment [43], and in simulation optimization, many simulations may be run in parallel [25]. We therefore consider the partially adaptive case, in which batches of samples are chosen in rounds, and outcomes are observed between rounds. The second aspect is the target number of items to select. Choosing the single best one ($k = 1$) is the most common and most studied case, but some of the applications above instead call for selecting several items. Our method can handle any target selection size $k$, giving it a wider range of applicability.

There are two common metrics for how “good” a sample allocation procedure is. One is its probability of correct selection (PCS), the probability that the selected items are the true best items [5, 21, 37]. The other is its expected opportunity cost (EOC), sometimes referred to as simple regret, which is the expected shortfall in value between the selected set and the true best set [11, 19]. PCS treats every incorrect selection equally, while EOC makes bigger mistakes more costly. For many tasks, EOC, which we adopt as our objective, is the more natural fit since a big mistake can be far worse than a marginal one. That said, the two are closely related. In simulations, a method that performs well on one metric tends to perform well on the other, and our theoretical guarantees apply to both objectives.

Given the objective, the decision maker’s problem can be formulated as a finite-horizon dynamic program. In each round, given the current state (the posterior distribution over item values) and the available budget, one chooses the optimal batch of samples by accounting for all possible outcomes and future states. However, the complexity of this program grows rapidly with the budget size, so solving it exactly is out of reach [15, 25, 51]. This has motivated a large body of literature, which we group into two broad categories. The first set of methods allocates samples by solving a sequence of static optimization problems, one per round, in place of the dynamic program. The second group constructs allocation rules directly from the current state without solving an optimization problem. We take the first approach and construct a per-round surrogate objective whose minimizer gives the sample allocation.

There are several ways to form a sequence of one-round static objectives. In a Bayesian formulation, one common approach is to maximize the expected value of information (EVI), the improvement in the quality of the final selection that collecting the next batch would produce [11, 9]. When samples follow a Gaussian distribution and are chosen one at a time, Frazier et al. [15, 16] show that this can be computed exactly for the top-1 EOC objective, yielding the knowledge gradient (KG) policy. In general, however, even the one-step look-ahead objective is intractable, so the literature resorts to approximations. Examples targeting PCS include AOAP for top-1 identification [40], and AOAm [51] and EMP [38] for top-$k$, all designed for Gaussian samples. For AOAm, Zhang et al. [51] also prove that the sampling proportions converge to the proportions that maximize the asymptotic error-decay rate, assuming their necessary sample-proportion conditions have a unique solution.

A second approach in this category, referred to as optimal computing budget allocation (OCBA) [4], assumes that in each round the current estimates of the means are the true means and finds the sample allocation that would be optimal under this assumption [5]. However, even with this assumption, the problem is generally intractable. It is therefore commonly approached by constructing a bound on the objective, a lower bound for PCS or an upper bound for EOC, and deriving allocations that optimize the bound as the budget goes to infinity. Most of this work concerns Gaussian samples and top-1 selection under PCS, with extensions to top-$k$ [3, 20], and to top-1 and top-$k$ selection under EOC [24, 19].

An alternative approach targets learning the allocation that maximizes the exponential rate at which the probability of incorrect selection decays as the budget grows, the large-deviation rate [21]. Algorithms of this kind run sequentially, again treating the current mean estimates as the truth, and choose allocations using the optimality conditions for this rate. Chen and Ryzhov [8] establish such an algorithm that applies to general distributions in the $k=1$ setting and prove that the sampling proportions converge to the optimal proportions as the budget approaches infinity.

Next, we review methods that construct sampling rules directly from the current state. Methods similar in spirit to our approach are confidence and gap-based frequentist approaches, which allocate samples based on how close items are to the current estimated top-$k$ cutoff and how uncertain their estimated values are, such as gap-e [18, 2] and KL-LUCB [35]. Our allocation algorithm is based on the same two quantities, but differs in how the top-$k$ cutoff is defined, how the distance from it is measured, and how uncertainty is used.

Elimination algorithms [1, 33] divide sampling into rounds, allocate samples uniformly among the remaining items within a round, and eliminate the empirically weakest items according to an elimination schedule. The successive accepts and rejects (SAR) approach [2] also “accepts” empirically strong items, avoiding spending samples on items already determined to likely be in the top-$k$.

Algorithms based on the expected improvement (EI) rule, which originates from the Bayesian optimization literature [28], have been designed for top-1 selection problems [45, 41, 7]. For each item, the algorithm computes the expected amount by which the item’s value exceeds that of the current best item, and samples the item with the largest improvement score. In the Gaussian case, Chen and Ryzhov [7] show that a variant of EI achieves the optimal sampling proportions asymptotically.

A family of Bayesian algorithms referred to as top two algorithms allocates a sample in each round by identifying a “leader” item and a “challenger” item, then flipping a weighted coin to decide which of the two to sample [44]. Various methods have been proposed for choosing the leader and challenger [29], for example, based on Thompson sampling [44] or expected improvement scores [41]. For $k=1$ and a broad class of sampling distributions, Russo [44] proves that with an unbiased coin, the posterior probability of an incorrect selection decays at an exponential rate of at least half the best achievable rate. Russo [44] further shows that with suitable tuning of the coin flip probability, it attains the optimal decay rate. Qin et al. [41] establish the corresponding results for top two expected improvement with Gaussian samples. You et al. [50] and Qin and You [42] extend top two methods to the top-$k$ setting and give a rigorous way to weight the coin such that in the top-1 Gaussian setting, the posterior probability of an incorrect selection decays at the optimal exponential rate without a tuning parameter, though their general top-$k$ sampling rules do not come with posterior error decay rate guarantees. Establishing optimality for general $k$ is more difficult because the conditions governing how samples must be allocated among items have a more complicated structure [51, 42]. Russo [44] conjectures that his top-$k$ extension of top two Thompson sampling achieves a posterior error decay rate of at least half the optimum.

Given the rich literature on this problem, our motivation for this work stems from a related project. In Goldberger and Rudi [23], we study the top-$k$ identification problem in a different setting with three key components that existing methods do not address in combination. First, the decision maker can choose among different binary sample forms (in particular, a binary rating of an item or a comparison of two items), each with a different cost and signal quality. Methods that allocate directly from the current state have no clear way to handle this, whereas methods that solve a static problem may, because the costs and signal quality of each sample form can be incorporated into the objective. Second, that paper considers “small” budget settings, so analysis in the asymptotic regime is not of primary importance. Third, samples may be collected in batches. Among the methods above, OCBA-style and elimination-based algorithms support batching [31]. However, OCBA algorithms generally rely on asymptotic analysis, and elimination-based algorithms are often outperformed in simulation by other methods such as gap-e and KL-LUCB [2, 35, 51]. This motivated us to develop a new method for that specific problem. In this paper, building on insights from that work, we develop a simpler, far more general method for Bayesian top-$k$ identification.

Our approach. We take the approach of solving a sequence of tractable static optimization problems, which we develop in two steps. First, we give an upper bound on EOC, structurally different from those in the OCBA literature on EOC, and discuss why this structure has the potential for an improved algorithm. Second, to address the intractability of minimizing this bound as a one-round static problem, we approximate the future posterior distribution after a batch via expected moment matching rather than relying on asymptotic analysis. This idea originates in the EVI and OCBA literature with Gaussian samples by Chen et al. [6], He et al. [24], and Zhang et al. [51], and we generalize it to other distributions. These two steps yield a simple, interpretable per-round static surrogate objective. A key property of the surrogate is that it is separable across items, enabling efficient algorithms for solving the batch problem without relying on heuristics.

Results and theoretical guarantees. Similar to the top two and modified expected-improvement methods for $k=1$, our method is conceptually simple and intuitive, yet remarkably comes with strong theoretical performance guarantees for general $k$. For a broad class of sampling distributions, we prove that posterior EOC (as well as the posterior probability of an incorrect selection) decays at an exponential rate of at least half the best achievable rate. This establishes the type of general-$k$ guarantee conjectured by Russo [44], but for our sampling rule.

In a range of simulations, we demonstrate that our method is competitive with and often improves on the current state-of-the-art approaches, including those specialized to $k=1$ as well as those for general $k>1$, those specialized to Gaussian as well as other distributions, and both at small and large budgets. These comparisons are conducted in the non-batched setting to include as many existing methods as possible. In the batched setting, we show that our method significantly outperforms existing batch algorithms and its performance degrades less as the batch size grows.

We believe this method can be valuable to both researchers and practitioners. It is flexible across a wide range of practical settings, is competitive with the best algorithms in each of these settings, and extends naturally to richer settings, such as those with varying sample forms and sampling costs.

The rest of the paper is organized as follows: § 2 formulates the model and objective, § 3 derives the surrogate objective and the batch allocation algorithm, § 4 establishes the theoretical guarantees, § 5 presents the numerical experiments, and § 6 concludes.

Model and objective

We consider $I$ items, indexed $i=1,\ldots,I$, each with an unknown mean $\theta_i \in\mathcal{J}$, where $\mathcal{J} \subseteq\mathbb{R}$ is an open interval. The goal is to select the $k<I$ items with the largest means. Each item $i$ has a prior distribution $\Pi_{i,0}$ on $\theta_i$ with density $\pi_{i,0}$. For item $i$, let $\mathcal{Y}_i$ be the outcome space, and let $\{P_i(\cdot\mid\theta):\theta\in\mathcal{J}\}$ be a known mean-parametrized family of distributions on $\mathcal{Y}_i$ with density or mass function $p_i(\cdot\mid\theta)$. Conditional on $\theta_i$, the $n$th observation $Y_{i,n}$ from item $i$ is drawn i.i.d. from $P_i(\cdot\mid\theta_i)$, with $\mathbb{E}[Y_{i,n}\mid\theta_i]=\theta_i$. We make the following regularity assumption.

Remark 1 gives examples of standard conjugate pairs (ESCP) satisfying Assumption 1 that we will refer to throughout the paper.

• Gamma with known shape: $Y_{i,n} \mid\theta_i \sim\operatorname{Gamma}(\kappa_i, \theta_i/\kappa_i)$ with scale $\theta_i/\kappa_i$ and $\kappa_i$ known ($\kappa_i = 1$ gives the exponential distribution), prior $\theta_i \sim\operatorname{InvGamma}(a_i(0), b_i(0))$ with $a_i(0) > 2$, $b_i(0) > 0$, $\mathcal{J} = (0, \infty)$.

• Lognormal with known log-variance: $Y_{i,n} \mid\theta_i \sim\operatorname{Lognormal}(\log\theta_i - \sigma_i^2/2, \sigma_i^2)$ with $\sigma_i^2$ known, prior $\log\theta_i \sim N(m_i(0), s_i^2(0))$, $\mathcal{J} = (0, \infty)$.

Sampling proceeds in $T+1$ rounds, an optional warm-up round indexed $t = 0$ followed by $T$ adaptive rounds.

In the warm-up round, each item is sampled $n_0 \in\mathbb{Z}_{\ge0}$ times. In rounds $t = 1, \ldots, T$, we choose a batch of $B_t \in\mathbb{Z}_+$ samples, where the per-round batch sizes $\{B_t\}_{t=1}^T$ are pre-specified. The choice of samples in round $t$ is denoted by the vector $\Delta^t = (\Delta_1^t, \ldots, \Delta_I^t)$, where $\Delta_i^t \in\mathbb{Z}_{\ge0}$ is the number of samples assigned to item $i$.

Let $N_i(t)$ denote the total samples allocated to item $i$ up to but not including round $t$, so that $N_i(0) = 0$, $N_i(1) = n_0$, and