Learning-Augmented Search Data Structures

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. [KBC⁺18] first showed that ML could be incorporated to create data structures that support faster look-up operations while also saving an order-of-magnitude of memory compared to optimized data structures oblivious to such ML heuristics. Subsequently, learning-augmented algorithms [MV20] 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 [Mit18, LLW22], algorithms with faster runtimes [DIL⁺21, CSVZ22, DMVW23], mechanisms with better accuracy-privacy tradeoffs [KADV23], online algorithms with better performance than information-theoretic limits [PSK18, GP19, LLMV20, WLW20, WZ20, BMS20, AGP20, ACL⁺21, AGKP21, IKQP21, LV21, ACI22, AGKP22, APT22, GLS⁺22, KBTV22, JLL⁺22, SGM22, AGKK23, ACE⁺23, SLLA23], streaming algorithms with better accuracy-space tradeoffs [HIKV19, IVY19, JLL⁺20, CIW22, CEI⁺22, LLL⁺23], and polynomial-time algorithms beyond hardness-of-approximation limits, e.g., NP-hardness [EFS⁺22, NCN23, CLR⁺24].

In this paper, we focus on the consolidation of ML advice to improve data structures for the fundamental problem of searching for elements among a large dataset. For this purpose, tree-based structures stand out as a 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.

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 [LLW22] to study learning-augmented binary search trees noted that although previous results already characterized the statically optimal tree if the underlying distribution is known [Knu71, Meh77], these methods do not handle dynamic insertion operations. In contrast, skip lists, introduced by  [Pug90a], maintain balance probabilistically, offering a simpler implementation while delivering substantial speed enhancements [Pug90b]. Skip lists are generally built iteratively in levels. The bottom levels of the skip list is an ordinary-linked 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 an ordered link list. Traditional skip lists are built by promoting each item in a level to a higher level randomly with a fixed probability $p\in(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 an 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 a lower list. It can be shown that the expected number of steps in the search is $\mathcal{O}\left(\frac{1}{p}\log_{1/p}n\right)$ so that $p$ 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 [SL00, LJ13], high efficiency for dynamic datasets [GZ08, PT10], network routing [HPJ03, ASS20], and real-time analytics [BBB⁺20, ZCL⁺23]. Thus while binary search trees have been a long-standing choice for querying ordered elements, skip lists offer a 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 a more efficient alternative to binary search trees for searching in higher-dimensional 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 is repeated until every data point is included in the tree, creating a 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 a 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 a 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 $N(k;\alpha,N)$ is inversely proportional to its rank $k$, raised to the power of $\alpha$ (where $\alpha$ is a positive parameter), in a dataset of $N$ elements. In particular, we have $N(k;\alpha,N)=\frac{1/k^{\alpha}}{\sum_{n=1}^{N}(1/n^{\alpha})}$. The value of $\alpha$ determines the steepness of the distribution so that a smaller $\alpha$ value, i.e., closer to $0$, makes it more uniform, while a larger $\alpha$ 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 technology. Indeed, they appear in many applications such as word frequencies in natural language [WW16, BHZ18], city populations [Gab99, VA15], biological cellular distributions [LVM⁺23], income distribution [San15], etc.

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 on the incoming queries. Although we 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.

In this section, we describe our construction for a learning-augmented skip list and show various 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.

We first describe our learning-augmented skip list, which utilizes predictions $p_{i}$ for each item $i\in[n]$, from an oracle. Similar to a traditional skip list, the bottom level of our skip list is an ordinary-linked 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 a level to a higher level randomly with a fixed probability $p\in(0,1)$, we automatically promote the item $i$ to a level $\ell$ if its predicted frequency $p_{i}$ satisfies $p_{i}\geq\frac{2^{\ell}-1}{n}$. Otherwise, we promote the item with probability $\frac{1}{2}$. 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.

It is worth addressing a 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 to 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. A potential approach to resolve this issue would be to use a separate threshold probability is set for each level so that only nodes with a 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., $\frac{1}{2^{n}}$. We can thus first filters out the low-frequency elements and place them remain on the bottom level of the skip list and then proceed with using a 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 an expected search time of $\Omega(n)$. Hence, we ensure that each element still has a 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 [Knu71, Meh77] do not apply, as observed by [LLW22]. 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 Section 3.

We first show an upper bound on the expected search time of our learning-augmented skip-list.

To achieve Theorem 2.1, we first show that each item $i\in[n]$ must be contained at some level $\max\left(0,1+{\left\lfloor\log(np_{i})\right\rfloor}\right)$, depending on the predicted frequency $p_{i}$ of the item. We also show that with high probability, the total number of levels in the skip list is at most $\mathcal{O}\left(\log n\right)$. This allows us to upper bound the expected search time for item $i$ by at most $2C+2\min\left(\log\frac{1}{p_{i}},\log n\right)$. We can then analyze the expected search time across the true probability distribution $f_{i}$. Putting these steps together, we obtain Theorem 2.1.

We also prove a lower bound on the expected search time of an item drawn from a probability distribution $f$ for any skip list.

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

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 a Zipfian distribution.

Next, we show that our learning-augmented skip list construction is robust to somewhat inaccurate oracles. Let $f$ be the true-scaled frequency vector so that for each $i\in[n]$, $f_{i}$ is the probability that a random query corresponds to $i$. Let $p$ be the predicted frequency vector, so that for each $i\in[n]$, $p_{i}$ is the predicted probability that a random query corresponds to $i$. For $\alpha,\beta\in(0,1)$, we call an oracle $(\alpha,\beta)$-noisy if for all $i\in[n]$, we have $p_{i}\geq\alpha\cdot f_{i}-\beta$. Then we have the following guarantees for an $(\alpha,\beta)$-noisy oracle:

To achieve Lemma 2.5, we parameterize our analysis in Theorem 2.1. Due to the guarantees of the $(\alpha,\beta)$-noisy oracles, we can write $p_{i}\geq\frac{\alpha}{2}\cdot f_{i}$, which allows us to express the search time $\log\frac{1}{p_{i}}$ in terms of the true entropy of the distribution and a small additive constant that stems from $\log\frac{1}{\alpha}$. In fact, we remark that even when the predictions are arbitrarily inaccurate, our learning-augmented skip list still has expected query time $\mathcal{O}\left(\log n\right)$, since the total number of levels is at most $\mathcal{O}\left(\log n\right)$ with high probability. Since the expected query list of an oblivious skip list is also $\mathcal{O}\left(\log n\right)$, then the expected query time of our learning-augmented skip list is within a constant multiplicative factor, even with arbitrarily poor predictions.

In this section we present details on our novel approach to KD tree construction. First, we present the algorithm that constructs a learning-augmented KD tree. We focus on the setting where queries can be made on the search space rather than the items in the dataset, which is much more interesting for high-dimensional datasets, since even building a balanced tree on 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_{i}$ for element $i$ of our dataset, the intuition of our method is straightforward. Whereas a learning-augmented binary search tree would attempt to 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.

Thus, instead of relying on standard techniques to determine the splitting point of our dataset, we find a specific dimension in which there exists a balanced split such that the probability of a query being on either branch of the tree is 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 to reject these negative queries more quickly.

We prove the following guarantees on the performance of our learning-augmented KD tree, first assuming that our oracle is perfect.

The analysis of Theorem 3.1 corresponding to our learning-augmented KD tree follows from a similar structure as the proof of Theorem 2.1. However, the crucial difference is that the universe size is now $N$, which is exponential in $d$. Thus constructions that consider distributions over all of $[N]$ may suffer $\mathcal{O}\left(\log N\right)=\mathcal{O}\left(d\log\Delta\right)$ query time, which can be prohibitively expensive for large $d$, e.g., high-dimensional data. Hence, our algorithm requires a bit more care in the truncation of queries with low probability and instead, we build a balanced KD tree for any item with less than $\frac{1}{n^{2}}$ probability of being queried, so that each of their query times is at most $\mathcal{O}\left(\log n\right)$. We further remark this implies robustness of our data structure to arbitrarily poor predictions, by a similar argument as in Section 2.

We next prove a lower bound on the expected search time of an item drawn from a probability distribution $f$ for any KD tree.

By Theorem 3.1 and Lemma 2.3, we thus have the following corollary for the expected query time on our learning-augmented KD tree on a set of search queries that follows a Zipfian distribution.

Finally, we show near-optimality when given imperfect predictions from a $(\alpha,\beta)$-noisy oracle:

In this section, we describe a number of empirical evaluations demonstrating the efficiency of our learning-augmented search data structures on both synthetic and real-world datasets. We provide additional experiments in Appendix D.

Skip lists on CAIDA dataset. In the CAIDA datasets [CAI16], 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 1, the CAIDA datasets can be approximately characterized by an $\alpha$ 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-augmented and the oblivious skip lists in Figure 2 is roughly $1.86\times$, which is surprisingly and perhaps coincidentally close to our theoretical speed-up of roughly $1.81\times$ on a Zipfian dataset with exponent $1.37$.

Next, we demonstrate that our proposed algorithm still manages to outperform the classic skip list even 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 during a 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 T1 contains a set of elements S1, and some future query stream contains a 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 to build S2; otherwise, if the element has not shown up during T1 (i.e., in S1), then we assume its probability to be 0. After this, the probability vector is normalized to sum to 1, resulting in a predicted probability vector to 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 3(a) and we show that even when the prediction is not perfect, the augmented skip list still performs better than a conventional skip list.

Figure 3(a) 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-augmented 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 3(a) means:

• •

  10_2: the first 10 minutes of data are used to create the reference (i.e., oracle) and the last 2 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 2 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 to create the reference and the last 6 minutes are used for testing.

Further analysis of the temporal change of item frequency shows the reason behind the good performance of the history-based oracle. Figure 3(b) 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 3(b) 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^{12}$ points in 3-dimensional space, with frequencies given by a fixed Zipfian distribution with parameters $a=5,b=2$ – parameters at which our method greatly outperforms a standard KD tree. In order to simulate constructing the tree on noisy data, we multiply the ground truth query probabilities by numbers sampled uniformly from 1 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^{14}$ 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 a traditional KD tree.

Next, we generate datasets of $2^{12}$ points in 3-dimensional space, with frequencies given by a Zipfian distribution with parameters $a,b$. In the left plot, we assign these Zipfian weights randomly. In the right plot, however, we assign Zipfian weights with ranks decreasing with the distance to some random data point. We then query the tree $2^{14}$ 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 5.

KD trees on 3D point-cloud datasets. Finally, we evaluate our method on point cloud data generated from the Stanford Lucy mesh [Sta96], with dimensions $\sim 1000\times 500\times 1500$. We first uniformly sample $2^{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 a new set of $2^{16}$ surface samples on 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 a traditional KD tree yields an average lookup depth of 17.6.

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.