Binary Search with Distributional Predictions
Authors: Michael Dinitz
Sungjin Im
Thomas Lavastida
Benjamin Moseley
Aidin Niaparast
Sergei Vassilvitskii
Venue: NeurIPS 2024 poster (2024)
Keywords: learning-augmented algorithms
algorithms with predictions
distribution predictions
Primary area: optimization
TL;DR: This paper extends the algorithms with predictions model to the case where predictions are distributions for the problem of searching in a sorted array.

Abstract: Algorithms with (machine-learned) predictions is a powerful framework for combining traditional worst-case algorithms with modern machine learning.  However, the vast majority of work in this space assumes that the prediction itself is non-probabilistic, even if it is generated by some stochastic process (such as a machine learning system).  This is a poor fit for modern ML, particularly modern neural networks, which naturally generate a *distribution*.  We initiate the study of algorithms with *distributional* predictions, where the prediction itself is a distribution.  We focus on one of the simplest yet fundamental settings: binary search (or searching a sorted array).  
 This setting has one of the simplest algorithms with a point prediction, but what happens if the prediction is a distribution?  We show that this is a richer setting: there are simple distributions where using the classical prediction-based algorithm with any single prediction does poorly.  
 Motivated by this, as our main result, we give an algorithm with query complexity
 $O(H(p) + \log \eta)$, where $H(p)$ is the entropy of the true distribution $p$ and $\eta$ is the earth mover's distance between $p$ and the predicted distribution $\hat p$.  This also yields the first *distributionally-robust* algorithm for the classical problem of computing an optimal binary search tree given a distribution over target keys.
 We complement this with a lower bound showing that this query complexity is essentially optimal (up to constants), and experiments validating the practical usefulness of our algorithm.

Binary Search with Distributional Predictions
Michael Dinitz* Johns Hopkins University mdinitz@cs. juu.
Sungjin Im UC Merced s amm@ucmerced. edu
Thomas Lavastida
University of Texas at Dallas
juu. edu
thomas lat.vastid@@uddallllel edu
Benjamin Moseleyf Carnegie Mellon University moseleyb@andrev, cmu. edu
Aidin Niaparast Carnegie Mellon University ani aniapara@andrev. cmu. edu
Sergei Vassilvitskii
Google Research
sergeivdgoogle.cm
Abstract
Algorithms with (machine-learmed predictions is powerful framework for com- bining traditional worst-case algorithms with modern machine learning. However, the vast majority of work in this space assumes that the prediction itself is non- probabilistic, even if it is generated by some stochastic process (such as machine learning system). This is a poor fit for modern ML, particularly modern neural net- works, which naturally generate distribution. We initiate the study of algorithms with distributional predictions, where the prediction itself distribution. We focus on one of the simplest yet fundamental settings: binary search (or searching a sorted array). This setting has one the simplest algorithms with a point prediction, but what happens if the prediction is a distribution? We show that this is richer setting: there are simple distributions where using the classical prediction- based algorithm with any single prediction does poorly. Motivated by this, as main result, we give an algorithm with query complexity (H) + log 17). where (p) is the entropy of the true distribution p and 1) is the earth mover's 's distance between p and the predicted distribution p. This also yields the first distributionally- optimal binary search tree
algorithm for the classical problem of computing given a distribution over target keys. We complement this with lower bound showing that this query complexity essentially optimal (up constants), and experiments validating the practical usefulness of our algorithm.
Introduction
Algorithms with predictions, or algorithms with machine-learned advice, has proved to be a useful framework for combining machine learning (which is extremely useful in the usual case but can be quite bad when the worst case) with traditional worst-case algorithms (which are quite good in the worst-case but do not do well have appeared in many places in the past, the formal study of this setting was pioneered by Lykouris) particularly motivated by the practical success of learned index
might hope in the usual case). While similar ideas
and Vitsilvitskii 2021], and structures Kraska et al. 2018). The goal usually to design an algorithm for some classical and also provided with some type of "advice" "prediction"
important problem in the setting when (presumably given by some machine learning algorithm) as to what the instance is like. If the advice "good" then want our algorithm do extremely well, while if the advice is "bad" then
"Supported in part by NSF awards 1909111 and 2228995.
1 'Supported in part NSF awards 1844939, 2121745, and 2423106, and by ONR grant N00014-22-15 1-20011 "Supported in part by Google Research Award, an Infor Research Award, a Carnegie Bosch Junior Faculty Chair, NSF grants CCF-2121744 and CCF- CCF-1845146 and ONR Grant N000142212702.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

want to inherit the traditional worst-case guarantee. In other words, we want the best of both worlds: good performance in the average case thanks to machine learning, but also robustness and good performance in the worst-case from traditional algorithms.
Consider the basic problem of searching for key in sorted array. This problem is the first example considered in the survey of Mitzemmacrre and Vassilvitskii [2021], as it is perhaps one of the simplest yet also best-motivated settings. Given sorted array with n elements and target key i, we can of course do a binary search for i using only Olog n) comparisons to find its location a(i). But suppose that we additionally receive a prediction d(i) E [rr] of the location of the target key i in the array, possibly from some machine learning system which attempts to predict the correct location for each key. If the prediction is perfect (a(i) = d(i)) then it is easy to use this prediction: only one comparison is needed! On the other hand, if the prediction is meaningless, then we can run classical binary search. But what if the prediction is "closee" turns out that "doubling binary search" from the predicted point can be used to design an algorithm which makes at most ((log(áá(i) - a)) comparisons [Mizzenmacher and Vassilvitskii 202T]. So if the prediction is very close to the true location then we make very few queries, while if is very far away then we recover the traditional binary search comparison bound.
The ability to obtain these types of results has led to an explosion of interest in algorithms with predictions; see Section 1.2 for some references. Sometimes the prediction itself is simple, as in the search problem, while sometimes it is quite complex, for instance encompassing a multi-dimeensonall vector. However, with only a few exceptions such as Diakonikolas et al. (202T) Angelopoulos et al. [024], which will be discussed in Section 1.2 in detail, all of these papers share an important feature: the prediction itself is nn-pprbbabilistic. That is, the prediction
single (potentially
high-dimensional) point (or maybe a small number of such points). While making the setting simpler, this is not a good match with the actual output of most ML systems (particularly modern neural networks), which inherently output a distribution. The question we study in this work is how to take advantage of the full richness of the prediction. Of course, we can always turn a distribution into single point in any number of ways (using a max likelihood estimator (MLE), sampling from the distribution, etc.). But is that always the right thing to do? Or, can we in fact do better by taking full advantage of the entire predicted distribution?
1.1 Our Results and Contributions
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 Ollog n.) comparisons binary search for array of size n, we recall the "median" or "bisection" algorithm (first described by Knuth
1971] and analyzed by |Mehhhrrn 1975)) 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 11 (p) 1, where (p) is the entropy of the distribution [Mchlhom 1975]. (Note that when the distribution is uniform over n. elements, this recovers the Ollog n.) binary search bound). This is in fact essentially optimal: is known that every algorithm requires least H (p)/3 3 queries in expectation when target keys are drawn from [Mehlhorn 1975].
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 22(n) comparisons in expectation. So our main question is the following: how can we best utilize prediction 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?
Reduction to point distributions. We first show in Section 2.1 that the obvious approach, of reducing P to a point prediction (whether by sampling, using a max- likeliikeldood prediction, or some other method) and then using previous algorithms, performance. In addition to ruling out a natural class of algorithms, this gives additional motivation to our study of distributional predictions: as discussed, essentially all previous work studies the case in which the prediction is single point (in this case, location); yet, most machine-learning systems will actually output a distribution. So our lower bound implies that any of these traditional point-based algorithms, no matter how good the bound obtained compared to their prediction, must
bad idea that can lead to poor worst-case
2

---

suffer fundamentally poor performance in the real world where target keys actually come from a distribution.
Main algorithm. We then give our main result in Section phases of the "median" algorithm and classical binary search to obtain a query complexity of (p) + log where di the earth mover's distance (EMD, also known as the Wasserstein Wi metric) between p and p. See Section 3 for precise details. Note that (p) S (loogn) and log 1) S log n. So if our prediction is close to p. then our algorithm has performance essentially equal to the best possible bound H (p). On the other hand, if our prediction ý is far from (so provides essentially no information), we do not suffer the poor performance of naively believing in P and running the median algorithm on it, instead recovering a bound of O [ (p) + log 17) = Ollog n).
algorithm which interleaves
While there are many notions of "distance" between distributions, EMD is natural one in our setting. Many other notions of distance, like {,, do not take the geometry of the line into account. For example, consider some distribution p over [n.] and let p' be the distribution obtained from p by moving 7/2 probability mass from 1 to 2, and let p" be the distribution obtained from p by moving 7/2 probability mass from to n. Then the & distance between p and p' is 7, and so is the & distance between p and y". Yet clearly p' is a "more accurate" prediction for p. The earth mover's distance recognizes this fact, and is more appropriate measure than f1. Similarly, popular measures such as KL-divergence (which are not technically metrics, but do give notion of distance) suffer the same flaws as f while also being extraordinarily sensitive to mismatches in the support (the KL-divergence can be infinite if the supports of the two distributions do not agrec).
Distributional Robustness of Optimal Binary Search Trees. We have so far discussed a distri- butional prediction setting where target key arrives with predicted distribution of its location. This is a strict generalization of Mitzenmacher and Vassilvitskii (2021). where the prediction is a single location in the array. In their model, the location error prone and in ours the distribution over locations is error prone. Our goal is to construct an effective search strategy given the target and the prediction.
But there is another related setting: there is a single (unknown) distribution over target keys, and are given (possibly erroneous) prediction of this distribution and are asked to design search algorithm with minimum expected lookup time when target keys are drawn from the true distribution. In other words, instead of each target key coming with predicted distribution ý over locations in [n] and the true location being drawn from some true distribution p. we are given ahead of time a predicted distribution ý over kargel keys [n], and are asked design lookup algorithm for inputs that use this distribution. But then these target keys in the input are actually drawn from p rather than p. and do not come with target-specific predictions.
Since any comparison- based search algorithm is equivalent to a particular binary search tree, if ů then this is precisely the classical problem of computing optimal binary search tree [Mchlhorn] 1975]. So we can interpret our results as providing distributionally efficiently compute BST where the expected lookup time under the true (but unknown) query distribution is at most O (p) + log r)). Surprisingly, given the classical nature of computing optimal (or near-optimal) BSTs, this simple question of "what if my distribution is incorrect?" has not been considered in the data structures and algorithms literature.
optimal BSTs: given p. we
Worst case lower bound. We complement our algorithmic development with a lower bound in Section 3.2 proving that no algorithm can use fewer than 2(log 1)) queries in the worst case. Since (p) is a known lower bound as well even if p is known perfectly [Mehlhorn] 1975). J. this implics that our algorithm is asymptotically tight. So if we measure accuracy of the prediction via EMD, no algorithm can make asymptotically better use of a distributional prediction.
Portfolios of predictions. There has been recent interest in the study of algorithm with multiple predictions, sometimes called prediction portfolios. See, for example, [Balcan et al.) 2021 Dinitz et al. 2022 Anand et al., 2022 Kevi and Nguyen| 2023). The goal is usually to do as well as the best of the predictions in the portfolio, with the difficulty being that we do not know priori which of these predictions is best. We extend our main algorithm to this setting in Section 4 showing that possible to use multiple distributional predictions effectively.
3

---

Experiments. Finally, in Section 5 we give empirical evidence of the efficacy of the algorithm we propose. We first use synthetic data to demonstrate the effect of distribution error on the performance of the algorithm. We then evaluate it on a number of real world datasets.
1.2 Other Related Work
Machine learning augmented algorithms have found applications in various areas - for example,
online algorithms [Lykouris and Vassilvitskii 2021 Purohit et al. [2018]. combinatorial algorithms [Dinitz al.
2021 Davies et al.]]0033. differential privacy [Amin et al. 2022]. data structures [Lin)
et al.l.| 2022 Vaidya et al. 2021 McCauley et al. 2024, Mccauley et al. [2024] and mechanism design [Agrawal et al. 2022]. to name a few. In particular, online algorithms have been extensively studied with ML advice for various problems, such as online caching [Lykouris and Vassilvitskii, 202T], ski-rental (Purohit et al. 2018), scheduling [Lattanzi et al., 2020]. knapsack [Im et al., 2021], set cover (Bamas et al. 22020]. J. and more [Lindermayr and Megow 2022]. Due to the vast literature,
only provide a few samples.
Particularly relevant to our setting is the work of Lin et al. predictions for binary search trees. This work investigates how to improve treap's guarantees when item frequencies follow distributions such as the Zipfian distribution.
2022 which initiated the study of
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. [202I] studies the ski-rental problem and prophet inequalities with access to Li.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 distributional prediction and develop new ideas 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 known distribution is well studied under the rubric of two-stage stochastic optimization [Swamy and Shmoys] 2006)] Ahmed 2010). with techniques like Sample Average Approximation (SAA) [Kim et al.) 2015 having 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., [Bertsimas and Goyal 2010 Dütting and
Kesselheim 2019 Bertsimas et al. 2022 Besbes et al. 2022]. 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 al. 2024 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 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 close connection to EMD.
2 Preliminaries
We now formally describe our problem and setting. Let 01 < ... < On be a set of n keys, and p = (Pi..... Pn)) be a probability distribution over the keys. Our goal is develop a search strategy (or search algorithm), which takes a target key as input, and finds a position i such that a; = We aim to find search strategies with low expected search cost when the target key a is sampled according to the distribution p. In our analysis we will consider the number of comparisons, also known as the query complexity, as the main metric of study. This metric captures the information theoretic complexity of the problem, and ignores computational overhead. Formally, let the search

---

cost C(as) of finding be the number of comparisons done by the algorithm when the target key is a, The expected search cost is then Σi= (aC)
To aid in this goal, we are given a prediction, f (fi,,.... .r......PPP...... of p. To account for the fact that the predicted distribution may be incorrect, we let 1j denote the earth mover's 's distance (EMD) between p and p. The EMD between two distributions / and Q is the solution to the optimal transport problem between them, or more formally, it is inf inf=-1((P))) F(z.g)>>/(( y)]. where [I(P, Q) is the set of joint distributions with marginals P and Q.
Finally, for distribution p. we let H (p) = - Pr log(ps) be the entropy of p. Throughout the paper, all logarithms are in base 2.
Given the breadth of work on point predictions, it is tempting to try and reduce the distributional prediction problem to point predictions. We show that this approach does not lead to good results.
2.1 Point Predictions from Distributions
Given prediction of p. suppose the algorithm first computes some point á from á and then uses the doubling binary search algorithm from Mitzenmacher and Vassilvitskii [202]] with prediction á. This will have expected running time of Olog([āá al)), where is the true location of the key. A natural question is whether there is some & so that O(log(Iá al)) is comparable to (H + log n).
Unfortunately, this is not possible, necessitating our more involved algorithm and analysis. Let be the distribution on two atoms, with Pn/1 = 1/2 and P3n/1 = 1/2, and let ý = p. Clearly 1/ = so any competitive algorithm must terminate after constant number of comparisons in expectation. On the other hand, consider some à E [n]. If å S n/2,, then since = 3n/4 with probability 1/2 we have that Ellog(iá al] 2 1 (log(n/4) = 2(log n). Similarly, if á 2 n/2. then since a = n2/4 with probability 1/2 we have that E [olo[(á l al)] ]] 2 (llog Hence converting ý to a point prediction and then using the algorithm Mitzenmcheerr and Vassilvitskii [022]] as black box is doomed to failure.
and H
2
3 Algorithm
To develop robust approach, recall the two baseline algorithms traditional binary search with an Ollog n) running time and the algorithm that recurses on the median element of the distribution, with ((H(p) (p)) running time (assuming it has access to the true distribution p).
In our algorithm we interleave these two approaches to get the best of both worlds. Let E {aj.... an) be the target key. We proceed recursively, keeping track of an active search range [e, r] (if = a, we always have E ()))). Initially, we start with & = and n. The algorithm proceeds in iterations. Each iteration for i = 0, 0, has two phases
Bisection. Divide the search range in half based on the predicted probabilities p. Formally, find index k, < S such that Σ}= P5 1 and Σj Pj J where S - Pj. Compare a to ag. If they are equal, return Otherwise, based on the result of the comparison, continue the search on the ranges [e, k: - 1] or [K: + 1, Continue this process for 2' steps, and if not found, begin the second phase.
Binary Search at the Endpoints. Let [] be the current search range. Set d = min(22" r f). Check if is in the range [e, l + d) or [r d,r] (by comparing to ag+d and 4)- If a is in one of these ranges, say [e, + d]. do a regular binary search (by choosing the middle point of the range each time) the range [e, + d]. until is found. Otherwise, start the next iteration with the new search range 1L + d + 1, d 2-11 1].
The algorithm continues until a is found.
3.1 Analysis
The goal show the following theorem with respect to the algorithm.
5

---

Theorem 3.1. The expected query complexity of the described algorithm is at most (H + 8 max(log(n) + +, 1) + 8 ( max(lo))
Before formally proving the theorem, we give key intuition about the analysis. In iteration k. the Bisection phase is continued for 2 steps. In each step, one comparison is made, which makes the cost of this phase 2*.
Consider the Binary Search at the Endpoints phase. Two comparisons made during the phase unless a E [e, & d] or a E [r d, r|. those cases, we run a traditional binary search on an interval of length d + 1, whose cost is log d = log 2 = =
For each key a, takes at most log( 1.) iterations of the Bisection phase get to a search range that has a predicted probability mass of at most p./2. We can charge the total cost of the iterations up to this point to the term p: log( }.) in (p).
Either we find ag earlier, in which case the total cost of the iterations be charged to Hp)), or there is an at least p2/2 mass that was predicted to lie outside the interval, allowing us to lower bound We make this argument formal below.
With probability pt, the target key that we are looking for is 0z. The goal to bound the expected cost of the algorithm, which is Σn= paC() where C(a;) is the cost of the algorithm when a = aq. Let k; be the first iteration at which a is found, assuming = ag As mentioned earlier, the total cost of the first phase of the iterations O to Ki is Σラ 2 < 2k++1 Also, the cost of the second phase in each iteration before k; is 2, and in iteration k is at most 24, So the total cost of the algorithm for iterations to k; is at most
Proof of Theorem 3.7
24. + 2k, S 4. 2*. We partition the keys based on k; into two sets, and bound the cost of each set separately. Let 11 = {i: h I (log(())//))) and 12 == {i: K2 > log(log( (4p:)))
First, we bound the cost of indices in I by constant factor of H(p):
Σ C.CC(a,) HEN
< Σ p
Σ p. log (lgg(//p))==)/// S 4H()) 8. veh
Now we bound the cost of the indices in I2 by constant factor of log(1)) Let i E Iz. We know that during iteration j of searching for a;, in the Bisection phase, the predicted probability mass in the search range decreases by a factor of at least 2² Therefore the predicted probability mass in the search range [e, r] at the end of the first phase in iteration k; 1 is at most
2
2 Pi
<
R 2
22
2
4/p:
2
where the inequality holds because E 12. So Ej S p./ Let D; min(i f, i). In the transportation problem corresponding to the earth mover's distance between p and p. probability mass of at least p:/2 2 needs to be moved from point to the outside of the interval [f, r). The cost of this movement in the objective function of the transportation problem is least D; Pt/ 2. Therefore
we have 2 Eie,, D; 22.2. In the Binary Search at the Endpoints phase of iteration k, we probe indices within distance d = 2²t around the two endpoints of the search range. Since a; is not found before iteration h; we conclude that 22 < Dj, which means that 2 < 2og D;). Let p(12) = Eier, Ps. We have
Σ EpC((a)) 54 Σ p,2% iel.
<8
Σ pr log(D.)
=8 Em log( + (1 (0((2))())))))))())())))) log(())
(
sEl2
iELz
ViElz
By concavity of the log(-) function and Jensen's 's incquality we have
8
Σ p. log(A)) + (1 p((2)) log(1) < 8 log i
Σv.D, + (1 (1-p(1a)) < max (2)
'To account for the case where E [0, 1) where log(n) < O, we impose a bound by taking the maximum of log(n) and O.

---

The last inequality follows from the fact that if EiEl, PrD; S we have EiEl, P;D; (1-p(((1)(15)) S and otherwise we have
log Σ +1- (1- (1-p(b)))) 1iEh
< log
EpiDi 1E
ΣEd
+ +15 log(2n) + log(n) +
3.2 Lower Bound
is well known that there is a lower bound S2 (p) on the expected query complexity for binary search [Mehlhorn] 1975]. We now show that there is lower bound of (l v) the expected query complexity as well, even on instances with H (p) = O. This shows that there must fundamentally be log 7) dependence on the earth mover's distance, even for instances where it is not absorbed by the dependence on the entropy.
Theorem 3.2. For any 7 E [n]] any comparison-batt based (deterministic or randomized) algorithm must make (llog 11) queries some instance where H (p) 0 and the earth mover's ' distance between and p is O())
Proof. Thanks Yao's principle, it suffices to give distribution over instances of this problem and argue that any deterministic algorithm has a large expectation over this distribution. Let the set of keys be [n]. We present family of problem instances 7,,-.. ....... I,- where each instance can happen with probability " In instance I,, the true access distribution is singleton location i, i.e., in I, have p; = 1 and p; = 0 for each
[n]] (i). In all the problem instances J.....; Ja the prediction
is the uniform distribution over [n] Note that for each instance 7;, we have H(p) 0. Also, the earth mover's distance between ý and p is at most 17.
Our claim is that any deterministic algorithm has an expected cost of S2 1) over this distribution. To see this, note that the expected cost of any deterministic algorithm over this distribution of instances exactly equals the cost of that algorithm on an instance /* where the true access distribution uniform over [n). Now by the lower bound of Mehlhorn [1975], the cost of any deterministic comparison- based algorithm on I is Q (p")), where p" Is the uniform distribution over [n]. To conclude the proof, note that H (pp)) (2(log
Combining Theorem 3.2 with the (2((P)) lower bound due to Mehlhorn 1975 results in the following worst-case lower bound, asymptotically matching Theorem
Corollary 3.3. Any comparison-basee algorithm for binary search with distributional predictions has worst-case expected query complexity Q (p) log(1)).
4 A Portfolio of Predictions
In the previous section we showed an algorithm that optimal given a single distributional prediction. Here we extend this result to the setting where there are m different distributions given as a prediction. of p given. Let The be the earth mover's s distance between P& and p. The goal is design an algorithm competitive with single best distribution Pa. That is, comparable to minkeema) [m]] log Me and H (p).
That is, for E {1, (1.,....,,)) there are predictions Pa = (fi,k..... .......P.... PP.kk
4.1 Algorithm for Multiple Predictions
We proceed in 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,.... iteration has phases.
Let [E, r] be the current search range. Let S₁ Ej=ePs be the remaining probability mass in the k'th prediction.
Bisection.

---

Let Eg be such that Σj=( Py.k I S and Ej PI. I 1 Sk.. That is, ts is the median 2
of the k'th distribution. Sort the indices k E [m] so that i ..... ..... bm For convenience, let to = and fm+1 = r. Perform a binary search on Utg dt, Otg..... atm+1 to find the interval where a E (t,1) for some j 1, ...., m). The new search range is I]. 1|.
[t + 1,12-1-1
Continue this for 2' steps, and if is not found, begin the second phase described below.
Binary Search the Endpoints. Let [e, r] be the current search range. Set d min(2? - -)- f). Check if a is in the range [e, l + d] or [r - d, r] (by comparing a to ac+d and ar-d). . If is in one of these ranges, say [e, d]. do regular binary search (by choosing the middle point of the range each time) on the range [e. + d]. until a found. Otherwise, start the next iteration with the new search range + d + 1,r -d d-1].
The algorithm continues until is found.
4.2 Analysis for Multiple Predictions
We now state the following theorem regarding the algorithm for multiple predictions. The overhead of using multiple predictions is log m factor. The proof is very similar to the proof of Theorem 3.] and has been deferred Appendix A
Theorem 4.1. Given different distributions, the expected query complexity of the algorithm is log(m) (H) + max(minke((m) log V). 0)).
Experiments
We now present empirical evaluation of the proposed algorithms both synthetic and real datasets. Our goal is to show how predictions can be used to improve the running time of traditional binary search approaches. Since our theoretical results are about query complexity, and keep the results pppllementatioonnndeeeeddde our main metric will be the number of comparisons performed by each method. Our implementation can be found at https:/gghhob.cooo//ddddaa/pssse
We compare the performance of the following algorithms:
Classic The prediction agnostic approach that recursively queries the midpoint of the array.
Bisection - The bisection algorithm recursively queries the median of the predicted distri- bution (when the predicted probability in the search range is O, this algorithm queries the midpoint of the array). This strategy is nearly optimal when the predicted distribution is correct Mchlhorn 1975]: however, it is not robust to errors in the predicted distribution.
Learned BST The algorithm described in Section 3 the parameter larger to broaden the search space in the very early iterations, setting d to min(28.2' -
We make one modification, setting
This is heuristic approach to make the Bisection algorithm more robust. Given a prediction á we generate a new distribution, = Ap + (1 A)u, where u is the uniform distribution [n]. We then run the Bisection algorithm q. In experiments, À 0.5 is used.
Convex Combination
5.1 Synthetic Data Experiments
We begin with experiments synthetic data where vary the prediction error in controlled environment to show the algorithms sensitivity and robustness to mispredictions.
In this setting, let the keyspace be the integers in [-105, 10 We then generate t = 10 independent points from normal distribution with mean and standard deviation 10, rounding down each the nearest integer. This results in a concentrated distribution in very large key space. The points form the predicted distribution, p. To generate the test distribution, we proceed the same manner, but shift the mean of the normal distribution away from 0 by some value > 0. Note that for 100 the train and test distributions have 0 overlap with high probability. For each value 8, we repeat the experiment times and report the average and standard deviation of the costs.

---

Cassic Sisect
25
Learned 5ST
Conver Combiration
10
50
1.00 350
200
250
350
250
Figure 1: Results for synthetic data experiments. The y-axis measures the average cost (query complexity) each algorithm and the x-axis measures the amount shift in the test distribution. The training and test data are regenerated 5 times. The solid lines the mean and the clouds around them are the standard deviation of these experiments.
1240
4.0
680
0.0
4.0
36.0
10
0.44
0.0
200
808
(a) AskUbuntu
(b) SuperUser
(c) StackOverfloo
Figure 2: The train and test distributions when 50 for the three datasets.
Our results for this setting can be found in Fig. I 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 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 O, performs the best, as expected, while the Learned BST approach exhibits some overhead due to hedging against possible errors. However, 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 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 slightly worse than the Classic baseline.
5.2 Real Data Experiments
Dataset Description. In order to test our approach on real-world data, we use temporal networks from Stanford Large Network Dataset Collection These datasets represent the interactions stack exchange websites StackOverflow, AskUbuntu, and SuperUser [Paranjape et al.. 2017]. In all cases, use the answers-to-quesiin user u answering user v's question at time t. In this interaction, u is the source and v is the target user. Our data sequences consist of the source users from each interaction sorted in increasing order of timestamp, and we restrict the dataset to the first one million entries.
dataset, which contains entries of the form (u, v, t), which represents
Keys, Predictions, and Test Data. For each data sequence, the set of elements in the first 10% of the sequence is used as the set of keys of the binary search trees. Let A be the remaining 90% of the sequence and let 01 < Q2 < ... < an be the set of keys. For each element x E A, if a; S < a,+1, 50, we use the first percent of A as training data and the
we replace by a, For t = 5. 10. rest as test data. The training and test data are used to obtain the predictions (0) and actual access distribution (p), respectively. To obtain these distributions we use the normalized frequencies of each key in the training and test data.
For completeness, we show the distributions of the keys when t = 50 both for the training set and the test in 2

---

Damt
20
(a) AskUbuntu
(b) SuperUser
(c)
StackOverflow
Figure 3: Results for real data experiments. The y-axis measures the average cost of each algorithm and the x-axis indicates the fraction of the dataset used for training
(a) AbuUbuntu
(b) SuperUser
(c) Sverkoverflow
Figure 4: Results for real data experiments. The y-axis measures the average cost of each algorithm and the x-axis indicates the logarithm of the earth mover's distance between ý and p.
We present the results on these experiments Figures and Figure 3 we plot the average cost of the algorithms against the size of the training data. As expect, the size of the training data increases, the performance of all distribution-depeedeen algorithms get better, as the distribution error decreases. We make this more precise Figure where we plot the average against log of the EMD error.
We note a few observations. The learning agnostic, Classic, is suboptimal in all but handful of cases, showing that there is value using the distribution of the data to improve performance. Second, we validate the theory, showing that the learned BST's performance degrades smoothly log 7) increases. Third, the convex combination heuristic is not very effective on real world data, giving only marginal improvements over the bisection method.
Finally, on both AskUbuntu and SuperUser datasets, the learned BST approach performs significantly better than all of the baselines, saving 20-25% comparisons average. Unlike the Bisection algorithm is also never worse than the Classic baseline. On the StackOverflow dataset our approach about 10% worse than bisection method, owing the distribution being less concentrated around the median. In these cases, the overhead of learned BST is apparent, given that the second phase is unlikely to be fruitful in the first few iterations.
Overall, these results show that the Learned BST method robust against errors, and performs well against other approaches. Further improving the constant factors so that the learned approach has strong worst-case guarantees and performs well against other learned approaches remains challenging open problem.
Conclusion
There has been growing line of work showing how to improve optimization algorithms using machine learned predictions. Predominately, prior work has leveraged non- non-probabilisti predictions, despite the fact that most ML systems, such as neural networks, output a distribution.
This work introduces a model where the prediction is a distribution. We show that algorithms can perform better by taking full advantage of the distributional nature of the prediction, and that reduction to point prediction is insufficient to provide competitive algorithms.
Given the breadth of work in the Algorithms with Predictions area Mitzenmacher and Vassilvitskii [2021) there wide variety open questions concerning how adapt algorithms to the setting of distributional predictions.
10

---

References
Priyank Agrawal, Eric Balkanski, Vasilis Gkatzelis, Tingting Ou, and Xizhi Tan. Learning-augee augmented mechanism design: Leveraging predictions for facility location. In Proceedings of the 23rd ACM Conference on Economics and Computation, pages 497-528, 2022.
Shabbir Ahmed. Two-stage stochastic integer programming: A brief introduction. Wiley encyclopedia of operations research and management science, pages 1-10, 2010.
Kareem Amin, Travis Dick, Mikhail Khodak, and Sergei Vassilvitskii. Private algorithms with private predictions. CoRR, abs/2210.11222, 2022. doi: 10.48500ARXI::2001112 URL
https://doi. http://....rr/......66660///1/// 2210. 11222
(01550//rXivv
Keerti Anand, Rong Ge, Amit Kumar, and Debmalya Panigrahi. Online algorithms with multiple predictions. In Proceedings of the 39th International Conference on Machine Learning, ICML 2022, 2022.
Spyros Angelopoulos. Marcin Bienkowski, Christoph Dürr, and Bertrand Simon. Contract scheduling with distributional and multiple advice. In Kate Larson, editor, Proceedings of the Thirty- International Joint Conference
Artificial Intelligence, IJCAI-24, pages 3652-3660, International Joint Conferences on Artificial Intelligence Organization, 8 2024. doi: 10.24963/ijeai Main Track.
URL https: org/ 10. 24963/292111 ijcai. 2024 404
Maria-Florina Balcan, Tuomas Sandholm, and Ellen Vitercik. Generalization in portfolio-based algorithm selection. In Thirty-Fisth AAAI Conference Artificial Intelligence, AAAI 2021, Thirty- Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021, pages 12225-1222222 12232. AAAI Press, 2021. URL https://ojs aaai. org/ index. php/ AAAI http://aatceeeeveeeeeeeeeeeeeeee 17451
Étienne Bamas, Andreas Maggiori, and Ola Svensson. The primal-dual method for learning aug- mented algorithms. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Iisuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, Decem- ber 6-12, 2020, virtual, 2020. URL www.tbasstt......... html 055441144a317799cccc66666 abtp://ooobbbboo---bbtatt.......
neurips cc/ wwwppppppp//aeecccooocccceeeeeee (eapeer20200haaa.
Dimitris Bertsimas and Vineet Goyal. On the power of robust solutions in two-stage stochastic and adaptive optimization problems. Mathematics of Operations Research, 35(2):284-305 2010.
Dimitris Bertsimas, Shimrit Shtern, and Bradley Sturt. Two-stage sample robust optimization.
Operations Research, 70(1):224-44,, 2022.
Omar Besbes, Will Ma, and Omar Mouchtaki. Beyond IID: data-driven decision-makinn in heteroge- neous environments. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NewrIPS 2022, New Orleans, LA, USA, November 28 December 9, 2022, 2022. URL http: /ttp:r//pppppeeeeeeee nips cc/paper files//aper/20222 -wwtraat-
hash hwwp://sbbbboooooododeeoooe
html
Sami Davies, Benjamin Moseley, Sergei Vassilvitskii, and Yuyan Wang. Predictive flows for faster ford-fulkerson. In International Conference on Muchine Leurning, pages 7231-7248. PMLR, 2023.
Iaas Diakonikolas, Vasilis Kontonis, Christos Tzamos, Ali Vakilian, and Nikos Zarifis. Learning online algorithms with distributional advice. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 133 of Proceedings of Machine Learning Research, pages 2687-2696. PMLR, 18-24 Jul 2021. URL htLiitss://proceedingg
mlr. Allr.pres//ee wwteess/v139//e htt.11////caakooooololll.... html
Michael Dinitz, Sungjin Im, Thomas Lavastida, Benjamin Moseley, and Sergei Vassilvitskii. Faster matchings via learned duals. In Mare" Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin,
11

---

Percy Liang, and Jennifer Wortman Vaughan. editors, Advances in Neural Information Process- ing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 64, 2021, virtual, pages 10393-10406, 2021. https: pttpttpp::///occeeddddi 2021
Michael Dinitz, Sungjin Im, Thomas Lavastida, Benjamin Moseley, and Sergei Vassilvitskii. Algo- rithms with prediction portfolios. In Sanmi Koyejo. S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA. USA, November 28- December 9, 2022, 2022. URL //aperss nips wwwpaper_ppapeeeeeeeeeeee ffllls// wwwessst-tttatctooo erence. htmll 22pp1177777221/122225566333333a2
Paul Dütting and Thomas Kesselheim. Posted pricing and prophet inequalities with inaccurate priors. In Anna R. Karlin, Nicole Immorlica, and Ramesh Johari, editors, Proceedings of the 2019 ACM Conference on Economics and Computation, EC 2019, Phoenix, AZ, USA, June 24-28, 2019, pages 111-129, ACM, 2019. doi: 10.1145/332522.33295766 URL luttps://dan org/ 1145/ 3326526.3329576
Sungjin Im, Ravi Kumar, Mahshid Montazer Qaem, and Manish Purohit. Online knapsack with frequency predictions. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pages 2733-2743, 2021. URL URL neurips html
Eniko Kevi and Kim Thang Nguyen. Online covering with multiple experts. CoRR, abs/2312.145664 2023. doi: 10.4550/ARXXW2232224556 14564
Sujin Kim, Raghu Pasupathy, and Shane G Henderson. A guide to sample average approximation. Handbook of simulation optimization, pages 207-243, 2015.
Donald E. Knuth. Optimum binary search trees. Acta Informatica, 1:14-25, 1971. doi: 10. 1001/ BF00264289. /ttp:/ttt/////ooorrrrr
1007 1007/BF0026442889
Tim Kraska, Alex Beutel, Ed H Chi, Jeffrey Dean, and Neoklis Polyzotis. The case for learned index structures. In Proceedings of the 2018 International Conference Management of Data, pages 489-504, ACM, 2018.
Silvio Lattanzi, Thomas Lavastida, Benjamin Moseley, and Sergei Vassilvitskii. Online scheduling via learned weights. In Shuchi Chawla, editor, Proceedings of the 2020 ACM-SIAM Symposium on Discrete Algorithms, SODA 2020, Salt Lake City, UT, USA, January 5-8, 2020, pages 1859- 1877. SIAM, 2020. doi: 0911337/997811119759999 9781611975994.11 114
org/ 1137/1
Honghao Lin, Tian Luo, and David Woodruff. Learning augmented binary search trees. In Interna- tional Conference on Machine Learning, pages 13440. PMLR, 2022.
Alexander Lindermayr and Nicole Megow. Algorithms with predictions. -ooorrtthhh--witt---eeeeeeee httprriiinnnnttuut.i..o github.io/ 2022.
https://
Thodoris Lykouris and Sergei Vassilvitskii. Competitive caching with machine learned advice.
Journal of the ACM (JACM), 68(4):1-25, 2021.
Samuel McCauley, Ben Moseley, Aidin Niaparast, and Shikha Singh. Online list labeling with predictions. Advances in Neural Information Processing Systems, 36. 2024.
Samuel Mccauley, Benjamin Moseley, Aidin Niaparast, and Shikha Singh. Incremental topological ordering and cycle detection with predictions. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 4/st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 35240-35254. PMLR, 21-27 Jul 2024. URL URI.tttss:/ htttts:::/poceeeiinnnn mlr, alr.press//e wtt.pprwwss/////// http:///33/3mmccccceeee mcc.mley22aaaaaa. html
12

---

Kurt Mehlhorn. Nearly optimal binary search trees. Acta Informatica, 5:287-295, 1975.
Michael Mitzenmacher and Sergei Vassilvitskii. Algorithms with Predictions, page 646-662. Cam- bridge University Press, 2021. doi: 10. 1011917/ 108637435.037.
Ashwin Paranjape, Austin R Benson, and Jure Leskovec. Motifs temporal networks. In Proceedings of the tenth ACM international conference web search and data mining, pages 601-610, 2017.
Manish Purohit, Zoya Svitkina, and Ravi Kumar. Improving online algorithms via ml predictions. In Advances in Neural Information Processing Systems, pages 9661-9670, 2018.
Chaitanya Swamy and David B Shmoys. Approximation algorithms for 2-stage stochastic optimiza- tion problems. ACM SIGACT News, 37(1):33 46, 2006.
Kapil Vaidya, Eric Knorr, Michael Mitzenmacher, and Tim Kraska. Partitioned learned bloom filters. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, OpenReview.net, 2021. http://ppe//oeeeeeeeeeeeeee../ep wtwp://rrrreeeeeeeeeermrrrmmiaae 6BRLOfrMhW
13

---

A Omitted Proofs
With probability Pt, the key a that we are looking for is aç. The goal is to bound the expected cost of the algorithm, which is Ei_ P,C(),), where C(a;) is the cost of the algorithm when a az. each step of the Bisection phase, a binary search is done on the medians of the predicted probability distributions on the current search range. When the binary search is done, for each prediction k, the median l the prediction falls outside of the new search range. This means that the probability mass of Pk the new search range has dropped by at least factor of compared to the initial search range before the binary search. Let k = arg min Th. From the above discussion, the probability mass of Py- in the search range drops by a factor of at least 2 in each step of the Bisection phase, which results in a total drop of at least 2 in iteration j of the algorithm. The cost of each Bisection step is log m, which makes the total cost of the Bisection phase in iteration j equal to (log m.) 2'. Let T, be the first iteration at which a is found, assuming a = 0g. The total Th cost of the first phase of the iterations 0 to T; is (log m) Σラ= 2 < (logm) 22:+1, Also, the cost
Proof Theorem 4.7 4.7
of the second phase in each iteration before T; 2, and in iteration T, at most 2T: So the total cost of the algorithm for iterations 0 to T; is at most (log m) 2 + +2
- 27; O(llo m) 2 We partition the keys based on T; into two sets, and bound the cost of each set separately. Let I1 = {i:
T; 2 log(log(4/p:)) and /2 == {i: T; > log(log(4/p.))
First, we bound the cost of indices in I by (log m) (pH())
Em(((a== = (Er, (flog m) 2%) (log m). 0 E Ps log(4/p)) = (logm)-O(H())) VEE
VEL
NEL
Now we bound the cost of the indices in I2 by (logm) m) (max(log(7x- ), 1)). Let i 1z. We know that during iteration j of searching for a, in the Bisection phase, the predicted probability mass pD-- in the search range decreases by a factor of at least 2 Therefore the predicted probability mass pD- in the search range [E, r at the end of the first phase in iteration T, is at most
2
2 Pi
S
Th-12 22
22
22
4/p:
2
where the inequality holds because E 1g. So Ej=ePjA S p2/2. Let D. := min(i l,r i). In the transportation problem corresponding to the earth mover's distance between p and pz., probability mass of at least p needs to be moved from point to the outside of the interval [e, r]. The cost of this movement in the objective function of the transportation problem is at least D; ps/22. Therefore we have 7ą- 2 Ever, D: ps/2. In the Binary Search at the Endpoints phase of iteration
j. we probe indices distance of d = 2 around the two endpoints of the search range. Since a; is not
found before iteration T;, we conclude that 22 < Di, which means that 2 < 2 2og( D;). Let p(12) = Eier, Ps. We have
Σ p.C(a.) S (logm) 0
Σ p.2
(1)
ieh
vE
S (log m)
Σ Ps log( Di))
iElz
< (log
0
Σ pi log( Di) + (1-p((())) log(1)
(3)
S (log m)
log
Ep:D++
(4)
vie (max(log(ne> 1)),
I (log m)
(5)
where inequality results from concavity of log(-) function and Jensen's inequality, and inequal- ity (5) is because of the following
14

---

If EiEt, PrD; I then we have
log (1:D:+(1-p(u)) log(2)
VEE
If Ei D;D; > then we have
log
Σ
PtD: (1 p(la))
log
EriDs VEE2
+1 +1= O(log(11 ))
VEL2
B Experimental Setup
We use Python 3.10 for conducting our experiments system equipped with an 11th Generation Intel Core CPU running 2.80GHz, 32GB of RAM, 128GB NVMe KIOXIA disk drive, and 64-bit Windows 10 Enterprise operating system. It's worth noting that the cost of the algorithms, i.e. the expected query complexity, is hardware-indee indeeeednn...
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the results are fully stated with assumptions.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: A formal model with all assumptions is explicitly stated.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: full proof is given all detail.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All details stated regarding experiments. Code is provided in public GitHub repository.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Data is publicly available. Code is provided in public GitHub repository.
6. Experimental Setting/Details
15

---

Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: details are given.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Error bars reported.
8. Experiments Compute Resources
Question: For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Answer: [Yes]
Justification:
Setup stated.
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: The authors have read the code of conduct and this research conforms.
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The work is well-known and widely used algorithm. There is need for safeguarding.
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Any code/data used is cited.
13. New Assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
Answer: [Yes]
Justification: Code is released a public GitHub repository.
Crowdsourcing and Research with Human Subjects
For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Question:
Answer: [NA]
16

---

Justification: No crowd-soureing or human subjects.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country institution) were obtained?
Answer: [NA]
Justification: No human subjects
17


Reviews:
[{"review_id": "BELGkUQqxK", "reviewer": "Reviewer_kf39", "summary": "This paper proposes a learning-augmented algorithm for searching in a sorted array. Different from all previous learning-augmented algorithms, it takes in distributional predictions. The main result is an algorithm with query complexity $O(H(p) + \\log \\eta)$, where $H(p)$ is the entropy of the true distribution and $\\eta$ is the Earth Mover's distance between the true and predicted distributions. The paper also includes proofs to show the theoretical optimality and experiments to validate the practical usefulness.", "questions": "- The algorithm in [Mitzenmacher and Vassilvitskii, 2021] that searches in a sorted list with predictions receive separate predictions for each query. Is this also the case for the setting discussed in this paper? \n    - If so, as I understand, the proposed algorithm needs to rebuild the binary search tree every time it receives a new query along with its distributional predictions. Then, this would lower bound the time complexity to answer each query with $O(n)$. Is that correct?\n- I am confused by the reference to \"binary search tree\" in the paper (even in the title). Does the proposed algorithm actually require building a binary search tree in its specification and implementation? Why does section 3 not contain any explanation related to the binary search tree?", "limitations": "Most of the limitations are discussed in the weakness section and questions section. In summary, I think the presentation of this paper is not clear enough. This creates difficulties for me to: 1. position the proposed algorithm in the context (of other recent learning augmented searching algorithms). 2. evaluate its novelty. 3. even understand its technical details (as mentioned in the questions section). If the author could provide a better explanation of the above-mentioned points, I would consider raising my score.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- The paper follows the recent line of work on \"learning-augmented algorithm\", or \"algorithm with predictions\". This is a promising new direction that tries to combine the theoretical soundness of classic algorithms with the learning ability of machine learning algorithms.\n- The section on theoretical analysis (though the ideas are simple) is effective.\n- The paper includes experimental results to back up the theory. The experimental settings are diverse. The performance of the proposed algorithm is strong compared to all the baselines.", "weaknesses": "- I think the presentation of the paper can be greatly improved. To list a few points:\n    - in line 50, the sentence \"That is, the prediction is a single (potentially high-dimensional) point (or maybe a small number of such points).\" is hard to read for me.\n    - in line 56, the sentence \"Or, can we in fact do better by taking full advantage of the entire predicted distribution.\" should end with a question mark.\n    - the abusive use of the word \"essentially\" greatly weakens the soundness of the paper (For example, in lines 66, 70, 78 87, 88). The expression \"essentially optimal\" should be clarified with collaboration on complexity and constants.\n- I think several key works in the field of learning-augmented algorithms are missing, which makes it hard to position this paper in the correct context. For example, I think the algorithms proposed in \"Learning-Augmented Binary Search Trees\" by Lin et al should be at least discussed and even compared against (now, this paper is only mentioned as a very general reference for learning-augmented data structures).\n- It is not straightforward to me why the techniques used in the proposed algorithm are novel and not hard to come up with. I encourage the author to make a clearer point on \"which components of the proposed algorithm are novel and different from existing techniques\".", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the insightful comments.  As the reviewer astutely observed, the bounds we claim are on query complexity, rather than running time. We apologize for conflating the two and we will make this more precise.\n\nQuery complexity is a standard metric when studying data structures. The number of queries is often a more informative measure of complexity.  For example, it could be the case that each comparison requires an I/O to disk (or network!), or a comparison may require an experiment that needs to be performed, so the rest of the computation has negligible cost. \n\nAs the reviewer points out, there are two potential ways to interpret the results.\n- Model 1: Queries arrive over time and there is a distribution over locations which arrives with each the query.  This is a strict generalization of [MV 21], where the prediction is a single location in the array. In their model, the location is error prone and in ours the distribution over locations is error prone. Our goal is to construct an effective search strategy given the query and the prediction.\n- Model 2:  there is a single (unknown) distribution over queries and the goal is to design a single binary search tree when given a (possibly erroneous) prediction of this distribution.  Rather than generalizing [MV 21], this generalizes classical work on optimal binary search trees ([Mehlhorn 1975]). So in this model, our work can be thought of as computing an optimal BST that is robust to a distribution shift.  Surprisingly, given the classical nature of computing optimal (or near-optimal) BSTs, this simple question of “what if my distribution is incorrect?” has not been considered in the data structures and algorithms literature.\n\nFrom the perspective of query complexity, these models are identical: one can treat the prediction in model 1 as the given distribution in model 2.  Any search strategy (from model 1) is a binary search tree (from model 2), and vice versa.  We will discuss this equivalence and clarify all of our claims to be about query complexity in the camera-ready. We note that the revision is minor (including the above discussion and slight changes in theorem statements).\n\n1. [MV 21] … receives separate predictions for each query. Is this also the case … in this paper?\n\nYes: we are given the prediction with the query---model 1 described above.  If instead we have a global distribution over queries (model 2), then the prediction comes *before* all of the queries, and constructing the BST is just preprocessing.\n\n2.  On the need to rebuild the binary search for every query … lower bound the time complexity to answer each query with $O(n)$. \n\nYes and no.  Yes, the actual time complexity would be larger than the query complexity. This is in some sense unavoidable due to the richness of the predictions.  But no, we do not actually build the binary search tree for each query.  Rather, we give a search algorithm which uses comparisons, where the number of comparisons is bounded by our theorems.  Of course, any search strategy is also a binary search tree.  So we can think of there being an “implicit” binary search tree that our search algorithm is using, even though it is not actually explicitly constructing a binary search tree.  If our focus were model 2, where we want to actually construct a binary search tree, then we would explicitly construct it (but would only do this once as preprocessing, since in that model there is a global distribution rather than one per query).\n\n3. Comparison to prior work. \n\nThe work by Lin et al. is the first foundational work on how to use predictions for data structures, and is related to the second model of our results (discussed above). However, their theoretical results are somewhat narrow in scope: they predict the rank ordering of the elements rather than their full distribution, and their measure of error is specifically tailored to distances on permutations (it only corresponds to a distance in distribution under the assumption that the distribution is Zipfian). Moreover, they make an assumption that the error is small (the values of $\\epsilon$ and $\\delta$ are constant).  One of the goals of our work was to relate performance to a standard error measure between two distributions (earthmover distance), without any a priori assumptions on its value.\n\nBecause their prediction metric is different (orderings vs full distributions), Lin et al. cannot compare their performance to that of the optimal binary search tree, and must settle for comparing to the optimal treap. As such, our results are not directly comparable with theirs.\n\n4. Our contributions. \n\nA major contribution of the paper is not in the techniques themselves, but in showing that the current literature on learning augmented algorithms that does not consider distributional predictions is suboptimal.  Specifically, as we discuss in the introduction, since modern ML algorithms naturally output a distribution, they are a more natural type of prediction.  Second, as we show in Section 2.1, distributional predictions are **strictly** more powerful than point predictions, so if we want optimal performance we should actually use the full distribution rather than flattening it to a singleton.\n\nWhile the results and the specific techniques may appear simple in hindsight, It is not a priori clear how to combine the doubling search of [MV 21], which is robust to errors but cannot handle distributions, with the classical distribution-based algorithm achieving H(p) query complexity, which is not robust to errors in the distribution.\n\nOverall, our work studies distributional predictions in a simple setting in order to ease the formalization of the model and emphasize the differences from point predictions. Now that we have demonstrated the importance of distributional predictions, we believe that there will be significant amounts of followup work expanding the algorithms with predictions literature to the case of distributional predictions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vBY79ve3K3", "reviewer": "Reviewer_mBjQ", "summary": "In this paper, the authors study the classical problem of binary searching over a sorted array in the learning-augmented setting, with distributional advice. Binary search in both the learning-augmented setting with point advice and the classical setting with known query distribution is well studied, but the authors are the first to study binary search in the learning-augmented setting with distributional advice. The authors argue that distributional advice is more natural than point advice as outputs of machine learning models such as neural networks, and show that distributional advice can actually perform asymptotically better than point advice.\n\nSpecifically, the authors consider the model in which the target element of the binary search is drawn from an unknown distribution, and aims to minimize the expected cost of searching the element. The search algorithm is also given a prediction on the underlying distribution; the prediction error is measured in the earth mover distance between the prediction and the true distribution. Their search algorithm is not trivial by any means, but very intuitive: On a high level the algorithm executes in iterations, where in each iteration it conducts median search based on the predicted distribution a set number of times depending on the current iteration, and checks both ends of the search range, finishing with a normal binary search if the target is near the endpoints. If the target is not found near the endpoints, start the next iteration on the middle range. \n\nTheir analysis is simple and shows that in expectation the algorithm uses $O(H(p) + \\log \\eta)$ comparisons, where $H$ is the binary entropy function, $p$ is the true distribution, and $\\eta$ is the earth mover distance between $p$ and the prediction $\\hat p$. As corollaries, they show that the algorithm can adapt to multiple predictions, and conduct experiments to show the performance of their algorithm on real and synthetic data sets.", "questions": "As I've outlined above, I don't find division between point predictions and distributional predictions as large as the authors painted. It is a personal opinion though, so I would like to hear the author's thought on this matter.\n\nIt would be interesting to see what other problems can distributional advice excel on. Binary search is still a relatively simple problem, so application of the same methodologies on problems with more/less structure would be much more convincing.", "limitations": "The authors discuss their limitations properly. There are no ethical concerns.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is very cleanly written and technically sound, and really does not have a lot to complain about. I was initially somewhat skeptical to see a significant new result about learning-augmented binary search, which is a very well-studied subject, but the authors managed to bring me a welcomed surprise.\n\nTheir \"claim to fame\" is on the novelty of the advice model, using a distributional prediction as opposed to a traditional point advice. While I have some minor questions about their model, they do indeed achieve a very impressive performance, and back it up with a matching lower bound analysis. A large part of projects in the field of learning-augmented algorithms, in my understanding, is about finding the right form of advice to use in conjunction of a traditional algorithm, and using more probabilistic, distributional advice is a direction I am quite interested in as well.", "weaknesses": "I have no major complaints with the paper, but in my opinion, the author's usage of distributional advice is, on a high level, similar to how other learning-augmented use their respective form of advice. The author's search algorithms rely on the median of the distribution to conduct bisection of the search range, which is in of itself not a probabilistic operation; It also relies on the fact that the input to the algorithm is an element drawn from a distribution. One can arguably draw parallel to other non-distributional learning-augmented algorithms for online problems that uses predictions on its input sequence, which is also unknown to the algorithm beforehand. There are also scenarios in both online and offline algorithms where the input may not be drawn independently, or may not be drawn from a distribution at all, where non-distributional learning-augmented algorithms may still hold value.\n\nI do not think this invalidates the author's work - I still find the paper really impressive, but personally I do not find the divide as great as the author painted.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We agree that it is natural to adapt the algorithms with predictions model to a distributional prediction by simply using the median as the predicted value.  However, if the benchmark is the optimal solution on the true distribution, using the median as a point prediction will not give strong guarantees compared to optimal. We prove an even stronger statement in Section 2.1 –  a separation between the point prediction setting and the distributional prediction setting.  In other words, in Section 2.1 we show that if you restrict yourself to a point prediction (whether it is the median of the given distributional prediction or something else), you simply must pay more than if you actually use the information given by the distributional prediction.  So while in some settings there might not be much difference between point and distributional predictions, in our setting (the most basic prediction setting!) there is.  \n\n\nThis paper takes a step in understanding a rich and largely unexplored area. That is, how can one design and analyze an algorithm for a problem when (1) the true distribution is stochastic and the benchmark is the optimal solution and (2) the algorithm is given a noisy distribution.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lNs9xkw6B5", "reviewer": "Reviewer_EXym", "summary": "The paper studies the problem of binary search for an item in a sorted array in a learning-augmented setting, where the element to search for is drawn from some unknown distribution, and the algorithm has access to a prediction of this distribution. This contrasts with the majority of work in the learning-augmented setting (aka algorithms with predictions), most of which assume that the prediction is an exact value rather than a probability distribution.\n\nThe main result is that there is an algorithm performing this task in $O(H(p)+\\log \\eta)$ expected time, where $H(p)$ is the entropy of the true distribution and $\\eta$ is the Earthmover distance between the predicted and the true distribution. This extends the known $O(H(p))$ result from the case where the exact distribution is known ($\\eta=0)$ to the case where only a prediction of the distribution is known. A matching lower bound on the dependence on $\\eta$ is also given.\n\nThe basic skeleton of the algorithm is to run the known algorithm for the case of known distribution. As the search interval shrinks during this process, it is repeatedly (in exponentially growing steps) checked whether the search item is close to the left or right end of the current search interval. If so, the algorithm concludes with a classical binary search at the respective (left or right) end.\n\nThe algorithm is further extended to the setting where $m$ predictions are given, losing a factor of $O(\\log m)$ in the guarantee with respect to the best of the given predictions. The paper concludes with an experimental evaluation, yielding results in line with the theory.", "questions": "1. Please clarify whether/how the claimed running time bounds can be attained, or if the bounds hold for the number of comparisons instead.\n2. I suggest expanding the discussion of related work. For example, the paper \"Beyond IID: data-driven decision-making in\nheterogeneous environments\" by Besbes et al seems to consider a similar setting where an algorithm has access to a prediction of a distribution.\n3. The text says that Figure 2 is for t=5, but the Figure caption says t=50 instead.", "limitations": "Not specifically discussed. A discussion about the possible weaknesses mentioned above would be useful.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "- Extending learning-augmented algorithms to the more realistic case where the prediction is a distribution, which is lacking in results at the moment\n- Tight bounds for a fundamental problem\n- Clearly written", "weaknesses": "- It is not clear whether the bounds hold for the running time, as claimed, or only for the number of comparisons. The main algorithm has a step where the index of a median element is identified, which the analysis does not seem to be account for. Moreover, the portfolio algorithm has a step where m indices are sorted, which would seem to take time $O(m \\log m)$ and increase the bound by an additional factor $m$. One might treat this as preprocessing time that is only incurred once before many binary searches are conducted. However, for such preprocessing to become negligible would seem to require many subsequent binary searches, yielding many samples from the true distribution, which would seem to allow approximating the true distribution very precisely rather than relying on a (collection of) prediction(s).\n- The experimental evaluation is restricted to settings where the distributions are highly concentrated around a single value.\n\nConsidering the authors' comments in the rebuttal to these two points, I'm increasing my score from 5 to 6.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough review.  We hope we can address each of your comments below.\n\n1. Running time versus comparisons.\n\nWe are sorry for the confusion; putting “running time” in our theorem statements was a mistake.  We meant query complexity (or number of comparisons), as the reviewer realized. See the discussion in response to reviewer kf39.\n\n2. The experimental evaluation is restricted to settings where the distributions are highly concentrated around a single value.\n\nThe distribution needs to be reasonably concentrated for our algorithm to outperform binary search (the best worst case algorithm). If the distribution is too close to uniform, standard binary search is the better thing to do (although, as the theory shows, even if the distribution is close to uniform we will not be too far from optimal). \n\nWe remark that our algorithm outperforms both the classic binary search and the bisection algorithm on several real data sets as shown in the paper. Further, the distribution does not need to be concentrated in order for our methods to see gains. For instance, in experiments below, we show improvements on bimodal distributions. \n\n3. I suggest expanding the discussion of related work. For example, the paper \"Beyond IID: data-driven decision-making in heterogeneous environments\" by Besbes et al seems to consider a similar setting where an algorithm has access to a prediction of a distribution.\n\nThank you for suggesting the work by Besbes et al. A key point of difference is that our algorithm does not need to know the error of the prediction ahead of time. In more detail, the work of Besbes et al. considers stochastic optimization problems such as the newsvendor, pricing, and ski-rental problems under distribution shift and analyzes the asymptotic performance of the sample-average-approximation (SAA) policy in this setting.  For pricing and ski-rental, they show that SAA is not robust to small distribution shifts, but that shifting by a small amount (depending on the amount of distribution shift) in the appropriate direction is robust.  In contrast, we make no assumption about where the distribution prediction comes from and our algorithm does not need to know how incorrect the distribution is to remain robust.\n\n4. The text says that Figure 2 is for t=5, but the Figure caption says t=50 instead.\n\nThank you for spotting the typo, the correct training data percentage is $t=50$. We will fix the discrepancy.\n\nAdditional experiments:\n\nWe have included in this rebuttal new synthetic experiments on bimodal distributions. The setting is similar to that of section 5.1, with the difference that both the predictions and the actual access distributions are bimodal. This is the details of the setting:\nThe keyspace is the integers in $[-10^5,10^5]$. To generate the predictions, we generate $10^4$ independent points from each of the following two distributions (rounding down each to the nearest integer):\n\na. A normal distribution with mean 0 and standard deviation 10.\n\nb. A normal distribution with mean $d$ and standard deviation 10.\n\nWe ran the experiments for different values of $d$, which is the distance between the means of the two normal distributions (or equivalently the distance between the two peaks of the final bimodal distribution). For each value of $d$, to generate the actual access distribution (i.e., the test data) we follow the same procedure, but we shift the two peaks of the distribution by some value $s>0$, i.e., the test data is generated by sampling $10^4$ points independently from each of the following two distributions:\n\na. A normal distribution with mean $s$ and standard deviation 10.\n\nb. A normal distribution with mean $s+d$ and standard deviation 10.\n\nFor each value of $d$ and $s$, we repeated the experiment 5 times and calculated the average number of comparisons of each algorithm. Here we report the results for $d=100$. \n\n\n\n| | s=0 | s=50 | s=100 | s=150 | s=200 | s=250 | s=300 | s=350 | \n| --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| Classic | 16.47 | 16.71 | 16.65 | 16.71 | 16.65 | 16.69 | 16.66 | 16.69 |\n| Bisection | 5.50 | 20.96 | 17.01 | 27.39 | 28.13 | 28.89 | 28.52 | 28.70 |\n| Learned BST | 9.51 | 9.86 | 9.99 | 10.01 | 12.43 | 14.37 | 15.51 | 18.69 |\n| Convex Combination | 6.50 | 16.09 | 12.07 | 17.30 | 17.62 | 17.69 | 17.67 | 17.69 |\n\nAs the results show, our algorithm outperforms the baselines even when there is a large shift in the test distribution. Similar results were obtained for different values of $d$, including $d=50,200,1000$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8jOS6eZYD6", "reviewer": "Reviewer_ud7J", "summary": "The main question guiding the paradigm of algorithms with predictions is the following: if we want to solve a given instance of a classical computer science problem, but are given predictions of some form (typically by some machine learning algorithm which has seen a lot of similar problem instances as training data), can we use the predictions to our benefit? By \"to our benefit\", we want two properties: 1) consistency i.e., if the prediction is indeed very accurate, then we want to fully utilize this, and not suffer the worst-case performance that we would have suffered had we not used the predictions at all and used the default prediction-agnostic algorithm. 2) robustness i.e., if the prediction is fully bogus, then we still want the performance of our algorithm to be no worse than the worst-case performance of the prediction-agnostic algorithm. The classic example in this space is that of performing binary search on a sorted array. Given a key query, we want to find if the key is present in the array. Suppose we have access to a prediction that tells us where we should look at in the array. The classic prediction-agnostic binary search algorithm has a worst case running time of $O(\\log n)$. But if the prediction is accurate, note that we can finish in constant time---so we want to utilize this somehow. It turns out that by seeding our binary search at the predicted value, we can achieve this best-of-both-worlds guarantee.\n\nWhile most of the existing literature on algorithms with predictions uses predictions that are point-valued predictions, this work considers the problem setting where the prediction is a distribution. The justification is that, modern day generative models naturally output whole distributions. The problem under consideration is still that of searching a key in a sorted array of size $n$. But now, we assume that the queries come from a distribution $p$. A classical \"median bisection\" algorithm due to Knuth (1971) and Mehlhorn (1975) shows that the best expected running time achievable is $\\Theta(H(p))$, where $H(p)$ is the entropy of $p$. Suppose the distribution $p$ is heavily concentrated on a few middle elements in the array. Then the classical binary search algorithm takes (expected) running time $\\Omega(\\log n)$, which is much worse than $H(p)=O(1)$. What if we are given point predictions? The authors give a simple argument that even then, an algorithm that uses such point predictions along with binary search is doomed to suffer a running time of $\\Omega(\\log n)$. \n\nBut what if we are given a predicted distribution $\\hat{p}$, which is supposed to be an estimate of $p$? Note that if $\\hat{p}=p$, we should in theory be able to obtain a running time of $H(p)$. But we don't know how well $\\hat{p}$ approximates $p$. The authors show that it is possible to make use of the predicted distribution $\\hat{p}$ to our benefit. Concretely, they give an algorithm that makes use of the explicit form of $\\hat{p}$, and has an expected running time of $O(H(p)+\\log \\eta)$, where $\\eta$ is the Earthmover distance between $\\hat{p}$ and $p$. Note that when $\\hat{p}\\approx p$, this achieves a running time of $O(H(p))$ which is what we wanted (and which is the best possible). Also, when $\\eta=n$, which is the largest it can be, the guarantee is still $O(H(p)+\\log n) = O(\\log n)$, since $H(p) \\le \\log n$. This is in fact the worst-case running time of the classic prediction-agnostic binary search algorithm. The algorithm that the authors propose is a careful combination of classic binary search and the median bisection algorithm (which uses $\\hat{p}$ to compute its medians). The analysis of the algorithm is insightful and easy-to-follow.\n\nIs $O(H(p)+\\log \\eta)$ the best attainable guarantee? Note that even if $\\hat{p}=p$, we can't do any better than $H(p)$ by the argument of Mehlhorn (1975). Thus, the question is if the $\\log \\eta$ factor is necessary. Indeed, the authors show it is, and hence their algorithm is optimal. The authors also propose an extended version of their algorithm which can use multiple predicted distributions, and gets a guarantee with respect to the predicted distribution that has the least Earthmover distance. Finally, the authors demonstrate the efficacy of their algorithm against baselines={classic binary search, classic median bisection} in synthetic simulations as well as real-world datasets. The experimental results are compelling and validate the theory to a large extent, showcasing the benefit of using the full-power of the distributional form of the prediction at hand.", "questions": "I am a little confused by the proof of the lower bound. Where exactly are you using the form of the predictions $\\hat{p}$ in your argument? In line 245, you simply say that, \"the best thing that an algorithm could do is build a balanced binary search tree on $[\\eta]$.\" It seems to me that this sentence stems from the way the distribution over $p$'s is constructed, and is agnostic to the specific form of $\\hat{p}$. Put another way, could you point out what is flawed in the following proof, which gets a $\\Omega(\\log{n})$ lower bound?\n\nTake $\\hat{p}$ to be uniform on all of $[n]$. The distribution over true instances $p$ (which are distributions over keys) is simply: draw $i^* \\in [n]$ uniformly at random, and set $p_i=1[i=i^*]$. Then, by the reasoning in line 244-245, any deterministic algorithm will have at least $\\Omega(n)$ keys such that the cost on these keys is $\\Omega(\\log n)$, because again, \"the best thing an algorithm could do is build a balanced binary search tree on $[n]$\". Hence expected cost is $(1/n)(\\Omega(n))(\\Omega(\\log n)) = \\Omega(\\log n)$. Proof concluded by noting EMD between $\\hat{p}$ and $p$ for any $i^*$ is at most $n$.\n\nDepending on clarification about this, I will either increase/decrease my score...", "limitations": "I do not foresee any limitations that the authors have not addressed.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "As mentioned, this paper departs from the usual algorithms with predictions paradigm which assumes access to point-valued predictions, but instead considers the setting where the prediction is a distribution in itself. The authors make a strong case for the richness of this seemingly more general problem setting: even for the textbook example of searching in a sorted array, the classic binary search algorithm with a point-valued prediction no longer attains the best achievable performance, and a non-trivial algorithm, that makes explicit use of the full form of the distribution-valued prediction, is necessary to achieve performance gains. Moreover, it is commendable that the algorithm presented by the authors is provably optimal. The work of the authors makes a strong case for studying other classical algorithms-with-predictions problems from the literature under the setting of distribution-valued predictions. The theory is also well-supported by the experiments.", "weaknesses": "The main concern I have is with the Lower Bound in section 3.2. See Questions below. I would also tone down the claim of \"*initiating* the study of algorithms with distributions as predictions\" (line 6, line 59), especially since, as the authors themselves mention, such a setting was also studied by Angelopoulos et al. (2024).\n\n---\n\nMinor/typos: \\\nLine 60: It seems that there is a slight change in problem setting here for the algorithm of Knuth (1971), where now we think of queries being drawn from a distribution $p$. This should be stated beforehand, otherwise saying things like \"median of the distribution\" does not make sense, when no distribution has been introduced. Also, $p$ is not introduced notationally. \\\nLine 184: $S$ is not introduced notationally. I believe it should be the probability mass on $[l,r]$ i.e., $S=\\sum_{i=l}^r \\hat{p}_i$ (correct me if I am wrong) \\\nAt a lot of places, parenthetical citations \\citep{} ought to be used (e.g., lines 117-118 in the related work, etc.)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank you for the points you have brought up in your review.  We will incorporate the minor comments you raised into the paper and clarify the notation.  As to your main comment about the lower bound, we will improve the rigor and also hope that the following discussion clarifies the lower bound for you.\n\nIt is true that the construction you gave gives a lower bound of $\\Omega(\\log n)$, when the earth mover distance between $\\hat p$ and $p$ is $\\approx n$.  But one might worry that this lower bound is only possible because of the large EMD; a priori, it could be the case that $\\log \\eta$ is a lower bound when the EMD is extremely large, but once it is smaller (we have a more accurate prediction) the dependence can be improved.  Since we are most concerned with the case of reasonably accurate predictions, if this were the case then our lower bound would not be particularly meaningful.      To handle this, the lower bound we present in the paper generalizes this idea to handle any value of $\\eta$, instead of just $\\eta = n$.  \n\nIn other words, the theorem we prove is the following: “For any $\\eta \\in [0,n]$, any comparison-based algorithm must make $\\Omega(\\log \\eta)$ queries on some instance where $H(p) = 0$ and the Earth Mover’s Distance between $p$ and $\\hat p$ is $O(\\eta)$.”\n\nThis is a stronger lower bound than the one suggested by the reviewer, since it holds for all $\\eta$ rather than just for when $\\eta = \\Theta(n)$.  Moreover, since the bad instances have $H(p) = 0$, it shows that there must fundamentally be a $\\log \\eta$ dependence on the Earth Mover’s Distance; this does not only happen when it can be absorbed by the dependence on the entropy.  We will calrify these points upon revision.\n \nRegarding Angelopoulos et al. (2024), this contemporary and independent paper studies how to obtain an optimal tradeoff between consistency and robustness with distributional predictions. However, their solution space explored is considerably more limited than ours, consisting of geometric sequences with a multiplicative ratio of 2, each characterized by its starting point. Additionally, the gap between consistency and robustness is relatively small (4 vs. 2.77), implying less room to leverage predictions. Finally, their bound analysis is restricted to cases where the error is sufficiently small.\nIn 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. We will ensure that this previous work receives proper credit.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper studies the query complexity of searching through a sorted array when the keys are drawn from a given distribution which is approximately known. An algorithm is proposed which interleaves two approaches: the classic binary search algorithm and the "median search" approach from [1,2]. The *query complexity* of the proposed algorithm is shown to be $O(H(p)+\log(\eta))$ where $H(p)$ is the entropy of the ground truth distribution and $\eta$ is the Wasserstein distance between the estimated sampling distribution and the ground truth distribution. 

There is a consensus between all four reviewers  (and myself) that this is an interesting and novel problem setting and algorithm which is well deserving of publication. However were some concerns, especially from reviewer kf39, regarding the presentation. I partially agree with the reviewer: on the one hand, I did find the introduction and abstract difficult to parse without jumping straight to the more mathematical "preliminaries" section, and on the other hand, it is also the case that the authors have **incorrectly used** terms such as "**time complexity**" instead of "query complexity" throughout the paper. They have admitted their mistake during the rebuttal and discussion with reviewers ud7J and kf39 and promised to modify the paper accordingly. Reviewer kf39 argues that the required changes might be too much and that the paper might need a fresh round of reviewing. I think this is a subjective question, but my professional opinion is on the side of the authors: removing any explicit mention of time complexity in the theoretical results can be achieved before the camera ready deadline. I also would not find it particularly shocking if some of the *experimental* results demonstrated superior performance in terms of runtime, even if the theory itself only applies to the query complexity, as long as the distinction is made clear.

 There is also enough time for the authors to improve the presentation of the work, and I strongly encourage the authors to do so: this paper's contribution was clearly appreciated by many reviewers and has the potential to be of great interest to the community, if written in a more reader friendly way. 











References:

[1] Donald E. Knuth. Optimum binary search trees, Acta Informatica,  1971

[2] Kurt Mehlhorn. Nearly optimal binary search trees, Acta Informatica, 1975.
Decision: Accept (poster)