Binary Search with Distributional Predictions

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.

It 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$.

In 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)$.”

This 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.

Regarding 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. In contrast, our work demonstrates how a binary search algorithm can compare generally to Earth Mover's Distance (EMD) and a lower bound on the optimum for any predicted distribution. As a result, we develop novel algorithmic solutions that build upon a close connection to EMD. We will ensure that this previous work receives proper credit.

Thank you for your thorough review. We hope we can address each of your comments below.

1. Running time versus comparisons.

We 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.

2. The experimental evaluation is restricted to settings where the distributions are highly concentrated around a single value.

The 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).

We 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.

3. 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.

Thank 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.

4. The text says that Figure 2 is for t=5, but the Figure caption says t=50 instead.

Thank you for spotting the typo, the correct training data percentage is $t=50$. We will fix the discrepancy.

Additional experiments:

We 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: The 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):

a. A normal distribution with mean 0 and standard deviation 10.

b. A normal distribution with mean $d$ and standard deviation 10.

We 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:

a. A normal distribution with mean $s$ and standard deviation 10.

b. A normal distribution with mean $s+d$ and standard deviation 10.

For 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$.

As 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$.

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.

This 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.

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.

Query 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.

As the reviewer points out, there are two potential ways to interpret the results.

• 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.
• 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.

From 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).

1. [MV 21] … receives separate predictions for each query. Is this also the case … in this paper?

Yes: 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.

2. On the need to rebuild the binary search for every query … lower bound the time complexity to answer each query with $O(n)$.

Yes 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).

3. Comparison to prior work.

The 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.

Because 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.

4. Our contributions.

A 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.

While 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.

Overall, 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.

Thank you to the author for their detailed rebuttal and explanations.

• After changing all occurrences of "time complexity" to "query complexity," I now understand the basic models in the paper. I agree that query complexity is an important metric for search algorithms and data structures.
• I found the explanations for models 1 and 2 very helpful and believe they should be included in the paper. These explanations provide the necessary context to understand the algorithm, which I previously lacked.
• I appreciate the clarification on the differences between this paper and the work by Lin et al. While these two works are not directly comparable, I strongly recommend including this discussion in the related work section. The fact that both focus on augmenting binary search trees with learnable advice makes this discussion essential. However, I have some remaining concerns about the paper:

1. The discussion of basic baselines seems insufficient. On line 173, the author uses the two-point distribution to claim that "converting pˆ to a point prediction and then using the algorithm of Mitzenmacher and Vassilvitskii [2021] as a black box is doomed to failure." However, one immediate idea for this example is to run two MV21 searches (as "black boxes"!) in parallel, starting at the two points, and the key could be found in constant queries. I believe this is a straightforward way to generalize MV21 to distributional settings, and the given example is too simple to demonstrate that this algorithm does not work. The author should include a stronger example to illustrate the suboptimality of "running a few MV21 in parallel."
2. I still find the use of "binary search tree" somewhat confusing. While I understand that every search algorithm can be formalized as a BST, I don't see how this phrasing clarifies the algorithm. Just as MV21 can be described as a BST, it is more naturally and simply described as a search algorithm. In the case of the proposed algorithm, I agree that the Bisection phase is similar to descending a balanced BST, but the second phase seems more like checking the left and right boundaries. Therefore, I believe the algorithm and the paper's title should be something like "Searching with Distributional Predictions." Does the author agree?
3. There is a lack of justification for the novelty of "algorithms with distributional prediction." In the author's response, they express the potential of distributional predictions as a major contribution of their paper. However, as the author pointed out in the related work section, the idea of using distributional predictions is not new, and basic tasks like ski-rental have been studied by Diakonikolas et al. [2021]. Therefore, I question the author's claim that "we believe that there will be significant amounts of follow-up work expanding the algorithms with predictions literature to the case of distributional predictions."
4. A new question (minor): In line 303, why set the exponential coefficient of $d$ to 8? If this is purely a result of hyperparameter tuning, is there any justification for why this would be preferable in all cases (rather than overfitting to the experiments)?
5. Results on time complexity. I understand that the proposed algorithm has different time and query complexities, but I think the results could be strengthened and made more applicable if the author also addresses the time complexity of the proposed algorithm. This is especially relevant in model 2, where all queries share the same distributional prediction.

Overall, I agree that the algorithm proposed is interesting and, as a result, I have changed my rating from 3 to 5. However, I think the paper requires revision (at the very least, changing all mentions of "running time" to "query complexity" in the main algorithms and major results, adding model 1&2 explanation, and adding related works) to be publishable. I hope the author can further address the concerns I have raised above.

Thank you to the authors for their response! It has clarified most of the issues. I suggest that the authors add discussions on the $log^2$ multi-modal distribution, the differences compared to Diakonikolas et al., and some commentary on time complexity in the next version of the paper. These points were previously unclear.

In summary, I believe this paper presents a new algorithm with sufficient theoretical analysis and experiments. If presented clearly, it is worthy of publication at NeurIPS. However, I have a major concern about the presentation: the edited version after the rebuttal might be too different from the originally submitted version. I found the original version difficult to understand, and only after the authors promised to change or add significant content in the rebuttal discussion, did I find the presentation acceptable. But this level of change during the rebuttal period may be too substantial. For example, the authors agreed with me to change the paper’s title; they also admitted that the original references to "time complexity" were incorrect and will be changed to "query complexity." However, reviewer ud7J, who recommended acceptance, understood the paper completely based on the "time complexity" framework. This discrepancy has led to a misunderstanding in the reviewers' comments. Therefore, I will maintain my rating at 5 and leave it to the area chair to decide whether this level of change during the rebuttal phase is too extensive.

Thank you again very much for the feedback! We will certainly include the discussions of model 1 and 2, and more comparison to Lin et al., in the paper as you suggest.

“The discussion of basic baselines seems insufficient”

We agree that the proof of our specific claim, that one cannot black-box reduce to a single point prediction, does not directly imply that one cannot reduce to collections of point predictions. However, it is relatively straightforward to generalize our proof to this stronger claim – instead of a 2 point distribution, consider a multi-modal distribution with $d = \log^2 n$ points. The method suggested by the reviewer, of doing a black-box reduction to MV for each of a collection of point predictions (by running an instance of MV in parallel for each of them), fails for this example. Either there is no prediction for one of the $d$ points, in which case every instantiation of MV will do many queries, or we do have a prediction for each of the points, and thus, when we run them in parallel we would probe all $d$ of them, again resulting in suboptimal query complexity. Developing an algorithm that can handle such multi-modal distributions (and generalizations) in an optimal manner is a significant part of our contribution.

“I still find the use of "binary search tree" somewhat confusing.”

We completely agree that it would be best to change the title of the paper as you suggest.

“There is a lack of justification for the novelty of "algorithms with distributional prediction."

We will include a more detailed discussion of Diakonikolas et al. We would like to highlight that our use of distributions is different. Specifically, the work of Diakonikolas et al. does not consider what happens when the distribution is erroneous (a major theme of our work), rather their focus is on minimizing the number of samples from the true distribution that they need. Since predictions are often erroneous, there is a dire need to make sure our usage of predictions is robust. This is what we focus on in this work, and believe would be of interest to the general community. Of course, the empirical distribution of samples can be viewed as an erroneous prediction of the true distribution, but our setting is far more general, and allows for general distributional predictions with general errors (measured by EMD). So we do believe that our setting is novel and will lead to significant follow-up work.

“A new question (minor): In line 303, why set the exponential coefficient of $d$ to 8? If this is purely a result of hyperparameter tuning, is there any justification for why this would be preferable in all cases (rather than overfitting to the experiments)?”

Setting the exponential coefficient of d to a small constant is preferable for improved empirical results, and does not change the asymptotic complexity. To see why, recall that the algorithm explores segments of length $2^{2^i}$ in the $i$-th iteration. When i is very small, these segments are very small, making the iterations overly fast and unlikely to succeed. We found that setting the length to $2^{8 \cdot 2^i}$ allowed us to balance this trade-off better. The exact setting of the exponential coefficient is not important (4 and 16 worked almost equally well). Also, see the further experimental results in our response to Reviewer EXym which take place on bimodal instances.

“Results on time complexity.”

We note that in model 2, the time complexity is basically equivalent to the query complexity. In this model, there is a single distribution over queries, so the time to build our BST is just preprocessing – we expect to answer far more queries than the time it takes to build the BST (as long as the preprocessing time is at least somewhat reasonable, as ours is). And then once the BST is built, the time it takes to search in it is essentially equal (up to constants involving following pointers) to the query complexity that we analyze. So in this setting, the query complexity is the time complexity!