Binary Search with Distributional Predictions

In this paper we initiate the study of algorithms with distributional predictions, focusing on the basic search in a sorted array problem described above. In addition to the classic $O(\log n)$ comparisons binary search for an array of size $n$, we recall the “median” or “bisection” algorithm (first described by Knuth [19] and analyzed by Mehlhorn [27]) which probes the cell representing the median of the distribution, and recurses appropriately. When the target keys are indeed drawn from the given distribution, the expected query complexity (i.e. number of comparisons between elements in the array and the target) is bounded by $H(p)+1$, where $H(p)$ is the entropy of the distribution [27]. (Note that when the distribution is uniform over $n$ elements, this recovers the $O(\log n)$ binary search bound). This is in fact essentially optimal: it is known that every algorithm requires at least $H(p)/3$ queries in expectation when target keys are drawn from $p$ [27].

On the other hand, it is easy to see that if target keys are not drawn from $p$, then this algorithm can be arbitrarily bad: it can easily be made to use $\Omega(n)$ comparisons in expectation. So our main question is the following: how can we best utilize a prediction $\hat{p}$ which is not the true distribution $p$? Can we recover the near-optimality of the median algorithm without being subject to its worst-case performance?

Machine learning augmented algorithms have found applications in various areas — for example, online algorithms [24, 30], combinatorial algorithms [13, 11], differential privacy [3], data structures [22, 32, 25, 26] and mechanism design [1], to name a few. In particular, online algorithms have been extensively studied with ML advice for various problems, such as online caching [24], ski-rental [30], scheduling [21], knapsack [16], set cover [7], and more [23]. Due to the vast literature, we only provide a few samples.

Particularly relevant to our setting is the work of Lin et al. [22], which initiated the study of predictions for binary search trees. This work investigates how to improve a treap’s guarantees when item frequencies follow distributions such as the Zipfian distribution.

As discussed earlier, in ML augmented algorithms, predictions are typically given in the form of specific values, rather than distributions. Here, we discuss a few exceptions. The work by Diakonikolas et al. [12] studies the ski-rental problem and prophet inequalities with access to i.i.d. samples from an unknown distribution. Their focus lies on the sample complexity and not on the correctness of the distribution. Indeed, they obtain robustness by combining their consistent algorithm and the best worst-case algorithm in a black-box manner. In contrast, we assume full access to a distributional prediction and develop new ideas to obtain a tight trade-off between consistency and robustness, in conjunction with natural error measures involving entropy and the earth mover’s distance, which take the accuracy of the distributional prediction into account.

More broadly the question of how to improve performance of problems where either full instances, or some model parameters come from a known distribution is well studied under the rubric of two-stage stochastic optimization [31, 2], with techniques like Sample Average Approximation (SAA) [18] having a rich history. Similar to our setting, one can look at the robust setting, where the distribution available to the algorithm is different from the true distribution that examples are drawn from, see e.g., [8, 15, 9, 10]. Typically in these cases one assumes a bound on the difference between the two distributions, explicitly choosing what to hedge against, and then derives optimal strategies. In contrast, in this work we aim to find a smooth trade-off on the performance of the algorithm as a function of the distance between the two distributions, coupled with an upper bound on worst-case performance.

Finally, the very recent work of Angelopoulos et al. [5] is one of the few that consider distributional predictions. It shows an optimal tradeoff between consistency and robustness for a scheduling problem. However, their solution space explored is considerably more limited than ours, essentially consisting of geometric sequences with a multiplicative ratio of 2, each characterized by its starting point. Further, their bound analysis is restricted to cases where the error is sufficiently small. In contrast, our work demonstrates how a binary search algorithm can compare generally to earth mover’s distance (EMD) and a lower bound on the optimum for any predicted distribution. As a result, we develop novel algorithmic solutions that build upon a close connection to EMD.

Given prediction $\hat{p}$ of $p$, suppose the algorithm first computes some point $\hat{\alpha}$ from $\hat{p}$ and then uses the doubling binary search algorithm from Mitzenmacher and Vassilvitskii [28] with prediction $\hat{\alpha}$. This will have expected running time of $O(\log(|\hat{\alpha}-\alpha|))$, where $\alpha$ is the true location of the key. A natural question is whether there is some $\hat{\alpha}$ so that $O(\log(|\hat{\alpha}-\alpha|))$ is comparable to $O(H(p)+\log\eta)$.

Unfortunately, this is not possible, necessitating our more involved algorithm and analysis. Let $p$ be the distribution on two atoms, with $p_{n/4}=1/2$ and $p_{3n/4}=1/2$, and let $\hat{p}=p$. Clearly $\eta=0$, and $H(p)=1$, so any competitive algorithm must terminate after a constant number of comparisons in expectation. On the other hand, consider some $\hat{\alpha}\in[n]$. If $\hat{\alpha}\leq n/2$, then since $\alpha=3n/4$ with probability $1/2$ we have that $\mathbb{E}[\log(|\hat{\alpha}-\alpha|)]\geq\frac{1}{2}\log(n/4))=\Omega(\log n)$. Similarly, if $\hat{\alpha}\geq n/2$, then since $\alpha=n/4$ with probability $1/2$ we have that $\mathbb{E}[\log(|\hat{\alpha}-\alpha|)]\geq\Omega(\log n)$. Hence converting $\hat{p}$ to a point prediction and then using the algorithm of [28] as a black box is doomed to failure.

The goal is to show the following theorem with respect to the algorithm.

Before formally proving the theorem, we give key intuition about the analysis. In iteration $k$, the Bisection phase is continued for $2^{k}$ steps. In each step, one comparison is made, which makes the cost of this phase $2^{k}$.

Consider the Binary Search at the Endpoints phase. Two comparisons are made during the phase unless $a\in[\ell,\ell+d]$ or $a\in[r-d,r]$. In those cases, we run a traditional binary search on an interval of length $d+1$, whose cost is $\log d=\log 2^{2^{k}}=2^{k}$.

For each key $a_{i}$, it takes at most $\log(\frac{1}{p_{i}})+1$ iterations of the Bisection phase to get to a search range that has a predicted probability mass of at most $p_{i}/2$. We can charge the total cost of the iterations up to this point to the term $p_{i}\log(\frac{1}{p_{i}})$ in $H(p)$.

Either we find $a_{i}$ earlier, in which case the total cost of the iterations can be charged to $H(p)$, or there is an at least $p_{i}/2$ mass that was predicted to lie outside the interval, allowing us to lower bound $\eta$. We make this argument formal below.

It is well known that there is a lower bound of $\Omega(H(p))$ on the expected query complexity for binary search [27]. We now show that there is a lower bound of $\Omega(\log\eta)$ on the expected query complexity as well, even on instances with $H(p)=0$. This shows that there must fundamentally be a $\log\eta$ dependence on the earth mover’s distance, even for instances where it is not absorbed by the dependence on the entropy.

Combining Theorem 2 with the $\Omega(H(p))$ lower bound due to [27] results in the following worst-case lower bound, asymptotically matching Theorem 1.

We proceed in a similar manner, alternating the two phases. However, we change the algorithm so that in the first phase the algorithm performs a binary search on the medians of each distribution. The goal of this is to ensure that each distribution has its probability mass drop by at least half in each step.

As before, the initial search range is $[1,n]$. The algorithm proceeds in iterations. For $i=0,1,\ldots$, iteration $i$ has two phases.

• •

  Bisection. Let $[\ell,r]$ be the current search range. Let $S_{k}=\sum_{j=\ell}^{r}\hat{p}_{j,k}$ be the remaining probability mass in the $k$’th prediction.

  Let $t_{k}$ be such that $\sum_{j=\ell}^{t_{k}-1}\hat{p}_{j,k}\leq\frac{1}{2}S_{k}$ and $\sum_{j=t_{k}+1}^{r}\hat{p}_{j,k}\leq\frac{1}{2}S_{k}$. That is, $t_{k}$ is the median of the $k$’th distribution. Sort the indices $k\in[m]$ so that $t_{1}\leq t_{2}\ldots\leq t_{m}$. For convenience, let $t_{0}=\ell$ and $t_{m+1}=r$. Perform a binary search on $a_{t_{0}},a_{t_{1}},a_{t_{2}},\ldots a_{t_{m+1}}$ to find the interval where $a\in(a_{t_{j}},a_{t_{j+1}})$ for some $j\in\{0,1,2,\ldots m\}$. The new search range is $[t_{j}+1,t_{j+1}-1]$.

  Continue this for $2^{i}$ steps, and if $a$ is not found, begin the second phase described below.

• •

  Binary Search at the Endpoints. Let $[\ell,r]$ be the current search range. Set $d=\min(2^{2^{i}},r-\ell)$. Check if $a$ is in the range $[\ell,\ell+d]$ or $[r-d,r]$ (by comparing $a$ to $a_{\ell+d}$ and $a_{r-d}$). If $a$ is in one of these ranges, say $[\ell,\ell+d]$, do a regular binary search (by choosing the middle point of the range each time) on the range $[\ell,\ell+d]$, until $a$ is found. Otherwise, start the next iteration with the new search range $[\ell+d+1,r-d-1]$.

The algorithm continues until $a$ is found.

We now state the following theorem regarding the algorithm for multiple predictions. The overhead of using multiple predictions is a $\log m$ factor. The proof is very similar to the proof of Theorem 1 and has been deferred to Appendix A.

We begin with experiments on synthetic data where we can vary the prediction error in a controlled environment to show the algorithms sensitivity and robustness to mispredictions.

In this setting, let the keyspace be the integers in $[-10^{5},10^{5}]$. We then generate $t=10^{4}$ independent points from a normal distribution with mean $0$ and standard deviation $10$, rounding down each to the nearest integer. This results in a concentrated distribution in a very large key space. The $t$ points form the predicted distribution, $\hat{p}$. To generate the test distribution, we proceed in the same manner, but shift the mean of the normal distribution away from $0$ by some value $s>0$. Note that for $s=100$ the train and test distributions have 0 overlap with high probability. For each value of $s$, we repeat the experiment $5$ times and report the average and standard deviation of the costs.

Our results for this setting can be found in Fig. 1, where we plot the average search cost (query complexity) of each algorithm against the shift amount for the test distribution. At one extreme, where there is no shift in the test distribution, we observe that all three algorithms which utilize the predicted distribution perform well. Since the bisection algorithm is optimal when the error is 0, it performs the best, as expected, while the Learned BST approach exhibits some overhead due to hedging against possible errors. However, a perturbation to the predicted distribution causes the bisection algorithm to perform worse than classical binary search. Both the convex combination and learned BST algorithms demonstrate a smoother degradation in performance, with our proposed method (learned BST) giving more robust performance to even large shifts in the test distribution. When the erorr becomes very high, then the additional overhead of the learned BST algorithm makes it slightly worse than the Classic baseline.