Learning-Augmented Search Data Structures
Authors: Chunkai Fu
Brandon G. Nguyen
Jung Hoon Seo
Ryan S. Zesch
Samson Zhou
Venue: ICLR 2025 Poster (2025)
Keywords: learning-augmented algorithms
data structures
Primary area: learning theory

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.

Published as a conference paper at ICLR 2025
LEARNING-AUGMENTER SEARCH DATA STRUCTURES
Chunkai Fu" Brandon G. Nguyen* Jung Hoon Seo"
Ryan Zesch" Samson Zhou
ABSTRACT
We study the integration of machine learning advice 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 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 optimal up constant factor, even if the oracle is only accurate within constant factor. We also demonstrate robustness by showing that our data structures achieves expected search time that 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.
skip lists and KD trees are still
INTRODUCTION
As efficient data management has become increasingly crucial, the integration of machine learning (ML) has significantly improved the design and performance of traditional algorithms for many big data applications. Kraska al. (008)) first showed that MI could be incorporated to create data structures that support faster look-up operations while also saving an order-of-magnitude of memory compared optimized data structures oblivious to such ML heuristics. Subsequently, learning- augmented algorithms (Mitzenmacher & Vassilvitskij 2020) have been shown to achieve provable worst-case guarantees beyond the limitations of oblivious algorithms for a wide range of settings. For example, ML predictions have been utilized to achieve more efficient data structures (Mitzenmacher
2018. Lin et al. (022). algorithms with faster runtimes (Dinitz et al.] 2021 Chen et al. 2022c
Davies et al.) 2023), mechanisms with better accuracy -rrivcyyyyyyyyyyoceo tradeoffs (Khodak et al. 2022), online
algorithms with better performance than cntermatiionthhoreede limits (Purohit et al. |0101 Golllll lapudi
& Panigrahi 2019; Lattanzi et al. 2020
Wang et al.
2020 Wei & Zhang] 2020 Bamas et al. 2020
Anand et al. 2020 Almanza et al.
2021 Anand et al. 2021 Im et al. || Lykouris & Vassilvitskil
2021 Aamand et 20.2 Anand et al.| [022]] Azar et al. [022] Grigorescu et al. al.[2222)
Khodak
ct al. 2022 Jiang et al., 2022] Scully et al. [022] Antoniadis et al.) 22023 a (Shin et al. (002))
streaming algorithms with better accuracy-space tradeoffs (Hsu al. 2019 Indyk et al. 2019
Jiang et al. 2020: Chen et al.] 2022bb a Li et al. (023). and polynomial -ime algorithms beyond
20enn
hardness-of o limits, e.g.. NP-hardness (Ergun al. 2022 Nguyen al.) 2023 C.
ct al. 2024).
In this paper, we focus on the consolidation of ML advice improve data structures for the fun- damental problem of searching for elements among a large dataset. For this purpose, tree-based structures stand out as popular choice among other structures, particularly for their logarithmic average performance. However, these structures often have weaknesses for specific use cases that make them sub-optimal for various applications, which we now discuss.
"Fqual contribution, Texas A&M University, chunka Mhunka33698mma 3698gma1 i
com, bog@ttamu.edd
j. ses000117 (1o009977@mamii com, azeschâtamu. edu Texas A&M University, samsonzhou@ggaii coo

---

Published as a conference paper at ICLR 2025
Skip lists. One weakness of tree-based structures is that they need to be balanced for optimal performance, and thus their effectiveness is often closely tied to the order of element insertions. For example, the motivation of et al. (2022) to study
binary search trees noted that although previous results already characterized the statically optimal tree the underlying distribution known (Knuth 197T Mehlhorn 1977), these methods do not handle dynamic insertion operations. In contrast, skip lists, introduced by Pugh (1990a), maintain balance probabilistically. offering a simpler implementation while delivering substantial speed enhancements (Pugh) 1990b). Skip lists are generally built iteratively in levels. The bottom levels of the skip list is an ordinary- -inkee list, in which the items of the dataset are organized in order. Each higher level serves to accelerate the search for the lower levels, by storing only a subset of the items in the lower levels, also as ordered link list. Traditional skip lists are built by promoting each item in level to a higher level randomly with a fixed probability E (0,1).
Querying for a target element begins at the first element in the highest level and continues by searching along the linked list in the highest level until finding item whose value is at least that of the target element. If the found item is greater than the target element, the process is repeated after returning to the previous element and dropping to lower list. It can be shown that the expected number of steps in the search is O 1081/p n) so that serves as a trade-off parameter between the search time and the storage costs.
In many modern applications, skip lists are used because of their excellent search runtime and their space efficiency. Skip lists are often preferred over binary search trees due to their simplicity of implementation, their support for efficient range query, and their amenability to concurrent processes (Shavit & Lotan 2000 Lindén & Jonsson| 2013), high efficiency for dynamic datasets (Ge & Zdonik 2008; Pittard & Tharp 2010), network routing (Hu et al. 2003 Avin et al. 2020), and real-time analytics (Basin et al., 2020; Zhou et al., 2023). Thus while binary search trees have been a long-
standing choice for querying ordered elements, skip lists offer simpler, more efficient, and in some cases, necessary alternative.
KD trees. Another weakness of tree-based data structures is that they generally require the data to obey an absolute ordering. However, in many cases, e.g., geometric applications or multidimensional data, the input points can only be partially ordered. Thus in 1975, KD trees, which stand for k- dimensional trees, were proposed as more efficient alternative to binary search trees for searching in digee--dimesiional spaces in procedures such as nearest neighbor search or ray tracing for applications in computational geometry or computer vision. A KD tree works by picking a data point and splitting along some spatial dimension to partition the space. This process repeated until every data point is included in the tree, creating hierarchical tree structure that enables quick access to specific data points or ranges within the dataset.
Skewed distributions. Traditional search data structures treat each element equally when promoting the elements to higher levels. This balancing behavior facilitates good performance in expectation when query to the skip list is equally likely to be any dataset element. On the other hand, this behavior may limit the performance of the data structure when the incoming queries are from an unbalanced probability distribution.
Real-world applications can feature a diverse range of distribution patterns. One particularly common distribution is the Zipfian distribution, which is probability distribution that is a discrete counterpart of the continuous Pareto distribution, and is characterized by the principle that a small number of events occur very frequently, while a large number of events occur rarely.
In a Zipfian distribution, the frequency of an event (k; a;, N) is inversely proportional to its rank k raised to the power of (where a is positive parameter), in dataset of N elements. In particular, 1/K"
we have (k;; a, N) = ΣN _= (1/105' The value of a determines the steepness of the distribution so that smaller value, i.e., closer to 0. makes it more uniform, while larger a increases skewness.
Zipfian distributions provide a simple means for understanding phenomena in various fields involving rank and frequency, ranging from linguistics to economics, and from urban studies to information tech- nology. Indeed, they appear in many applications such as word frequencies in natural language (Wang & Wang] 2016; Blocki et al. 2018). city populations (Gabaix 1999 Vitanov & Ausloos (015). biological cellular distributions (Lazzarii et al. 2023), income distribution (Sandmo (0155). etc.
2

---

Published as a conference paper at ICLR 2025
Unfortunately, although Zipfian distributions are common in practice, their properties are generally not leveraged by traditional search data structures, which are oblivious to any information about the query distributions. To improve this performance bottleneck, we propose the augmentation of traditional skip lists and KD trees with "learned" advice, which (possibly erroneously) informs the data structure in advance about some useful statistics the incoming queries. Although model the data structure as having oracle access to the advice, in practice, such advice can often easily be acquired from machine learning heuristics trained for these statistics.
1.1 OUR CONTRIBUTIONS
We propose the incorporation of ML advice into the design of skip lists and KD trees to improve upon traditional data structure design. For ease of discussion in this section, we assume the items that may appear either in the data set or the query set can be associated with an integer in [N] := [I,...,N] which also in the case of high-dimensiones data, may be associated with a k-dimensional point. We allow the algorithm access to a possibly erroncous oracle that, for each E [N]. outputs quantity ps, which should be interpreted as an estimation for the proportion of search queries that will be made to the data structure for the item i. Hence, for each
[N], assume that p; [o, 1] and p1 +... + Pn. = 1. Note that these constraints can be easily enforced upon the oracle as pre-processing step prior to designing the skip list or KD tree data structure. We also assume that the oracle is readily accessible so that there is no cost for each interaction with the oracle. Consequently, we assume the algorithm has access to the predicted frequency p; by the oracle for all E [N]. On the other hand, we view sequence of queries as defining a probability distribution over the set of queries, so that Si is the true proportion of queries to item i, for each i E [N]. Although S; is the ground truth, our algorithms only have access to pi, which may or may not accurately capture fi-
skip list and KD trees, which gives expected search time at most 2C + = Se min (1og 1, log m). for some constant C > O. On the other hand, we show that any skip list or KD tree construction requires an expected search time of at least the entropy (.f) of the probability vector f. We recall that the entropy (f) is defined as H (f) >m= fi log T,
Consistency for accurate oracles. We introduce construction for a
Thus, our results indicate that within nearly factor of two, our learning auggmened search data structures are optimal for any distribution of queries. provided that the oracle is perfectly accurate. Moreover, even if the oracle on each estimated probability p; is only accurate up to a constant factor, search data structures are still optimal, up to a constant factor.
then our
Implications to Zipfian distributions. We describe the implications of our results to queries that follow a Ziplian distribution; analogous results hold for other skewed distributions, e.g., the geometric distribution. It is known that if the r-th most common query/item has proportion e for some >>> 1, then the entropy of the corresponding probability vector is a constant. Consequently, if the set of queries follows Zipfian distribution and the oracle is approximately accurate within constant factor, then the expected search time for an item by our search data structures is only a constant, independent of the total number of items, i.e., O (1). By comparison, a traditional skip list KD) tree will have expected search time O (logn).
Robustness to erroneous oracles. So far, our discussions have centered around an oracle that either produces estimated probabilities such that p; = fi or p. is within a constant factor of f.. However, in some cases, the machine learning algorithm serving as the oracle can be completely wrong. In particular, model that is trained on dataset before a distribution change, e.g., seasonal trends or other temporal shifts, can produce wildly inaccurate predictions. We show that our search data structures are robust to erroneous oracles. Specifically, we show that our algorithms achieve an expected search time that within a constant factor of an oblivious skip list or KD tree construction data structure achieves both consistency, i.c.,
when the predictions are incorrect. Therefore, good algorithmic performance when the oracle is accurate, and robustness, i.c., standard algorithmic performance when the oracle is inaccurate.
Empirical evaluations. Finally, we analyze our both synthetic and real-world datasets. Firstly, we compare the performance of traditional skip lists with our learning-augnnnn augmented skip lists on synthetically generated data following Zipfian distributions with various tail parameters. The dataset is created using four distinct a values ranging from 1.01
search data structures list on
3

---

Published as a conference paper at ICLR 2025
to 2, along with uniform dataset. During the assessment, we query specified number of items selectively chosen based on their frequency weights. Our results match our theory, showing that skip lists have faster query times, with an average speed-up factor ranging from 1.33 up to 7.76, depending on the different skewness parameters.
We then consider various datasets for internet traffic data, collected by AOL and by CAIDA, observed over various durations. For each dataset, we split the overall observation time into an early period, which serves as the training set for the oracle, and a later period, which serves as the query set for the skip list. The oracle trained using the IP addresses in the early periods outputs the probability of the appearance of a given node, and then the position of each node is determined.
skip list outperforms traditional skip lists with an average speed-up factor of 1.45 for the AOL dataset and 1.63 for the CAIDA dataset. Moreover, the insertion time of our learning-ause -augmented skip list is comparable with that of traditional skip lists on both synthetic and weal-world datasets. We also observe that our history-based oracle demonstrates good robustness against temporal change, with little shift in the dominant element set. The adopted datasets show that the set of the top frequent elements does not change much across the time intervals in the datasets used herein.
Our
We similarly perform evaluations on our learnig--augmenter KD tree data structure. We evaluate our KD tree data structure on Zipfian distributions with various tail parameters, and provide a heatmap of average lookup times for elements. We find that for a large variety of Zipfian parameters, ourt method is able to provides large improvements over traditional KD trees. We perform similar experiment under Zipfian distributions with added noise, and find our data structure still provides considerable improvements in query time.
We additionally evaluate our method on point cloud samples taken from a 3D model. We bin these samples in space, and create our leaming-augmented KD tree these binned samples. When querying this tree with new binned point samples, we find a decrease in average query time as compared to traditional KD tree under the same conditions. In addition this experiment, we provide results on real world datasets of n-grams and neuron activation, and similarly find improvements over traditional KD trees.
Concurrent and independent work. We mention that concurrent and independent of our work, Zeynali et al. (2024) used similar techniques to achieve the same guarantees on the performance of augmented skip lists that are robust to erroneous predictions. However, they do not show optimality for their auumeneg skip lists and arguably perform less exhaustive empirical evaluations. They also do not consider KD trees at all, which forms significant portion of our contribution, both theoretically and empirically.
Comparison to Lin et al. (2022). Our work was largely inspired by Lin et al. (2022), who observed that classical literature characterizing statically optimal binary search trees (Knuth) 1971 Mehlhom 1977) no longer apply in the dynamic setting, as elements arrive iteratively over time. Thus, they designed the construction of dynamic aagginn- binary search trees (BSTs). Their analysis for the expected search time utilized the notion of pivots within their trees and thus specialized to BSTs. Therefore, [Lin et (2022) explicitly listed skip trees and advanced tree data structures as interesting open directions. Qualitatively, our results are similar to Lin al. (2022), as are those of Zeynali et al. (2024). This is not quite altogether surprising because the main difference between these data structures is not necessarily the search time, but the either the ease of construction in the setting of skip lists, or the ability to handle dileii-dimensoonll data in the setting of KD trees.
somewhat
LEARNING-AUGMENTIN SKIP LISTS
In this section, we describe construction for consistency properties of the data structure. In particular, we show that up to a factor of two, our algorithm is optimal, given a perfect oracle. More realistically. if the oracle provides a constant-factor approximation to the probabilities of each element, our algorithm is still optimal up to a constant factor.
skip list and show various
We first describe our learning augmented skip list, which utilizes predictions p; for each item i e [n]. from an oracle. Similar to traditional skip list the bottom level of our skip list is ordinary-linked
4

---

Published as a conference paper at ICLR 2025
list that contains the sorted items of the dataset. As before, the purpose of each higher level is to accelerate the search for an item, but the process for promoting an item from a lower level to a higher level now utilizes the predictions. Whereas traditional skip lists promote each item in level to higher level randomly with a fixed probability p E (0, 1), we automatically promote the item to level l if its predicted frequency p; satisfies p; 2 2'=1, Otherwise, we promote the item with probability 5. This adaptation ensures that items with high predicted query frequencies will be promoted to higher levels of the skip list and thus be more likely to be found quickly.
X2
is worth addressing number of other natural approaches and their shortcomings. For example, one natural approach would be to use the "median" frequency across the items as a threshold to promote elements higher levels. However, this promotion scheme is not ideal because computing the median frequency at each time would either require an additional data structure for fast update time or increase the insertion time. potential approach to resolve this issue would be to use a separate threshold probability is set for each level so that only nodes with probability higher than the corresponding threshold are promoted to the next level. However, this approach seems to result in an unnecessarily large number of created levels if some item appears with small probability, e.g., We can thus first filters out the low- cow-frequency elements and place them remain on the bottom level of the skip list and then proceed with using separate threshold probability for each level. Unfortunately, this approach utterly fails to even match the search time performance of oblivious skip lists when the distribution is uniform, because all items will be in the same level, resulting in expected search time of (2)). Hence, we ensure that each element still has chance of being promoted to higher levels even when their probability is less than the corresponding threshold.
We again emphasize that due to the dynamic nature of the updates, existing results on statically optimal binary search trees (Knuth 197I, Mehlhorn 1977) do not apply, as observed by [in et al. (2022). We give the full details in Algorithm 1] For the sake of presentation, we focus on the setting where the queries are made to items in the dataset. However, we remark that our results generalize to the setting where queries can be made on the search space rather than the items in the dataset, provided the oracle is also appropriately adjusted to estimate the query distribution, using the approach we describe in Secciin 3
Algorithm Seaaring-auugmenee skip list
Require: Predicted frequencies Pi....s Pn. for each item in [n]] skip list
Ensure:
1: Insert all items at level ()
2: for each f do
3: if there are no items at level l - 1 then
4:
return the skip list
5: else
6:
for each i [m]] do
7: if predicted frequency pi 2 2f then Insert i into level
8:
9: else if is in level & - then
Insert i into level é with probability
We first show an upper bound on the expected search time of our learning augmented skip-list.
Theorem 2.1. For each E [n], let St and pt be the proportion of true and predicted queries to item i. Then with probability at least 0.99 over the randomness of the construction of the skip list, the expected search time over the choice of queries at most 20 + 2 E Sa min (oo J, loggn
To achieve Theorem 2.1
first show that each item i [n]] must be contained at some level max (0, 1 + [log(np)]].. depending on the predicted frequency p: of the item. We also show that with high probability, the total number of levels in the skip list is most O (log n). This allows us to upper bound the expected search time for item by at most 2C + 2 min (log p. 1,g We can then analyze the expected search time across the true probability distribution f.. Putting these steps together, obtain Theorem 2.1
5

---

Published as a conference paper at ICLR 2025
We also prove lower bound on the expected search time of an item drawn from probability distribution f for any skip list.
Theorem 2.2. Given a random variable X [n] so that X = i with probability fú. let T(X) denote the search time for X in a skip list. Then E (T)) 2 H(f), where (f) is the entropy of
Theorem 2.2 uses standard entropy arguments that have been previously used to lower bound the optimal constructions of data structures such as Huffman codes. We next upper bound the entropy of a probability vector that satisfies a Zipfian distribution with parameter s.
be fixed constants and let ʃ be a frequency vector such that Sa + for all i E [m). 8f8 > 1, then [J) = O (1) and otherwise if S 1, then H() & log n.
Lemma 2.3. Let 8,
By Theorem 2.1 and Lemma 2.3] we thus have the following corollary for the expected search time of our learning- augmented skip list on a set of search queries that follows Zipfian distribution.
Corollary 2.4. With high probability, the expected search time on a set of queries that follows a Zipfian distribution with exponent is most O (1) for
and O (log n) for S 1.
Next, we show that our peeming-augmentedd skip list construction is robust to somewhat inaccurate oracles. Let be the true-scaled frequency vector so that for each E [n]. fa is the probability that random query corresponds to i. Let p be the predicted frequency vector, so that for each i [ni], ps is the predicted probability that random query corresponds to i. For a, 8 E (0, 1), we call an oracle (a, -niis if for all E [n), we have 2 0. f; - 3. Then we have the following guarantees for an (a, (-)nn oracle:
Lemma 2.5. Let be a constant and 3 < in A learning augmented skip list with set of (a, B)- noisy predictions has performance that matches that of a learning learning-augmentee augmented learned with a perfect oracle, up an additive constant.
To achieve Lemma 2.5, (a,, -noiss oracles, we can write p; 2 6 fi. which allows us to express the search time log 1 in
we parameterize our analysis in Theorem 2.1
Due to the guarantees of the
terms of the true entropy of the distribution and small additive constant that stems from log a In fact, we remark that even when the predictions are arbitrarily inaccurate, our skip list still has expected query time O (log n)), since the total number of levels is at most O (log n) with high probability. Since the expected query list of an oblivious skip list is also O (logn) then the expected query time of our even with arbitrarily poor predictions.
1
-aggenned
skip list is within constant multiplicative factor,
3 LEARNING-AUGMENTTR KD TREES
In this section we present details on our novel approach to KD tree construction. First, we present the algorithm that constructs a learning-augmentee can be made on the search space rather than the items in the dataset, which is much more interesting for high- digh-diimnsional datasets, since even building balanced tree the search space could result in prohibitively high query time, as the height of the tree would already be at least the dimension d. Nevertheless, assuming that we have a query probability prediction p; for element i of our dataset, the intuition of our method is straightforward. Whereas a learning-ause augmentedd binary search tree would attempt find a value such that the probability of a query being on either branch of the tree is balanced, high-dimensional datasets do not have an absolute ordering.
KD tree. We focus on the setting where queries
Thus, instead of relying on standard techniques to determine the splitting point of our dataset, we find a specific dimension in which there exists balanced split such that the probability of query being on either branch of the tree balanced. However, there can still be high frequency queries that are not in the dataset, which can cause significantly high query time if not optimized. Hence, we also add to the tree construction high frequency queries that are not data points, in order reject these negative queries more quickly.
We prove the following guarantees on the performance of our learning KD tree, first assuming that our oracle is perfect.
Theorem 3.1. Suppose [A]d the space of possible input points and queries. Let N = A and p; be the probability that a random query is made to i [N], given the natural mapping between

---

Published as a conference paper at ICLR 2025
Algorithm 2 Learning-auugmentee KD tree construction
1:
function BUILDNODE(2)
2:
T 0
3:
if [ then
4:
T=x x return T best <-
5:
6:
7:
for each dimension i of x do for each element x do Compute the probability of points to the left of on axis If the probability is closer to 0.5 than best, update best
8:
9:
10:
11:
Taxis best.axis
12:
T.left (UILDNODE(2: a s[axis] & best.value)
13:
Trright BUILDNODE(r: x saxxis] best.value)
14:
return
Algorithm 3
KD tree
1: function BUI.D(dataset, queries)
2: dataset, + {a dataset x.prob
3: queries, - {x queries x.prob >
4: T Bun.DNooorddattsett queriess)
5: Insert {ar e dataset a.probb I n² into T using standard balanced KD tree construction
6: return
[N] and [A] Letp- (PN.....PPN) E RN be the probability vector and Hp) be its entropy. Then given a set n input points, the expected query time for the tree T created by Algorithm O (min H (p), (p),,ll)) n)).
The analysis of Theorem 3.1 corresponding to our learning augmented KD tree follows from similar structure as the proof of Theorem 2.1] However, the crucial difference that the universe size is now N, which is exponential in d. Thus constructions that consider distributions suffer O (log N) = O (dlog A) query time, which can be prohibitively expensive for large d, e.g.. high-dimensional data. Hence, our algorithm requires bit more care in the truncation of queries with low probability and instead, we build balanced KD tree for any item with less than , probability of being queried, so that each of their query times is at most O (logn) n). We further remark this implies robustness of our data structure to arbitrarily poor predictions, by similar argument as in Secciion 2
all of [N] may
We next prove lower bound on the expected search time of an item drawn from probability distribution f for any KD tree.
Theorem 3.2. Given random variable X E [n] so that X with probability fo. let D(X) denote the depth for X in a learning-augee augmeeted KD tree. Then E [D()) 2 H(f), where H(f) is the entropy of
By Theorem 3.1 and Lemma 2.3
we thus have the following corollary for the expected query time on our Baaming-augmentt KD tree on a set of search queries that follows a Zipfian distribution.
Corollary 3.3. With high probability, the expected query time on of queries that follows Zipfian distribution with exponent is at most O (1) for & > and O (log n) for s S
Finally, we show neai-optimality when given imperfect predictions from a (a, nn-is oracle:
Lemma 3.4. Let be constant and let 3 < 22 Then the query time for our learning-auue augmented KD tree with (a, B)-noisy prediction matches the performance of constructed using a perfect oracle up additive constant.
KD tre
7

---

Published as a conference paper at ICLR 2025
4 EMPIRICAL EVALUATIONS
In this section, we describe a number of empirical evaluations demonstrating the efficiency of our search data structures both synthetic and real-world datasets. We provide additional experiments in Appendix D
Skip lists on CAIDA dataset. In the CAIDA datasets (CAIDA] 2016). the receiver IP addresses from one minute of the internet flow data are extracted for testing, which contains over 650k unique IP addresses of the 30 million queries. Given that the log-log plot of the frequency of all nodes in the CAIDA datasets follows approximately a straight line in Figure I the CAIDA datasets can be approximately characterized by an a factor of 1.37. The insertion time is similar between classic and augmented skip lists, while Figure 2 shows that query time is almost halved when using the learning augmented skip lists at different query sizes. These results assume that the predicted frequency of all items in the query stream is accurate, i.e., the probability vector that is used to build the skip list matches exactly the query stream. The speed-up between the query times of the largest learning-augmentedd and the oblivious skip lists in Figure is roughly 1.86 x, which is surprisingly and perhaps coincidentally close to our theoretical speed-up of roughly 1.81x x on a Zipfian dataset with exponent 1.37.
CAIDA receiver IP frequency distribution
CAIDA receiver frequency distribution
3.0
10a
2.5
Peouancy
2.0
Requency
10
1.5
102
21.0
100
0.5
0.0
10a
10a
10a
10'
10a Indlex
10
10a
10"
10
10
10
10'
100
10
10*
Index
(a) CAIDA data distribution
(b) Zipfian fit (a = 1.37)
Figure I: CAIDA datasets distribution characterization in Figure Ia The nearly straight-fitted curve in Figure 1b immpliss that a Zipfian distribution with a distribution.
1.37 is a good fit to the CAIDA dataset
Insertion Time of Vanilla and Learned Skiplist
Query Time of Vanilla and Learned Skiplist
Vanilia Insertion Time
Vanilla Time
Insertion time (secords)
Learned Insertion Time
Query Time (seconds)
120
Leamedd Query Time
100
80
15
60
10
40
6
20
16.521
23.630
12,805
10.148
633
4.019
2.755
1522
dees
0.490
0.330
0.163
0,009
0.044
0.011
0.0055
0-
2.987
2.950
2.013
2.501
2,682
2.815
2.162
2,682
2007
..10
0.67
0.674
2.
2.35
3.83
3.
Number of unique IPs (Unit: 10k)
loo Irequent Query Count (Unit: 10MM)
(a) Insert time on CAIDA
(b) Query time on CAIDA
Figure 2: Comparison of insertion and query time on CAIDA for classic and augmmented skip lists. This figure compares the insertion and query times under varying numbers of top frequently accessed unique IPs between classic and augmented implementations. The horizontal axis in the two subfigures depicts the same scheme of IP selection, represented in two different ways, e.g., the top 29.9 million queries contain 665210 unique IPs, the next 29.5 million queries comprise 296384 unique IPs, etc.
Next, we demonstrate that our proposed algorithm still manages to outperform the classic skip list when temporal change exists in the probability vector by comparing the query time for the same set of query elements with different probability vectors being used to guide the building of the structure. For the skip list augmented by a noisy probability vector, the probability vector of elements
8

---

Published as a conference paper at ICLR 2025
during period of T1 is used as the predicted frequencies. The skip list being augmented by this probability vector has its own set of elements to be organized into the target skip list. Suppose the historic data from TI contains set of elements S1, and some future query stream contains set of elements S2. For each element in our target set S2, if the element is present in S1, then the occurrence probability of this element from S1 will be used build S2; otherwise, if the element has not shown up during TI (i.e., in SI), then we assume its probability to be O. After this, the probability vector is normalized sum to 1, resulting in a predicted probability vector be used to build a skip list based on the historic element frequency. Since there is temporal changes in the frequency of elements being queried, the predicted probability vector will show a discrepancy with the true probability vector. The results are presented in Figure 3a and we show that even when the prediction is not perfect, the augmented skip list still performs better than conventional skip list.
Pairwiss ise Intersection Incex Metrix
125010
Comparison of Query Times: Learned vs Vanilla
130000 0.9
0,90 960 960 965 1000 950 000.000.0...
Ouery time of Vanils Sciplit
1300000
9:970 570 960 960 96) 950 930
0.99
1000
Query Time Learned Skiples
Query Time Porfect Learned Seiplis
130200
130300
9896.900.0100
0.98
130400
5.500.9 0900.990.90
0.900
600
130700
1990.570.06
130800
0,500.0
130000
950
sin
970
0:50
200
131200
oep obo sta 970.
19000
199799
100000
3000
131000
10.2
ReterenceeTent Scheme
(a) Robustness test on CAIDA datasets
(b) Oracle credibility on CAIDA datasets
skip list erroneous oracles. In Figure 3b the labels on the axis indicate the time stamp that the internet trace data is collected, e.g., 130100 means the collection starts 13:01:00 :00 and lasts for 1 minute.
Figure 3: Robustness of our
Figure 3a shows that the skip list with perfect learning shows the best performance, while the skip list augmented with noisy learning performs very close to the scenario with perfect predictions. Moreover, the closer the test data is to the reference data chronologically, the closer the noisy-augmenter skip
list will perform to the perfect learning skip list. The CAIDA datasets used in this study contain 12 minutes of internet flow data, which totals around 444 million queries. The indices on the x-axis in Figure 3a means:
10 2: the first 10 minutes of data are used to create the reference (i.e., oracle) and the last
minutes are used to build and test the total query time using the former as reference.
2.2: the 9th and 10th minutes data is used as reference and the last minutes are used for testing.
3 3: the 1st, 2nd and 3rd minutes of data are used to create reference and the 4th, 5th and 6th minutes of data are used for testing.
6 6: the first 6 minutes are used create the reference and the last minutes are used for testing.
Further analysis of the temporal change of item frequency shows the reason behind the good perfor- mance of the history- based oracle. Figure 3b shows the change of intersection index between any 2 given minutes among the 12 minutes of CAIDA data. The intersection index is defined as the ratio of the number of shared queries to the total number of queries of any given 2 minutes of queries. Figure shows that the number of intersects queries has decreased by about 6% after 12 minutes, which indicates that the probability of the majority of the elements will be predicted with good accuracy, resulting in good oracle performance.
KD trees on synthetic datasets. KD Trees are commonly used in the field of computer graphics, with applications in collision detection, ray-tracing, and reconstruction. We first generate datasets of 2 points in 3-dimensional space, with frequencies given by a fixed Zipfian distribution with parameters a 5, = - parameters at which our method greatly outperforms a standard KD tree. In order to simulate constructing the tree noisy data, we multiply the ground truth query probabilities by
9

---

Published as a conference paper at ICLR 2025
Iraditional KD-Iree
Learning Augmented KD-Iree
6.
0U
Addrive Nose
Addrive Noizs
uodio AvarD owe
1.0
1.0
tblicctvvee
Multipiicativ Noise
Figure 4: Query time comparison for standard and boarning-augmeee KD trees with various noise.
numbers sampled uniformly from to M, and then add numbers uniformly sampled from 0 to A, before renormalizing to form a valid probability distribution. We query the tree 2 times, with point queries selected by the ground truth Zipfian distribution. We repeat this process 32 times, and report the median of the average query depth across all runs in Figure 4 We find that our method continues to outperform traditional KD trees under moderate amounts of noise, and at worst, performs on-par with traditional KD tree.
Next, generate datasets of 2 points in 3-dimensional space, with frequencies given by Zipfian distribution with parameters a, b. In the left plot, we assign these Zipfian weights randomly. In the right plot, however, assign Zipfian weights with ranks decreasing with the distance some random data point. We then query the tree 214 times, with point queries selected by the same Zipfian distribution. We repeat this process 32 times, and report the median of the average query depth across all runs. We find that, when points frequencies are distributed smoothly over space, our method's performance increases on less skew distributions, as seen in this Figure
Random Dist
Smooth Dist in Space
a.0
a.0
Deshh Aarro 3is
10.0
10.0
Figure 5: Comparison of query time on
KD trees with and without smooth spatial
distribution across various Zipfian parameters
KD trees on 3D point-cloud datasets. Finally, we evaluate our method on point cloud data generated from the Stanford Lucy mesh (Stan), with dimensions 1000 500 1500. We first uniformly sample 22 points along the mesh surface, and bin points with resolution 10, and assign lookup frequencies by the number of bin occupants. This results in 32k bins. Note, the resulting frequency distribution for binned cells is not highly skewed.
We then generate
of 26 surface samples the mesh binning them and assigning frequencies in the same way. When looking up with the new samples, our method yields an average query depth of 15.1, while traditional KD tree yields average lookup depth of 17.6.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENTS
Samson Zhou is supported in part by NSF CCF-2335411. The work was conducted in part while Samson Zhou was visiting the Simons Institute for the Theory of Computing as part of the Sublinear Algorithms program.
REFERENCES
Anders Aamand, Justin Y. Chen, and Piotr Indyk. (optimal) online bipartite matching with degree information. Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems, NeurlPS, 2022.
Nir Ailon, Bernard Chazelle, Kenneth L. Clarkson, Ding Liu, Wolfgang Mulzer, and C. Seshadhri. Self-improving algorithms. SIAM J. Comput., 40(2):350-375, 2011.
Laurence Aitchison, Nicola Corradi, and Peter E Latham. Zipf's law arises naturally in structured, high-dimensional data. arXiv preprint arXiv: 1407.7135. 2014.
Brian Allen and J. Ian Munro. Self-organizing binary search trees. J. ACM, 25(4):226-535, 1978.
Matteo Almanza, Flavio Chierichetti, Silvio Lattanzi, Alessandro Panconesi, and Giuseppe Re. Online facility location with multiple advice. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems, NeurlPS, pp. 4661 4673, 2021.
Keerti Anand, Rong Ge, and Debmalya Panigrahi. Customizing MI. In Proceedings of the 37th International Conference on Machine Learning, ICML, pp. 303-313, 2020.
MI. predictioons for online algorithms.
Keerti Anand, Rong Ge, Amit Kumar, and Debmalya Panigrahi. A regression approach to learning- augmented online algorithms. In Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems, NeurlPS, pp. 30504-30517, 2021.
Keerti Anand, Rong Ge, Amit Kumar, and Debmalya Panigrahi. Online algorithms with multiple predictions. In International Conference Machine Learning, ICML, pp. 582-598, 2022.
Antonios Antoniadis, Christian Coester, Marek Eliás, Adam Polak, and Bertrand Simon. Online metric algorithms with untrusted predictions. ACM Trans. Algorithms, 19(2):19:1119:1 2023a.
Antonios Antoniadis, Themis Gouleakis, Pieter Kleer, and Pavel Kolev. Secretary and online matching problems with machine learned advice. Discret. Optim., 48(Part 2): 2023b.
Chen Avin, losif Salem, and Stefan Schmid. Working set theorems for routing in self-adjusting skip list networks. IEEE INFOCOM 2020-IEEE Conference on Computer Communications, pp. 2175-2184, IEEE, 2020.
Yossi Azar, Debmalya Panigrahi, and Noam Touitou. Online graph algorithms with predictions. In Proceedings of the 2022 ACM-SIAM Symposium on Discrete Algorithms, SODA, pp. 35-66, 2022.
Étienne Bamas, Andreas Maggiori, and Ola Svensson. The primal-dual method for learning aug- mented algorithms. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems, NeurIPS, 2020.
Dmitry Basin, Edward Bortnikov, Anastasia Braginsky, Guy Golan-Gueta, Eshcar Hillel, Idit Keidar, and Moshe Sulamy. Kiwi: A key-value map for scalable real-time analyties. ACM Transactions on Parallel Computing (TOPC), 7(3):1-28, 2020.
Jeremiah Blocki, Benjamin Harsha, and Samson Zhou. On the economics of offline password cracking. In IEEE Symposium on Security and Privacy, SP, Proceedings, pp. 853-871. IEEE Computer Society, 2018.
Karthik C. S., Euiwoong Lee, Yuval Rabani, Chris Schwiegelshohn, and Samson Zhou. On approx- imability of E min-sum clustering. CoRR, abs/2412.03332. 2024.

---

Published as a conference paper at ICLR 2025
CAIDA.
The caida ucsd anonymized internet traces. [ttpss//ww..aaccaoor03 patasets/passiige dataset 2016.
Xinyuan Cao, Jingbang Chen, Li Chen, Chris Lambert, Richard Peng, and Daniel Sleator. Learning- augmented b-trees, 2023.
Lawrence Cayton and Sanjoy Dasgupta. A learning framework for nearest neighbor search. In Advances in Neural Information Processing Systems 20, Proceedings of the Twenty-First Annual Conference on Neural Information Processing Systems, pp. 233-240) 2007.
Justin Y. Chen, Talya Eden, Piotr Indyk. Honghao Lin, Shyam Narayanan, Ronitt Rubinfeld, Sandeep Silwal, Tal Wagner, David P. Woodruff, and Michael Zhang. Triangle and four cycle counting with predictions in graph streams, In The Tenth International Conference on Learning Representations, ICLR, 2022a.
Justin Y. Chen, Piotr Indyk, and Tal Wagner. Streaming algorithms for support-aware histograms, In International Conference on Machine Learning, ICML, pp. 3184 3203, 2022b.
Justin Y. Chen, Sandeep Silwal, Ali Vakilian, and Fred Zhang. Faster fundamental graph algorithms via learned predictions. In International Conference on Machine Learning. ICML, pp. 3583-3602, 2022c.
Valentina Ciriani, Paolo Ferragina, Fabrizio Luccio, and S. Muthukrishnan. Static optimality theorem for external memory string access. In 43rd Symposium on Foundations of Computer Science (FOCS, Proceedings, pp. 219-227. IEEE Computer Society, 2002.
Thomas M Cover. Elements of information theory. John Wiley Sons, 1999.
Sami Davies, Benjamin Moseley, Sergei Vassilvitskii, and Yuyan Wang. Predictive flows for faster ford-fulkerson. International Conference on Machine Learning, ICML, volume 202, pp. 7231- 7248, 2023.
Michael Dinitz, Sungjin Im. Thomas Lavastida, Benjamin Moseley, and Sergei Vassilvitskii. Faster matchings via learned duals. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems, NeurlPS, pp. 10393-104066 10406, 2021.
Jon C. Ergun, Zhili Feng, Sandeep Silwal, David P. Woodruff, and Samson Zhou, dearning-auggmene k-means clustering. In The Tenth International Conference Learning Representations, ICLR, 2022.
C. Torgeson G. Pass, A. Chowdhury. 500k user session collection. https:/ www. kaggle. datasets/dineshyyldee user- session collection. -00kk 2006.
Xavier Gabaix. Zipf's law for cities: an explanation. The Quarterly journal of economics, 1 114(3): 739-767, 1999.
Tingjian Ge and Stan Zdonik. A skip-list approach for efficiently processing forecasting queries. Proceedings of the VLDB Endowment, 1(1):984-995, 2008
Sreenivas Gollapudi and Debmalya Panigrahi. Online algorithms for rent-or-buy with expert advice. In Proceedings of the 36th International Conference Machine Learning, ICML, pp. 2319-2322 2019.
Goog. Google books n-gram frequency lists.
https://gttubb.cmm// wttp://t///oooooooooooo/oo
- e--oooks-ngrram--rreeeee 2022.
Google. Google ngram viewer.
www.goowww....oggmmmaaaaaettttt. 2012. URL http: 7/books. goo.ks..oooll....m httpll/coooooggg///ttaaeaeeeeeee
Elena Grigorescu, Young- -an Lin, Sandeep Silwal, Maoyuan Song, and Samson Zhou. Learning- augmented algorithms for online linear and semidefinite programming. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems, NeurIPS, 2022.
12

---

Published as a conference paper at ICLR 2025
Chen- Yu Hsu, Piotr Indyk, Dina Katabi, and Ali Vakilian. Learning- besed frequency estimation algorithms. In 7th International Conference on Learning Representations, ICLR, 2019.
Yih-Chun Hu, Adrian Perrig, and David B Johnson. Efficient security mechanisms for routing
protocolsa. Ndss, 2003.
David A Huffmman. A method for the construction of minimum-redundante codes. Proceedings of the IRE, 40(9): 1098-1101, 1952.
Sungjin Im, Ravi Kumar, Mahshid Montazer Qaem, and Manish Purohit. Online knapsack with frequency predictions. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems, NeurIPS, pp. 2733-2743, 2021.
Piotr Indyk, Ali Vakilian, and Yang Yuan. Learning-basee low-rank approximations. In Advances in Neural Information Processing Systems 32; Annual Conference on Neural Information Processing Systems 2019, NeurIPS, pp. 7400-7410, 2019.
Zachary Izzo, Sandeep Silwal, and Samson Zhou. Dimensionality reduction for wasserstein barycen- ter. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS, pp.
2021.
Shaofeng H.-C. Jiang, Erzhi Liu, You Lyu, Zhihao Gavin Tang, and Yubo Zhang. Online facility location with predictions. In The Tenth International Conference on Learning Representations, ICLR, 2022.
Tanqiu Jiang, Yi Li, Honghao Lin, Yisong Ruan, and David P. Woodruff. Leamning-augmenter data stream algorithms. In Sth International Conference on Learning Representations, ICLR, 2020.
Marek Karpinski, Lawrence L. Larmore, and Wojeiech Rytter. Sequential and parallel subquadratic work algorithms for constructing approximately optimal binary search trees. In Proceedings of the Seventh Annual ACM-SIAM Symposium Discrete Algorithms, pp. 36-41, 1996.
Mikhail Khodak, Kareem Amin, Travis Dick, and Sergei Vassilvitskii. auggmneed private algorithms for multiple quantile release. In International Conference on Machine Learning, ICML 2023, pp. 11344-16370 2023.
Misha Khodak, Maria-Florina Balcan, Ameet Talwalkar, and Sergei Vassilvitskii. Learning predic- tions for algorithms with predictions. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems, NeurlPS, 2022.
Donald E. Knuth. Optimum binary search trees. Acta Informatica, 1:14-25, 1971.
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 on Management of Data, SIGMOD Conference, pp. 489-504, 2018
Silvio Lattanzi, Thomas Lavastida, Benjamin Moseley, and Sergei Vassilvitskii. Online scheduling via learned weights. In Proceedings of the 2020 ACM-SIAM Symposium Discrete Algorithms,
SODA, pp. 1859-1877, 2020.
Silvia Lazzardi, Filippo Valle, Andrea Mazzolini, Antonio Scialdone, Michele Caselle, and Matteo Osella. Emergent statistical laws in single-cell transcriptomic data. Physical Review E, 107(4): 044403, 2023.
Yi Li, Honghao Lin, Simin Liu, Ali Vakilian, and David P. Woodruff. Learning the positions in countsketch. In The Fleventh International Conference on Learning Representations, ICLR, 2023.
Honghao Lin, Tian Luo, and David P. Woodruff. Learning augmented binary search trees. In International Conference on Machine Learning, ICML, pp. 13431-13440, 2022.
Jonatan Lindén and Bengt Jonsson. A skiplist-based concurrent priority queue with minimal memory contention. In Principles of Distributed Systems: 17th International Conference, OPODIS 2013, Nice, France, December 16- 18, 2013. Proceedings 17, pp. 206-220. Springer, 2013.
13

---

Published as a conference paper at ICLR 2025
Thodoris Lykouris and Sergei Vassilvitskii. Competitive caching with machine learned advice. J. ACM, 68(4):24:1124-22 2021.
Kurt Mehlhorn. A best possible bound for the weighted path length of binary search trees. SIAM J. Comput., 6(2):235-233 239, 1977.
Michael Mitzenmacher. A model for learned bloom filters and optimizing by sandwiching. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems, NeurIPS, pp. 462-471, 2018.
Michael Mitzenmacher and Sergei Vassilvitskii. Algorithms with predictions. In Tim Roughgarden (ed.), Beyond the Worst-Case Analysis of Algorithmx, pp. 646-662. Cambridge University Press, 2020.
Thy Dinh Nguyen, Anamay Chaturvedi, and Huy L. Nguyen. Improved learning -uggmnnte al- gorithmms for k-means and k-medians clustering. In The Eleventh International Conference on Learning Representations, ICLR, 2023.
Jonathan J Pittard and Alan I. Tharp. Simplified self-adapting skip lists. In International Conference on Intelligent Data Engineering and Automated Learning, pp- 126-136. Springer, 2010.
William Pugh. Concurrent maintenance of skip lists. University of Maryland at College Park, 1990a.
William Pugh. Skip lists: A probabilistic alternative to balanced trees.
Commun. ACM, (3(6): jun 1990b. ISSN 0001-0782. doi: 10. 00.1145/79973788777 URL https://doi org
10. N0..145/78991 78977
Manish Purohit, Zoya Svitkina, and Ravi Kumar. Improving online algorithms via ML predictions. In Advances Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS, pp. 9684 9693, 2018.
Agnar Sandmo. The principal problem political economy: income distribution in the history of economic thought. In Handbook of income distribution, volume 2, pp. 3-65. Elsevier, 2015.
Ziv Scully, Isaac Grosof, and Michael Mitzenmacher. Uniform bounds for scheduling with job size estimates. In 13th Innovations in Theoretical Computer Science Conference, ITCS, pp.
114:1-114:30, 2022.
Claude Elwood Shannon. A mathematical theory of communication. ACM SIGMOBILE mobile computing and communications review, 5(1):3-55, 2001.
Nir Shavit and Itay Lotan. Skiplist-based concurrent priority queues. In Proceedings 14th In- ternational Parallel and Distributed Processing Symposium. IPDPS 2000, pp. 263-268. IEEE, 2000.
Yongho Shin, Changyeol Lee, Gukryeol Lee, and Hyung-Chan An. Improved learning- augmented algorithms for the multi-option ski rental problem via best posssibe competitive analysis. In International Conference Machine Learning, ICML, 31539-31561, 2023.
Stan. The stanford 3d scanning repository, 1996. URL htttttttt/gggahhcss... wtw.aiiiissstttaoo..ddd..d
data/ esta/3Dscanree)
Nikolay Vitanov and Marcel Ausloos. Test of two hypotheses explaining the size of populations in system of cities. Journal of Applied Statistics, 42(12):268-2693, 2015.
Ding Wang and Ping Wang. On the implications of zipf's law in passwords. Computer Security ESORICS 2016 2t European Symposhum on Research in Computer Security, Proceedings, Part volume 9878 of Lecture Notes in Computer Science, pp. 111-131. Springer, 2016.
Shufan Wang. Jian Li, and Shiqiang Wang. Online algorithms for multi-shop ski rental with machine learned advice. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems, NewrIPS, 2020.
14

---

Published as a conference paper at ICLR 2025
Alexander Wei and Fred Zhang. Optimal robustness robustness-consiiteend consistency trade-offs for Boaming-augmente online algorithms. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems, NeurlPS, 2020.
Frances Yao. Speed-up in dynamic programming. SIAM Journal on Algebraic Discrete Methods, 3 (4):532-540, 1982.
Ali Zeynali. Shahin Kamali, and Mohammad Hajiesmaili. Robust learning-augmented dictionaries. CoRR, abs/2402.0965 2024.
Xuanhe Zhou, Cheng Chen, Kunyi Li, Bingsheng He. Mian Lu, Qiaosheng Liu, Wei Huang, Guoliang Li, Zhao Zheng, and Yuqiang Chen. Febench: A benchmark for real-time relational data feature extraction. Proceedings of the VLDB Endowment, 16(12):3597-3609 2023.
A ADDITIONAL RELATED WORKS
In this section, we discuss number of related works in addition those mentioned in Section This paper builds upon the increasing body of research in learning-augmented algorithms, data-driven augmented algorithms have been
algorithms, and algorithms with predictions. For example, applied to a number of problems in the online setting, where the input arrives sequentially and the goal is to achieve algorithmic performance competitive with the best solution in hindsight, i.e., an algorithm that has the complete input on hand. Among the applications in the online model, learning- augmented algorithms have been developed for ski rental problem and job scheduling (Purohit et al., 2018), caching (Lykouris& & Vassilvitskii 202 )) and matching (Antoniadis et al. 2023b). Learning- augmented algorithms have also been used to improve the performance of specific data structures such Bloom filters (Mitzenmacher 2018), index structures (Kraska et al., 2018), CountMin and CountSketch (Hsu et al. 2019). Specifically, (Kraska et al.) (2018) proposes substituting B-Trees (or other index structures) with trained models for querying databases. In their approach, rather than traversing the B-Trre to locate a record, they use a neural network to directly identify its position. Our work differs in that we retain the desired data structures, i.e., skip lists and kd trees, and focus on optimizing their structures to enable faster queries, which allows us to continue supporting standard operations specific to the data structures such as traversal, order statistics, merging, and joining. among others. Our work uses the frequency estimation oracle trained in (Hsu et al. 2019) on the AOL search query dataset and the CAIDA IP traffic monitoring dataset.
Perhaps the works most closely related to ours in the area of learning- augmented algorithms are those of (Lin et al. 2022, Cao al. 2023 Zeynali et al., 2024). Lin et al. (2022) noted that traditional theory on statically optimal binary search trees (Knuth) 1971 Mehlhorn 1977 is no longer applicable in dynamic settings, where elements are added incrementally over time. Hence, they developed dearning-augmentte binary search trees (BSTs) and showed that their expected search time is near- optimal. Cao et al.] (2023) then extended these techniques to general search trees, allowing for nodes with more than two children. [Cao et al. (2023) also studied the setting where the predictions may be updated, while ultimately still utilizing a data structure that requires rebalancing as data is dynamically changing- Zeynali et al. (2024) also consider the performance of learning- mamninggauumenteed skip lists that are robust to erroneous predictions; we claborate more on the differences from Zeynali et al. (2024) in SSecction I.1] We also note that none of these works consider KD trees at all, which is an important data structure with applications in computer vision and computational geometry, thus forming basis of our work. For more comprehensive source of related works in learning- wwwniinng-uggnnttee algorithms, see https:: al /oaloor/loooooooooo thms - wi th prrdd consss....... qi witth. thub. 1o/]
Beyond the context of shamiig--uggmenn of data structures that are optimal for their inputs. For example, while standard binary search trees use O (logn) query time, optimal static trees can be constructed using dynamic programming or efficient greedy algorithms (Mchlhorn] 1977. Yao 1982 Karpinski et al. 1996). given access frequencies. However, the computational cost of these methods often exceeds the cost of directly querying the tree. As result, key objective is to construct tree whose cost is within constant factor of the entropy of the data. Several approaches have achieved this either for worst-case data (Mehlhorn 1977) or when the input follows particular distributions (Allen & Munro 1978).
algorithms, there a large body of works that study design
More recent works have considered using results from learning theory to estimate the query frequen- cies, rather than assuming explicit access to their values. For example, Cayton & Dasgupta 2007
15

---

Published as a conference paper at ICLR 2025
studied how to obtain such oracle for generalization bounds the context of learning theory, analyzing the number of samples from an underlying distribution necessary to produce oracle with small enror rate. On the other hand, Ailon et al. (201T) studied algorithms for sorting and clustering that can improve their expected performance given access to multiple instances sampled from fixed distribution. Although the high-level goal of improving algorithmic performance using auxiliary information is the same as ours, the specifics of the paper seem quite different than ours, as the paper focuses on techniques for sorting and clustering. Similarly, Ciriani et al. (2002) considers self self-adjusting data structures, including skip lists, which can dynamically change as the sequence of queries arrive. However, their methods are catered specifically to the setting where there is access to the queries, whereas our data structures must constructed without such access and must therefore be able to handle erroneous predictions. Finally, we remark that utilizing techniques for learning theory, there are standard results about the lcarnability of oracles for the purposes of 2022
data structures. In particular, they study
algorithms (Ergun et al.
Izzo et al. 2021).
B MISSING PROOFS FROM SECTION 2
In this section, we give the missing proofs from Section
B.1 EXPECTED SEARCH TIME
We first show that each item promoted to a higher level with probability at least 1.2
Lemma B.1. For each item i E [n] at level e, the probability that i is in level E + is at least
Proof. Note that if 2 2f n, then will be placed in level f + 1. Otherwise, conditioned on the item i E [n|] being at level E, then Algorithm places i at level + 1 with probability 1. Thus, the probability that is in level f + 1 is at least 2-
We next upper bound the expected search time for any item at any fixed level, where the randomness over the construction of the skip list.
Lemma B.2. In expectation, the search time for item E [n]] at level E is at most 2.
Proof. Suppose item i E [n] is in level f. Let S C [2]] be the subset of items in level & that are less than i. Note that by Lemma B.1 each item of S is promoted to level + 1 with probability at least 2. Thus, the search time for item at level is if and only if the previous items in St were all promoted, which can only happen with probability at most 2 Hence, the expected search time for
item E [n] at level is at most
1
8
E [
<
+
2
22
2
We now show that each item i must be contained at some level depending on the predicted frequency p: of the item.
Lemma B.3. Each item i is included in level max (0, 1 [log(np;)]).
Proof. First, observe that all items are inserted at level 0. Next, note that Algorithm I inserts item i into level l if p. 2-11 or cquivalently log(np:) f - 1. Thus, each item i is included in level max (0, 1 + (log(np)).).
We next analyze the expected search time for each item
Lemma B.4. Suppose the toiul number of levels is most C + log n. for some constant C > 0. Then the expected search time for item is at most 2C + 2 min log J. log
16

---

Proof. For each level f, let ny be the number of items i E [ri] that are deterministically promoted to Published as a conference paper at ICLR 2025
Proof. By Lemma B.3 item i is included in level max (0, 1 + [log(np)]]). By Lemma B.2 the expected search time at each level is at most 2. Thus, in expectation, the total search time is at most 2(C + log max (0, [log(np)))) 2C 2 min (log p; logn).
Finally, we analyze the expected search time across the true probability distribution fs.
Lemma B.5. Suppose the total number of levels is at most C + log n. for some constant C' > For
each i E [n), let Ja be the proportion of queries item Then the expected search time at most 2C + 2 Ji min (log L logn). p;
Proof. For each query, the probability that the query item ft. Conditioned on the total number of levels being at most C + log n. then by [ymmmaa B.4 the expected search time for item is at most 2C' 2 max llog 1 logn). Thus, the expected search time at most
log 1 P1
2C +... fn) + 2f min
logn .... + 2 min log log
pn.
_ log
= + fo min
log
pi
We now show that with high probability, the total number of levels in the skip list is at most O (log n)).
Lemma B.6. With probability at least 0.99, the total number of levels in the skip list is at most 10 + log n.
exactly level f, i.e., ps E 2 2²) Note that for each fixed E [n]. the highest level remains is geometric random variable with parameter 1, beyond the highest level at which it is deterministically placed. This is because the item is promoted to each higher level with probability 1. Hence with probability 1 1 i is not placed at least & levels above its highest deterministic placement. Therefore, the probability that an item at level i is placed at level 10 + log n is at most T024n Since no fixed i will have predicted frequency more than then no item will be deterministically placed at level 2 + log n. Hence by union bound over all é [2 + log ni], the probability that an item is placed at level 10 + log n is at most
2f
2+Hog
The 2
Σ
1024n
On the other hand, we have Di Pi 1, that
2+log n
ng 2 < 2n.
Therefore, with probability at least 0.99, the total number of levels in the skip list is at most 10 logn. n.
Thus, putting together Lemma B.5
and and Lemma B.6 we get:
Theorem 2.1. For euch E [n), let Jú and pe be the proportion of true and predicted queries to item i. Then with probability at least 0.99 over the randomness of the construction of the skip list, the
1 expected search time over the choice of queries at most 20 + 2 2 Si: -min (log p. log n
17

---

Published as a conference paper at ICLR 2025
B.2 NEAR-OPTIMALE
We first recall the construction of a Huffman code, type of variable-length code that is often used for data compression. The encoding for Huffman is known to be an optimal prefix code and can be represented by a binary tree, which we call the Huffman tree (Huffman 1952).
To construct a Huffman code, we first create a min-heap priority queue that initially contains all the leaf nodes sorted by their frequencies, so that the least frequent items have the highest priority. The algorithm then iteratively removes the two nodes with the lowest frequencies from the priority queue, which become the left and right children of new internal node that is created to represent the sum of the frequencies of the two nodes. This internal node is then added back to the priority queue. This process is continued until there only remains a single node left in the priority queue, which is then the root of the Huffman tree.
binary code is then assigned to the paths from the root to each leaf node in the Huffman tree, so that each movement along left edge in the tree corresponds appending 0 to the codeword, and each movement along right edge in the tree corresponds to appending a 1 to the codeword. Thus, the resulting binary code for each item is the path from the root to the leaf node corresponding the item.
Huffman coding is a type of byyboll coding, where each individual item is separately encoded, as opposed to alternatives such as run-length encoding. It is known that Huffman coding optimal among symbol-by-symbor coding with a known input probability distribution (Huffman 1952) and moreover, by Shannon's 's source coding theorem, that the entropy of the probability distribution is an upper bound on the expected length of codeword of a symbol-by-symbo coding:
Given a random variable Theorem B.7 (Shannon's source coding theorem). (Shamnon) 2001)
X E [n] so that X = with probability fú. let L(X) denote the length of the codeword assigned to X by Huffman code. Then E [L(x)) (f)), where H(f) is the entropy of f.
We now prove our lower bound on the expected search time of an item drawn from a probability distribution f.
Theorem 2.2. Given a random variable X E [n] so that X = i with probability fú, let T(X) denote the search time for X in a skip list. Then E [T)) 2 (f)), where (f) is the entropy of f.
Proof. Let L. be a skip list. We build a symbol-by encoding using the search process in L. We begin at the top level. At each step, we either terminate, move to the next item at the current level, move down a lower level. Similar to the Huffman coding, append a to the codeword when we move down to a lower level, and we append a I to the codeword when we move to the next item at the current level. Now, the search time for an item of the codeword of x in the
L corresponds to the length encoding. By Theorem B.7 and the optimality of Huffman codes among symbol-by-symbol encodings, we have that E (x)) HJ), where / the probability distribution vector of x.
B.3 ZIPFIAN DISTRIBUTION
In this section, we briefly describe the implications of our data structure to Zipfian distributions.
We first recall the following entropy upper bound for a probability distribution with support at most n.
Theorem B.8. (Cover 1999) Let f be a probubility distribution on a support of size [n). Then H() < logn.
We can then upper bound the entropy of a probability vector that satisfies a Zipfian distribution with parameter
Lemma 2.3. Let s, > be fixed constants and let f be a frequency vector such that f. = t for all
i [n]. Ifs 1, then H()) O (1) and otherwise if < 1, then H()) & log
18

---

Published as a conference paper at ICLR 2025
Proof. Since f is probability distribution the support of size [n], then by Theorem B.8 we have
that (f) & log n. Thus, it remains to consider the case where 1. Since S 1, we have
h())
=
log
logí
S 8 1
t=1
logí Note that there exists integer y 0 such that for i >> 7, have
& (0+1175
Since >> 1,
then s+1 > and thus
Σ
i 11/2 (1).
1x+1)/2
i=1
Hence,
log i
log i (1).
h(f) S
+
s
1
1
By [Theorem 2.1 and Lemma 2.3] we have the following statement about the performance of our Boaming-augmente skip list set of search queries that follows Zipfian distribution.
Corollary 2.4. With high probability, the expected search time on a set of queries that follows Zipfian distribution with exponent is most O (1) for s and O (log n) for S
B.4 NOISY ROBUSTNESS
In this section, we show that our Barning-auugmene skip list construction is robust to somewhat inaccurate oracles. Let be the true-scaled frequency vector so that for each E [n²], fi is the probability that random query corresponds to i. Let be the predicted frequency vector, so that for each E [n]. pa is the predicted probability that a random query corresponds to i. For a, 3 E (0, 1), we call an oracle (o,, nnisy if for all E [n). we have p; a St -3
skip list with set of (a,, B)- noisy predictions has performance that matches that of a learning -ugiiintd learned with a perfect oracle, up to an additive constant.
Lemma 2.5. Let be a constant and 3 < im A starning-augmente
Proof. Suppose the total number levels at most + log n for some constant that this occurs with high probability for learning-augmentee skip list with a set of (a, -n-iss predictions. For each i E [n], let f; be the proportion of queries to item and let p; be the predicted proportion of queries to item i. By Lemma B.3 the expected search time at most
0. Note
Σ fo min log 1 ps
2C +
logn
Since the oracle is (o,, nnoss then have p: 2 Se for all [n].
We first note that in the expected search time for is proportional to min (log J. log
logn). m). Thus, for
expected search time for item i, it suffices to assume ft in for all i.
Observe that for ft 1 and B then p;
B implies
2n
pi f. B
fi.
Hence, we have 1 < pi a J. so that the expected search time for item i most
2
2
2C + min
log J.
+ log log a
19

---

Published as a conference paper at ICLR 2025
Therefore, the expected search time is at most
n 2C + Σ fr
min
log
log m pi
20 220+2 Σ
[i s..min
log Ji
log
S: log
1=1
2 + fo min log
2C + 2log
llo
fi
Since the perfect oracle would achieve runtime + 2 fo min (log J. log then follows that earning-augmenter skip list with set of (a, nnoss predictions has performance that matches that of a dearning-augmenter learned with a perfect oracle, up to an additive constant.
C MISSING PROOFS FROM SECTION 3
First, we will show that the expected depth of a given query depends on the probability of that query, and that high frequency queries must be found close to the root of our tree.
Lemma C.I. Suppose [A] is the space of possible imput points and queries. Let N = A and pi be the probability that random query is made to E [N], given the natural mapping between [N] and [A]á. Then the level at which resides in the tree is at most O (log p4.
1
Proof. First, consider only the high-frequency query points and data points for which we base our construction off.
In constructing the Barming-augmenter KD tree, we balance the contents of the children nodes such that a query to that node has probability of of belonging to each of the children. Therefore, at depth of d, the probability of belonging to either child is 1 zz In particular, query with probability
P: at depth d must satisfy p;
24 Thus, we have that the depth of i is O [log 1 p. as desired.
Since the lowest probability of high-frequency data point is n2, this tree must have depth of at most log n.
Now, consider low-frequency data point, which add to the bottom of the tree. By construction, the learned point of our tree has depth at most 2 log n. Then, when inserting the additional data points
balanced KD tree, can accumulate most an additional depth of log n. Note, p / implies
log n < log J. Thus, this low-frequency data point will have depth of most 3 log n O (log] + as desired.
Similarly, if i is not data point and is low frequency, achieve the same bound of O (log 4 In this case, we simply terminate at a leaf node and determine that the desired query is not in the dataset.
p
In summary, any query which has high frequency can be found in O (log 1 p) time. Low-frequency
data points can similarly be found in O (log) 1 ) time, and low frequency queries can be determined to not exist in O [lgg] } time.
Lemma C.2. Suppose [A] is the space of possible input points and queries. Let N = Ad and pı be the probability that random query is made to E [NT, given the natural mapping between [N] and [ Then the level at which resides in the tree is at most O (log n).
Proof. This follows directly from the analysis in Llemma
Now, have demonstrated the the depth of given query point is bounded by both O (logn) and O log 2. p. Using this fact, we will now show that the expected query time of our algorithm is bounded by both the entropy of the dataset (p) in addition to log n.
We now analyze the performance of our Carning-augmentt KD tree.
20

---

Published as a conference paper at ICLR 2025
Theorem 3.1. Suppose [A] is the space of possible input points and queries. Let N A and Pj be the probability that a random query is made to i [N. given the natural mapping between [N] and [A]. Let p (P1..... .NN) E IN be the probability vector and (p) be its entropy. Then given a set of n. input points, the expected query time for the tree created by Algorithm is
O (min( H (p), n)).
Proof. Following C.I the points in [ with non-negligible probability pi 2 n. are guaranteed to exist in the Baaming-auggmett
KD tree T with depth most log 1 For points in the dataset [ni] P; with negligible probability, they exist the tree and have depth in (log n.). For all other points not contained in the KD tree, the query will terminate at a depth of O (log n.)
For any point in [A]a, the depth that query to i will terminate in the tree is
1
Depth(i) O min log ps
log n
(1)
Then, the expected search time T is the expected depth of given point,
E pc-O O (min llog - log )) O (min(H()) (2)
E [TT Σ p.. p.-Depth(i)-
pi
SE|A
vE[A]2
Thus, we have shown that the expected query time of our algorithm is bounded by the entropy of the dataset. In particular, when the dataset has highly skew distribution, ((p) can be far less than log
C.I NEAR-OPTIMALITY
Near-optimality of our learning-and auggmeted KD) trees uses a similar argument to the proof of the near-optimality of our maaning-augmment skip lists. particular, we again utilize Shannon's source coding theorem from Theorem B.7 We then have the following:
Theorem 3.2. Given random variable X E [n] so that X = with probability fo let D(X) denote the depth for X in a learning-and srrmiing-uggmeted KD tree. Then E [D(X) 2 H(f), where H(f) is the entropy of f-
Proof. In a dearnnng-uugmentn KD tree, the search path to an element can be encoded as a codeword, with entries indicating whether the lower upper branch is taken each node traversal. Moreover, the length of this codeword in the symbol-by-symbol encoding corresponds to the depth of element Then, by Theorem B.7 and the optimality of Huffman codes in symbol-by-symboo encodings, we have that E [X] Z H()), as desired.
C.2 NOISY ROBUSTNESS
In this section, we analyze the performance the leaming-augmented KD tree under noisy data conditions.
Previously, we analyzed the performance of the learning aaghinggegmmeteed KD-tree assuming access to a perfect prediction oracle.
Now we analyze the performance with a noisy oracle. That is, for each i E [A] there a true-scaled frequency f; that the point will be queried and that p; is prediction made by the noisy oracle.
First, we analyze the multiplicative robustness of the algorithm. In this case, the oracle predicts fa up to some multiplicative constant R such that f. = ap.
Lemma C.3. Suppose [A]d is the space of possible input points and queries. Let N = Ad and pa be the probability that random query is made to E [N] during tree construction. Suppose during runtime that the true probability of querying i is J or P; for some E BR+, Then the level at which
i resides in the tree is at most O log T. + log // 1
21

---

Published as a conference paper at ICLR 2025
Proof. If a 1, this is immediate. If this is the case, in construction expected i to be queried more often than it actually is, so our construction placed the point higher in the tree than is necessary. Thus, the depth is at most the previously shown 1
P. -
Now, suppose 1. In this case, we must have placed deeper in the tree than we should have, as our construction frequency is less than the true query frequency. Then, as in C.I the depth of is O (log] 1 Now, we have that
(10s o (oog - Pi
= (108g) (+)=0 =0 J, + log
=0
Having shown multiplicative robustness of the learning-augmente KD tree we next analyze the aeditive-multillicatie robustness of the method. An oracle is (a, nnniss if the prediction satisfies 2 a fi - 3 for constants a, 3 E (0, 1).
Lemma 3.4. Let a be a constant and let 3 S 2. Then the query time for our KD tree with (a, 3)-noisy prediction matches the performance of Dearning-uugmentn KD tree constructed using a perfect oracle up to an additive constant.
Proof. From Lemma C.1 we have that the depth of point in the aanniig KD tree is O
1 where p; is the predicted frequency of i. Pi
Then, with an (a, -noosy oracle we have that the prediction is bounded from below by p: 2 a fi 8. As a point in the tree, we can assume that the scaled frequency f. has a lower-bound
1/n².
By choosing B S a/n"² we ensure that the predicted is always nonnegative. Then, let 3 = a a/2n²
pt 2 a -
32 as
(4)
2n 2
Then, applying Lemma
C.I Igaain
2
+ log 1 x
O log = 0 log afi
=
log fi
(5)
D ADDITIONAL EMPIRICAL EVALUATIONS
D.1 SKIP LISTS
In this section, we perform empirical evaluations comparing the performance of our learning- augmented skip list that of traditional skip lists, on both synthetic and real-world datasets. Firstly, we compare the performance of traditional skip lists with our bearning-augment synthetically generated data following Zipfian distributions. The proposed sorming-augmentee skip lists are evaluated empirically with both synthetic datasets and real-world internet flow datasets from the Center for Applied Internet Data Analysis (CAIDA) and AOL. the synthetic datasets, a diverse range of element distributions, which are characterized by the skewness of the datasets, are evaluated to assess the effectiveness the learning augmentation. In the CAIDA datasets, the a factor is calculated reflect the skewness of the data distribution.
skip lists on
The metrics of performance evaluations include insertion time and query time, representing the total time it takes to insert all elements in the query stream and the time it takes to find all elements in the query stream using the data structure, respectively.
The computer used for benchmarking is a Lenovo Thinkpad P15 with intel i7- 11800H@2.30 2,3GH,, 64GB RAM, and ITB of Solid State Drive. The tests were conducted a Ubuntu 22.04.3 LTS OS. GNOME version 42.9.
22

---

Published as a conference paper at ICLR 2025
D.1.1 SYNTHETIC DATASETS
In the synthetic datasets, both the classic and augmented skip lists are tested against different element counts and a values. In terms of the distribution of the synthetic datasets, the uniform distribution and Zipfian distribution of a between 1.01 and with query counts up to million are evaluated. is worth noting that the number of unique element queries could vary for the same query count at different a values in the Zipfian distribution, which may affect the insertion time.
Table 1: Speed up factor of augmented skip list over classic skip list under different synthetic distributions
Query size of synthetic data (unit: thousand)
Distribution
0.5
10
100
500
1000
1500
2000
2500)
3000
3500
4000
Average
uniform
3.02
0.84
1.01
1.05
1.11
1.14
1.17
1.21
1.22
1.42
1.4
1.33
0=1.01
3,63
2.6
1,04
1.24
1.21
1.2
1.14
1.3
1.18
1.53
a=1.25
3.28
3.74
5.87
2.89
2.47
3.21
2.95
3.34
3.55
3.16
3.12
3.42
a=1.5
2.42
8.97
6.93
6.54
7.99
5.83
4.65
3.8
4.92
5.34
5.93
5.76
0=1.75
12.43
5.76
9.78
6.76
7.13
7.31
7.09
6.63
5.07
6.98
7.76
o=2
8.19
2,5
5,56
10.1
4.47
3.91
7.26
5.33
9.29
7.65
5.55
6.35
Table 2: Node count for each distribution configuration in the million dataset
Unique node count
1.01
2886467
259892
1.75
8386
2
2796
T the speed-up factor, defined the time taken by the augmented skip list over the classic skip list for the same query stream. We can observe a progressive improvement the performance of augmented skip lists as the dataset skewness increases. It also suggests that augmented skip list will perform at least as good as the traditional skip list and will outperform a traditional skip list by a factor of up times depending on the skewness of the datasets.
showws that the insertion time decreases with more skewed datasets for the same size of the query stream. This attributed to the reduced number of nodes in the datasets, as shown in Table 2
The query time of augmented skip lists is also reduced greatly compared to the classic skip lists as shown in Figure 7
In addition, we conduct experiments to compare the performance of standard binary search trees and standard skip lists. In particular, we generate datasets of size n (5000, 10000, 15000, 20000, 25000, 30000, 35000, 40000, 45000, 50000, }. For each fixed value of n, each element of the dataset generated uniformly at random [2m], i.e., uniformly at random from 21, 2m.. data structures will perform similar to oblivious data structures. We measure the construction time of the data structures, based on the input dataset. Our results demonstrate that as expected, skip lists perform significantly better than balanced binary search trees across values of n. due to the latter's necessity of constantly rebalancing the data structure. In fact, skip lists performed almost 4x better than BSTs in some cases, e.g., n 20000. We illustrate our results in Figure g
Because the dataset is generated uniformly at random, then
1.1.2 AOL DATASET
The AOL dataset (G. Pass 2006) features around 20M web queries collected from 650k users over three months. The distribution of the queries is shown in Figure 9 The AOL dataset is a less skewed dataset than CAIDA with an alpha value of 0.75.
Figure 9 shows the distribution of the AOL queries with estimated alpha value of 0.75. The AOL dataset resembles more to a slightly skewed uniform distribution with very few highly frequent items, which accounts for lower improvement as in the case of AOL shown in Figure 10 The total number of queries for items with higher than 1000 frequency accounts for only 5% of the total number of queries for the AOL datasets. The learning- skip list still outperforms the traditional skip
23

---

Published as a conference paper at ICLR 2025
Insertion Time of Vanilla Learned Skigllzt (Uriform))
Venilla Inser Time
Ven Insertion Time
W1t00t00 Time Geconds
Learned Insertion Time
Insortton Trme 1
Number Values (Uni 1M90
Number Values (Unit IMW0
(a) Uniform Distribution
(b) = 1.25
Time Varilla Learned Skiolist (ew1.5)
Skiplst (ow2)
Insertionn Thme (sconds)
15s0ttio Tarl Sppcorni
(Ulni: 1890)
00000
1890)
a=)a==15 1.5
d
Figure 6: Insertion time for synthetic datasets with a uniform distribution and under different a values of the Zipfian distribution for both classic and augmented skip lists. This figure illustrates the insertion time on the synthetic data for both the uniform distribution and the Zipfian distribution at different values. Generally, higher skewness of the datasets results in less insertion time when using the augmented structure. The decrease in insertion time is proportional to the increase in the a value, as higher a value leads reduction in the number of unique nodes, illustrated in Table 2
list on this slightly skewed dataset. This result is also in line with the results from the synthetic data shown in Table Twbree where lower alpha values have resulted in a lower speedup factor.
D.2 KD TREES
For KD Trees, first describe our methodology for evaluating our data structure on synthetic data. We then describe our empirical evaluations on real-worke datasets.
The computer used for KD tree benchmarking is a desktop machine with an Intel Core 19-14900KF@ 3200MHz, with 64GB RAM, and 2TB of Solid State Drive. The tests were conducted in Windows 10 Enterprise, version 10.0. 19045 Build 19045.
D.2.1 SYNTHETIC DATASETS WITH PERFECT KNOWLEDGE
First, consider dataset with Zipfian distribution. In order construct this dataset, first select n unique data points in [A] uniformly. We then generate Zipfian frequencies, fi P (1+6)= and randomly pair the frequencies to the data points to serve as both the construction and query frequencies. We construct either traditional KD tree, or our learned KD tree on this dataset. Then, evaluate the performance of querying by sampling the known datapoints with probabilities given by their Zipfian probabilities.
In Table 3 we use Zipfian distribution with parameters 1 and b = 2.7. This data demonstrates that our method outperforms a traditional KD tree, given that the Zipfian distribution has the given parameters. In Figure I] vary the Zipfian parameters of our dataset. As expected, our learned KD tree performance increases with how skew the distribution is. Moreover, we find that our method outperforms the traditional KD tree all tested Zipfian distributions.
24

---

Published as a conference paper at ICLR 2025
Time of Vonilo Leained Sktalez (Uniform)
Query
Quirry Time (scoind))
202
Venila Quary Time
Learned Query
Quary Tune Cocconene
Van Query
Number Values (Unit 1M90
Number Values (Unit IMW0
(a) Uniform Distribution
(b) = 1.25
Query Vanilla Skipliz (ow1.5
Query
Sciplist (1
Quary
Quary
Quary Time (acoodo
Query
QuoT Timee D0oo Roconng 000
Query
Tame S0O0 Taboc
(Uiii: 1890)
00000
Toae S0O0 TIDOOO
(Ulii: 1890)
=
ad)==2
Figure 7: Query time for synthetic datasets with a uniform distribution and under different a values of the Zipfian distribution for both classic and augmented skip lists. This figure compares the query time of the classic and augmented skip lists for both uniform distribution and Zipfian distribution at various values. Similar to insertion, the query time is significantly reduced under different query sizes with the implemented augmentation. The performance enhancement is especially pronounced for the high skewness of the dataset.
type traditional
dim
const time
avg query(s)
avg query depth
I
1.267E-01
2.645E-05
1.330E+01
learned
1
2.387E-01
2.181E-05
1,086E+01
traditional
2
1.266E-01
2.709E-05
1.341E+01
learned
2
3.232E-01
2.221E-05
1,086E+01
traditional
3
1.253E-01
2.700E-05
1,334E+01
learned
3
3.981E-01
2.264E-05
0.097
traditional
4
1.185E-01
2.724E-05
1.336E+01
learned
4
4.549E-01
2.256E-05
1,089E+01
traditional
5
1.300E-01
2.743E-05
1.336E+01
learned
5
5.286E-01
2.296E-05
1.096E+01
traditional
10
1.277E-01
2.896E-05
1.340E+01
learned
10
8.642E-01
2.403E-05
1,09
traditional
20
1.295E-01
3.121E-05
1.344E+01
learned
20
1.543E+00
2.564E-05
1,085E+01
traditional
40
1.361E-01
3.558E-05
1.340E+01
learned
40
2.911E+00
2.902E-05
1.078E+01
Table 3: We construct and query KD trees with our method and with a traditional KD tree on synthetic datasets of various dimensionality. We construct tree 10k points, and query IM times. We find that, independent of the data dimensionality, our method produces lower average query depths.
D.2.2 SYNTHETIC DATASETS WITH NOISY KNOWLEDGE
In reality, it is rarely the case that have perfect knowledge in constructing model. In order to evaluate the performance of our method on noisy data, we create a synthetic dataset with noisy training information.
As before, we generate world points and assign them Zipfian weights. When constructing the tree, each weight us updated to be Mp: A, where M and A are drawn from uniform distributions. This
25

---

Published as a conference paper at ICLR 2025
Comparison of BST and Skiplist
0.8
BST
Skiplist
0.7
0.6
0.5
Values
0.4
0.3
0.2
0.1
0.0
5000
10000
15000 20000 25000 30000 35000 40000 45000 50000
Labels
Figure 8: BST vs skip list construction times
1e4 AOL query frequency distribution
AQL query frequency distribution
3.0
10a
2.5
Requency
2.0
Frequency
1.5
1.0
0.5
0.0
10a
10
10
10a
10a
10
102
10
10*
Index
Index
(a) AOL data distribution
(b) Zipfian (a 0.75)
Figure 9: AOL datasets distribution characterization. This figure illustrates how the value of 0.75 obtained for the AOL dataset. The AOL dataset shows much smaller a value compared to the CAIDA dataset so AOL almost resembles uniform distribution despite very few high-frequency nodes. This also explains why the performance of the augmented skip list is close to the classic implementation.
new noisy distribution is normalized order to form valid probability distribution. Then, points are queried many times with their ground truth Zipfian probabilities. For the fixed Zipfian distribution with a = 5, b 2, we plot the effects of different ranges M and to demonstrate the effect noise has in Fig. 4 This figure demonstrates that, even with moderate amounts of noise, our method still outperforms a traditional KD tree. Moreover, our method still remains on par with the traditional KD tree when significant noise is present, due robustness guarantees.
26

---

Published as a conference paper at ICLR 2025
Insertion Time Vanilla and Learned Skiplist
Query Time Vanilia and Learned Skiplist
Vanilla Insertion
Vanilla
decords
101
Quary (vecones)
10:
20-1
10a
10
12.4500
4.226
1.2200
0aw011
00747
auzed
0.177
0.0027
0,8533
eoous
0.0081
33000
27708
2.07a
1802
L2278
0,00
0015
0.2283
0.1700
0.1144
Number
(uni: 100g
Frequent Quary Count cint. Myy
(a) Insert time on AOL
(b) Query time AOL
Figure Insertion and query time on AOL classic and augmented skip lists
Traditional KD
Learning Aucmented KD-Tree
0.c
00
10.0
Figure 11: We generate datasets of 22 points 3-dimentional space, with frequencies given by Zipfian distribution with parameters a, b. We then query the tree 2 times, with point queries selected by the same Zipfian distribution. We repeat this process 32 times, and report the median of the average query depth across all runs. We find that our method is able outperform the traditional KD tree method across all tested Zipfian distributions. Moreover, our performances increases as increases and b decreases, making the Zipfian distribution most skew. Notably, when 0, all points have uniform weights, at which point our method performs equivalently to the traditional KD tree.
D.2.3 WORLD DATASETS
In addition evaluating our method on synthetic data, also evaluate results real-world data.
First, we consider n-grams in various languages. test on a pre-processed subset of the Google N-Gram dataset (Goog Google 2012).
In order evaluate our method, we convert an n-gram to vector 7 with each entry indexing the words in the n-gram. We construct the barming-augmentee and traditional KD trees, and show the performance lookup with queries weighted by their ground truth frequency in Table In all cases, find that the -agmentd KD tree outperforms the traditional KD tree at average lookup depth.
Additionally, test method on dataset of neuron activity, provided by Aitchison al. (2014), which has shown to be Zipfian. This dataset consists vectors in {0, 1] indicating which of 30 cells fire at agiven time. As in their work, we bin in 20ms increments when constructing these
27

---

Published as a conference paper at ICLR 2025
vectors. We similarly ignore the time of observations, and build our learning -aggmented KD tree with frequencies given by the rate of appearance of vectors. We find that, when querying with probabilities equivalent to the training distribution, traditional KD tree has average query depth of 23.7. Using our aarming-augmented KD tree, however, we are able to achieve an average query depth of 14.9.
Dataset
Traditional Avg Query Depth
Learned Avg Query Depth
2grams wetiimssschieeeeeee simplliied
14,8388
12.1144
2grams_englisn english-fiction
14.216
11.6232
2grams_englhe
14.7253
11.6647
2grams french
14.5123
11.5999
2grams german
13.9473
11.7066
2grams_hebree
13.5322
12.0007
2grams_italian ittalian
14.5231
11,0027
2grams russian
14.5769
11,8082
2grams-spanish
15.1763
11,6582
3grams_ching whineseesimpllle
12.343
11.4211
3 5grams english-fiction
13.5264
11.3183
3 3grams english
15.002
11.3632
frrnch
14.9235
11,3529
ggrams-german
14.129
11.4285
3grams hebrew
13,0887
9,6663
3grams italian
13.3839
11,3616
3grams_russ jussiin
15.2475
11,1415
3grams 3grams-spanish
15,3855
11 113635
11.5823
9.8 9.8661
4grams_enggen .english-fiction
12.5141
9.7589
4grams_engge Lenglihh
13.4464
9.7793
4grams french
12.6383
9.835
4grams german
12.3911
9.8765
4grams_hebrew
9.3198
7,4493
4grams_itan italian
11.3456
9.6834
Agrams_russian
13.2406
9 9.6274
gegrams-spaniin
12.9712
9,8191
seraams_hinnee.seee chinese.simplified
11.5373
9,8982
5grams english-fiction
12.8027
9.8069
5grams english
12.5607
9.5967
frrams_french
12,402
9.835
-german
11,6912
9,8432
Sgrams hebrew
7.097
6,2143
5grams italian
ggrams
11.3061
9.7771
5grams_russin russian
12.885
9,6805
5grams ggrams_spannsh
12.111
9.8079
Table 4: We construct traditional and augmented KD trees for n- n-grams for various languages, and of various lengths n. Our method outperforms a traditional KD tree in all cases.
28


Reviews:
[{"review_id": "FmN8QInsf6", "reviewer": "Reviewer_btGJ", "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.\n\nSome details of data structure construction: \n\nSkip 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).\n\nKd 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.", "questions": "My main question is about the motivation of the kd tree setting as stated in the weaknesses section. \n\nAnother 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.\n\nMinor formatting comment: The font size used for plots needs to be made larger.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "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.\n\nMotivation: 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.\n\nSetting: 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.\n\nDetails 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "ilfwkc8Zyq", "reviewer": "Reviewer_Azhq", "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.", "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?\n2. 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?\n2. 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.\n3. I suggest adding additional experiment comparisons with other learning-augmented algorithms.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The integration of learned frequencies into skip lists and KD trees is well-constructed and achieves optimal performance.\n2. Experiments show that the proposed algorithm outperforms classical algorithms and is robust to noise.\n3. 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.\n2. The noise robustness measure, denoted as $(\\alpha,\\beta)$-noisy, is somewhat unconventional and may lack generality.\n3. The experiments do not include comparisons with other learning-augmented algorithms given frequency predictions.\n4. 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.\n\n[1] Cao X, Chen J, Chen L, et al. Learning-augmented b-trees[J]. arXiv preprint arXiv:2211.09251, 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "AxVxHWngjY", "reviewer": "Reviewer_i7vt", "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. \n\nThe 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.\n\nThe 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.\n\nProposed learning augmented skip list uses predicted probabilities to promote items to next levels with exponentially increasing required thresholds for predicted probabilities. \n\nWhile 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).\n\nThe 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.  \n\nEmpirically, the paper shows constant factor speedup of the augmented data structures, compared to the conventional approach", "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?\n- 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?\n- Would be nice to present graphs with theoretically expected improvement and actual improvement in various settings.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- Provided approach is an elegant way to improve search data structure performance by considering probabilistic nature of the data. \n- Given the abundance of machine learning and statistical methods to serve as an oracle, proposed approach can be widely used in practice.\n- 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).\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "vNYzrEJz5o", "reviewer": "Reviewer_zTZd", "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.", "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.\n\nAilon, N., Chazelle, B., Clarkson, K. L., Liu, D., Mulzer, W., & Seshadhri, C. (2011). Self-improving algorithms. SIAM Journal on Computing, 40(2), 350-375.\n\nCayton, L., & Dasgupta, S. (2007). A learning framework for nearest neighbor search. Advances in Neural Information Processing Systems, 20.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "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). \n\n-- The font size in the figures is beyond excusable", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)