Learning-Augmented Search Data Structures

Authors: Chunkai Fu, Brandon G. Nguyen, Jung Hoon Seo, Ryan S. Zesch, Samson Zhou
Primary area: learning theory
Keywords: learning-augmented algorithms, data structures

Abstract:
We study the integration of machine learning advice to improve upon traditional data structure designed for efficient search queries. Although there has been recent effort in improving the performance of binary search trees using machine learning advice, e.g., Lin et. al.  (ICML 2022), the resulting constructions nevertheless suffer from inherent weaknesses of binary search trees, such as complexity of maintaining balance across multiple updates and the inability to handle partially-ordered or high-dimensional datasets. For these reasons, we focus on skip lists and KD trees in this work. Given access to a possibly erroneous oracle that outputs estimated fractional frequencies for search queries on a set of items, we construct skip lists and KD trees that provably provides the optimal expected search time, within nearly a factor of two. In fact, our learning-augmented skip lists and KD trees are still optimal up to a constant factor, even if the oracle is only accurate within a constant factor. We also demonstrate robustness by showing that our data structures achieves an expected search time that is within a constant factor of an oblivious skip list/KD tree construction even when the predictions are arbitrarily incorrect. Finally, we empirically show that our learning-augmented search data structures outperforms their corresponding traditional analogs on both synthetic and real-world datasets.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 8

Reviews and discussion:
<thread>
creation date:
1730568459392

thread_type:
official_review

<thread_object>
creator: reviewer_ztzd

<content>
summary:
The paper studies search data structures augmented with an "oracle" (some ML-based prediction) of query frequencies, which can be used to improve over the classical setting that assumes no foresight into the query distribution. It focuses specifically on skip lists and KD-trees. It develops ML-augmented analogs, proves bounds on their performance under perfect and imperfect predictions, and presents an empirical evaluation.

strengths:
The problems are interesting and relevant, the theoretical results seem solid, and the empirical results show improvement. It is overall a solid submission.

weaknesses:
-- No related work section nor discussion, except for a couple of very specifically related recent/concurrent papers. No broader positioning in the literature; even though there is past work on very similar topics which goes unmentioned and unreferenced (like learning-augmented KD-trees in Cayton et al., see below). 

-- The font size in the figures is beyond excusable

questions:
How do your methods and results compare to the learning-based data structures of Ailon et al. Cayton et al., and other works related to those? They seem to study very similar problems in related settings.

Ailon, N., Chazelle, B., Clarkson, K. L., Liu, D., Mulzer, W., & Seshadhri, C. (2011). Self-improving algorithms. SIAM Journal on Computing, 40(2), 350-375.

Cayton, L., & Dasgupta, S. (2007). A learning framework for nearest neighbor search. Advances in Neural Information Processing Systems, 20.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730894306811

thread_type:
official_review

<thread_object>
creator: reviewer_i7vt

<content>
summary:
The  studies using machine learning models for improving other machine learning algorithms, namely search algorithms that use efficient data structures to organize data elements. 

The paper considers skip lists and KD-trees and augments them by adding an output of a learning algorithm. Skip lists organize data probabilistically in hierarchical and traversal is similar to binary search but with larger number of splits.

The paper consider unbalanced data distributions, mentioning that unbalanced distributions may result in suboptimal balancing behavior when underlying data structure construction algorithm expects balanced data. Authors propose to use an external machine learning algorithm that predicts statistics of the input data and provides an information to adjust the baseline search algorithm and data structure.

Proposed learning augmented skip list uses predicted probabilities to promote items to next levels with exponentially increasing required thresholds for predicted probabilities. 

While it may seem reasonable to use simple frequency estimation when number of possible items is not large, when number of items in the dataset is larger than number of encountered queries, having a machine learning oracle estimate expected frequency of previously unseen queries is a good idea. The work allows using such oracles to improve search data structures and proves better bounds on search time in general case and in case of an important and popular in real world datasets Zipfian distribution; and proves similar bounds for noisy oracle (representing real machine learning model with positive probability of incorrect guesses).

The idea behind learning augmented KD-tree is to make sure that each node or leaf has same probability on each tree height. Oracle predicted probabilities are used to ensure balanced splitting, improving search time if the oracle is usually correct. To limit algorithm complexity, low probability items are inserted into the tree using conventional approach. Similar bounds, as in case of skip lists, are proven for the augmented KD-tree.  

Empirically, the paper shows constant factor speedup of the augmented data structures, compared to the conventional approach

strengths:
- Provided approach is an elegant way to improve search data structure performance by considering probabilistic nature of the data. 
- Given the abundance of machine learning and statistical methods to serve as an oracle, proposed approach can be widely used in practice.
- New bounds were proven theoretically and experimental results partially support the theoretical findings.

weaknesses:
- From the paper it was not fully clear how to use proposed method with real machine learning oracle models (that non-trivially predict frequencies), instead of statistical oracles (that just calculate the table of frequencies).
- Some graphs, such as Figure 2, show a constant factor speed up. It would be nice to clarify what is theoretically predicted speedup in the cases of various datasets and parameters, and how theory aligns with practice.

questions:
- In the experiments, as I understood, simple history-based oracles were used. Is it the main use-case of the proposed method, or do we expect large oracles, such as deep learning models to be used as well?
- If the main use case will be using large machine learning oracles, I have the following question: for example, for image search, an oracle can be a neural network that takes an image and outputs expected query frequency; or for text search it could be a language model that outputs expected frequency of text or part of the text -- will the accuracy of real-world machine learning models be enough to guarantee improvement over conventional search structures?
- Would be nice to present graphs with theoretically expected improvement and actual improvement in various settings.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731097438675

thread_type:
official_review

<thread_object>
creator: reviewer_azhq

<content>
summary:
The paper proposes learning-augmented search data structures based on skip lists and KD trees. The authors achieve optimal expected search time and demonstrate robustness to inaccurate frequency predictors. Experimental results verify the effectiveness of the proposed algorithms.

strengths:
1. The integration of learned frequencies into skip lists and KD trees is well-constructed and achieves optimal performance.
2. Experiments show that the proposed algorithm outperforms classical algorithms and is robust to noise.
3. The paper is clearly written and easy to read. The ideas are simple, yet novel and effective. I appreciate the clear comparison

weaknesses:
1. The authors slightly overstate their contributions. For example, the claim of constant expected search time under the Zipfian distribution holds only when the exponent $s>1$, which is due to the entropy being of constant order (Lemma 2.3). In other words, for any data structure that achieves optimality has the same property. Additionally, in the abstract, the claim on the robustness when predictions are arbitrarily incorrect is not what you describe later.
2. The noise robustness measure, denoted as $(\alpha,\beta)$-noisy, is somewhat unconventional and may lack generality.
3. The experiments do not include comparisons with other learning-augmented algorithms given frequency predictions.
4. A closely related work [1] is not discussed in the paper, where the authors propose learning-augmented search trees that achieve optimality for arbitrary input distributions.

[1] Cao X, Chen J, Chen L, et al. Learning-augmented b-trees[J]. arXiv preprint arXiv:2211.09251, 2022.

questions:
1. In the abstract, the paper claims that for Zipfian distribution, the expected search time for an item is constant, while the traditional skip list or KD tree has an expected search time of $O(\log n)$. However, I could not find a detailed explanation for this claim in the main text. Specifically, is there theoretical support for vanilla skip lists (without learning) showing that the expected search time is $\Omega(\log n)$ under a Zipfian input distribution?
2. In the abstract, the paper claims for arbitrarily incorrect predictions, the algorithm is within a constant factor. Does this exactly come from your $(\alpha,\beta)$-noisy definition?
2. The noise robustness measure is unconventional. I suggest using KL divergence to quantify the distance between the predicted and true frequency distributions, as it may provide a more standard and interpretable measure of prediction accuracy.
3. I suggest adding additional experiment comparisons with other learning-augmented algorithms.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731173627435

thread_type:
official_review

<thread_object>
creator: reviewer_btgj

<content>
summary:
The paper considers the problem of augmenting search data structures like skip lists and kd trees with learning advice, in particular, an oracle that gives (possibly erroneous) estimates of the probability of query elements. When the oracle estimate is perfect, the search time for the proposed data structure is within a factor of 2 of the optimal expected search time (upto a constant additive slack). The estimator is also shown to be robust to errors in oracle estimates. For instance, even if the oracle is arbitrarily erroneous, the expected search time is still within a constant factor of the search time for the corresponding oblivious data structure. The superior performance of proposed data structures is also shown with experiments.

Some details of data structure construction: 

Skip lists: Skip lists are built in levels, with the bottom level as an ordered linked list. Higher levels store subsets of items to speed up search. Traditional skip lists promote items to higher levels randomly with a fixed probability (e.g., 1/2). In the proposed skip list, if an item's oracle probability exceeds a level dependent threshold, it is promoted up with probability 1 (this leads to faster access for more frequent items). For other items, they are still promoted up with a fixed probability (this leads to robustness to errors in oracle).

Kd tree: In traditional kd trees, each node splits the points according to the median across a selected dimension, iterating between different dimensions at different levels of the tree.  In the proposed data structure, one doesn't necessarily iterate between dimensions. The splitting point and dimension at any node is chosen such that the probability of the query being on either branch is as close to 0.5 as possible.

strengths:
The part on skip lists is a complete and meaningful contribution. The idea of promoting higher frequency elements to higher levels with more probability is both natural and easy to implement. Optimality as well as robustness of the proposed data structure is shown, and it is also shown to be superior in practice with both perfect and erroneous oracles.

weaknesses:
The part on kd trees felt rushed and left me confused about both the motivation and the details of the setting.

Motivation: In the paper, kd trees are considered for doing lookups for high dimensional points (as opposed to nearest neighbor search). But if one is just interested in lookups, it is unclear to me why we need kd trees. Why not just use something like the proposed skip lists after labeling the points from 1 to n? If we want to also support fast membership queries for frequent items not in the dataset, we can also add them to the skip list as is done in the current kd tree construction.

Setting: In the kd tree part, there is a big universe of items and a subset of it constitutes the dataset. The query can be any element in the universe. While in the skip list part, the  the query is always one of the dataset elements. Both settings seem reasonable but I was confused why different settings are considered in the two scenarios. It would be great if this can be clarified in the paper.

Details of the algorithm: It would be good to spend a paragraph discussing various algorithmic choices like was done in the skip list part (space doesn't seem to be a concern as the paper is already half page below the limit). For example, the motivation for handling low probability elements (prob less than 1/n^2) differently was unclear to me as there can be at most n of them. Also, if the authors can elaborate a bit more on how they are handling these points, that would be helpful.

questions:
My main question is about the motivation of the kd tree setting as stated in the weaknesses section. 

Another question: Would the proposed skip list data structure be able to handle deletions? For instance, if we delete some elements after constructing the dataset, would it still have the property that for the remaining elements, the data structure is close to optimal wrt the adjusted relative query frequencies? It would be good to include some discussion of deletion in the paper.

Minor formatting comment: The font size used for plots needs to be made larger.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732091360580

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank the reviewers for their insightful comments and for the positive feedback provided on the paper, including:
- The part on skip lists is a complete and meaningful contribution. (Reviewer btGJ)
- The idea of promoting higher frequency elements to higher levels with more probability is both natural and easy to implement. (Reviewer btGJ)
- Optimality as well as robustness of the proposed data structure is shown, and it is also shown to be superior in practice with both perfect and erroneous oracles. (Reviewer btGJ)
- Experimental results verify the effectiveness of the proposed algorithms. (Reviewer Azhq)
- The integration of learned frequencies into skip lists and KD trees is well-constructed and achieves optimal performance. (Reviewer Azhq)
- Experiments show that the proposed algorithm outperforms classical algorithms and is robust to noise. (Reviewer Azhq)
- The paper is clearly written and easy to read. The ideas are simple, yet novel and effective. (Reviewer Azhq)
- I appreciate the clear comparison (Reviewer Azhq)
- Provided approach is an elegant way to improve search data structure performance by considering probabilistic nature of the data. (Reviewer i7vt)
- Given the abundance of machine learning and statistical methods to serve as an oracle, proposed approach can be widely used in practice. (Reviewer i7vt)
- New bounds were proven theoretically and experimental results partially support the theoretical findings. (Reviewer i7vt)
- The problems are interesting and relevant, the theoretical results seem solid, and the empirical results show improvement. (Reviewer zTZd)
- It is overall a solid submission. (Reviewer zTZd)

We have uploaded a new version of the manuscript, incorporating reviewer suggestions and addressing the points raised, marking changes in blue text. In particular, we have included:
- An additional experiment comparing binary search trees and skip lists in the appendix.
- An extended discussion on related work in the appendix. 
- Additional intuition for the algorithmic design choices for learning-augmented kd-trees. 
- Commentary on the robustness of our schemes to arbitrarily bad performance by the oracle. 
- Replaced figures to increase the font sizes for improved legibility

We believe these revisions have strengthened the paper and look forward to further feedback. Below, we offer specific responses to the individual comments from each reviewer.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper augments two classic data structures, skip lists and kd trees, with advice from an ML oracle. The proposed solution is near-optimal when the oracle is accurate, and is robust to errors of the oracle. The paper provides theoretical guarantees for their approach, and the experiments show its effectiveness. All the reviewers had positive assessments of the paper, and I recommend acceptance.